From 7b3bb199d5291093596b65ef86d259d4185ade97 Mon Sep 17 00:00:00 2001
From: Michael Scherer <misc@redhat.com>
Date: Wed, 17 Feb 2021 19:15:48 +0100
Subject: [PATCH] Add support for extensions loading and creation.

Using POSTGRESQL_LIBRARIES and POSTGRESQL_EXTENSIONS allow
users to enable extensions (some bundled in the container) without
having to do anything but adding 2 variables.
---
 src/root/usr/bin/run-postgresql                 |  1 +
 .../container-scripts/postgresql/README.md      |  9 +++++++++
 .../container-scripts/postgresql/common.sh      | 17 +++++++++++++++++
 test/run_test                                   | 14 ++++++++++++--
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/root/usr/bin/run-postgresql b/src/root/usr/bin/run-postgresql
index 2367e57e..5ac3e928 100755
--- a/src/root/usr/bin/run-postgresql
+++ b/src/root/usr/bin/run-postgresql
@@ -47,6 +47,7 @@ if $PG_INITIALIZED ; then
     create_users
 fi
 
+create_extensions
 process_extending_files \
     "${APP_DATA}/src/postgresql-start" \
     "${CONTAINER_SCRIPTS_PATH}/start"
diff --git a/src/root/usr/share/container-scripts/postgresql/README.md b/src/root/usr/share/container-scripts/postgresql/README.md
index bc01aa9f..5b43c6a6 100644
--- a/src/root/usr/share/container-scripts/postgresql/README.md
+++ b/src/root/usr/share/container-scripts/postgresql/README.md
@@ -102,6 +102,15 @@ The following environment variables influence the PostgreSQL configuration file.
        Set to an estimate of how much memory is available for disk caching by the operating system and within the database itself
 
 
+The following environment variables deal with extensions. They are all optional, and if not set, no extensions will be enabled.
+
+**`POSTGRESQL_LIBRARIES`**
+       A comma-separated list of libraries that Postgres will preload using shared_preload_libraries.
+
+**`POSTGRESQL_EXTENSIONS`**
+       A space-separated list of extensions to create when the server start. Once created, the extensions will stay even if the variable is cleared.
+
+
 You can also set the following mount points by passing the `-v /host/dir:/container/dir:Z` flag to Docker.
 
 **`/var/lib/pgsql/data`**  
diff --git a/src/root/usr/share/container-scripts/postgresql/common.sh b/src/root/usr/share/container-scripts/postgresql/common.sh
index a8592ffe..a16ecde5 100644
--- a/src/root/usr/share/container-scripts/postgresql/common.sh
+++ b/src/root/usr/share/container-scripts/postgresql/common.sh
@@ -143,6 +143,13 @@ function should_hack_data_sync_retry() {
   return 1
 }
 
+function generate_postgresql_libraries_config() {
+  if [ -v POSTGRESQL_LIBRARIES ]; then
+    echo "shared_preload_libraries='${POSTGRESQL_LIBRARIES}'" >> "${POSTGRESQL_CONFIG_FILE}"
+  fi
+}
+
+
 # New config is generated every time a container is created. It only contains
 # additional custom settings and is included from $PGDATA/postgresql.conf.
 function generate_postgresql_config() {
@@ -160,6 +167,7 @@ function generate_postgresql_config() {
     echo "data_sync_retry = on" >>"${POSTGRESQL_CONFIG_FILE}"
   fi
 
+  generate_postgresql_libraries_config
   (
   shopt -s nullglob
   for conf in "${APP_DATA}"/src/postgresql-cfg/*.conf; do
@@ -478,3 +486,12 @@ process_extending_files()
     done
   done <<<"$(get_matched_files '*.sh' "$@" | sort -u)"
 }
+
+create_extensions()
+{
+  if [ -v POSTGRESQL_EXTENSIONS ]; then
+    for EXT in $POSTGRESQL_EXTENSIONS; do
+      psql -c "CREATE EXTENSION IF NOT EXISTS ${EXT};"
+    done;
+  fi
+}
diff --git a/test/run_test b/test/run_test
index 14860eae..7bc596be 100755
--- a/test/run_test
+++ b/test/run_test
@@ -29,6 +29,7 @@ run_s2i_enable_ssl_test
 run_upgrade_test
 run_migration_test
 run_pgaudit_test
+run_new_pgaudit_test
 "
 
 test $# -eq 1 -a "${1-}" == --list && echo "$TEST_LIST" && exit 0
@@ -38,6 +39,8 @@ test -n "${OS-}" || false 'make sure $OS is defined'
 
 CIDFILE_DIR=$(mktemp --suffix=postgresql_test_cidfiles -d)
 
+DOCKER_EXTRA_ARGS=
+
 volumes_to_clean=
 images_to_clean=()
 files_to_clean=
@@ -858,7 +861,7 @@ run_pgaudit_test()
   # create a dir for config
   config_dir=$(mktemp -d --tmpdir pg-hook-volume.XXXXX)
   add_cleanup_command /bin/rm -rf "$config_dir"
-  cp -r "$test_dir"/examples/pgaudit/* "$config_dir"
+  $DOCKER_EXTRA_ARGS || cp -r "$test_dir"/examples/pgaudit/* "$config_dir"
   setfacl -R -m u:26:rwx "$config_dir"
 
   # create a dir for data
@@ -867,6 +870,7 @@ run_pgaudit_test()
 
   DOCKER_ARGS="
 -e POSTGRESQL_ADMIN_PASSWORD=password
+$DOCKER_EXTRA_ARGS
 -v ${config_dir}:/opt/app-root/src:Z
 -v ${data_dir}:/var/lib/pgsql/data:Z
   " create_container "$name"
@@ -878,7 +882,6 @@ run_pgaudit_test()
   # Deliberately moving heredoc into the container, otherwise it does not work
   # in podman 1.6.x due to https://bugzilla.redhat.com/show_bug.cgi?id=1827324
   docker exec -i $(get_cid "$name") bash -c "psql <<EOSQL
-CREATE EXTENSION pgaudit;
 SET pgaudit.log = 'read, ddl';
 CREATE DATABASE pgaudittest;
 EOSQL"
@@ -898,6 +901,13 @@ EOSQL"
   grep -E 'AUDIT: SESSION,.*,.*,READ,SELECT,,,SELECT' "${data_dir}"/userdata/log/postgresql-*.log
 }
 
+run_new_pgaudit_test() {
+    DOCKER_EXTRA_ARGS="
+-e POSTGRESQL_EXTENSIONS=pgaudit
+-e POSTGRESQL_LIBRARIES=pgaudit"
+    run_pgaudit_test
+}
+
 function run_all_tests() {
   for test_case in $TEST_LIST; do
     : "Running test $test_case"