From 03c975d160d8f4ee40b9d6f9af4b81a26bee0d96 Mon Sep 17 00:00:00 2001
From: Logger <ido.samuelson@gmail.com>
Date: Fri, 6 Dec 2024 20:25:46 -0700
Subject: [PATCH] postgres: Add setupSchemaScript option for Schema Setup

Path to a script that will set up or update the PostgreSQL database schema.

This script must be idempotent, meaning it can be run multiple times without causing unintended side effects. If you change your schema dynamically, ensure that this script handles such cases gracefully to maintain database integrity.
---
 examples/postgres/devenv.nix          |  3 ++
 src/modules/services/postgres.nix     | 40 +++++++++++++++++++++++++++
 tests/postgresql-localhost/devenv.nix |  3 ++
 3 files changed, 46 insertions(+)

diff --git a/examples/postgres/devenv.nix b/examples/postgres/devenv.nix
index 46b277893..a0d2d5972 100644
--- a/examples/postgres/devenv.nix
+++ b/examples/postgres/devenv.nix
@@ -11,5 +11,8 @@
     initialScript = ''
       CREATE EXTENSION IF NOT EXISTS postgis;
     '';
+    setupSchemaScript = ''
+      echo "script to run to setup or update database schema. This script must be idempotent."
+    '';
   };
 }
diff --git a/src/modules/services/postgres.nix b/src/modules/services/postgres.nix
index e38e16717..79c099a4b 100644
--- a/src/modules/services/postgres.nix
+++ b/src/modules/services/postgres.nix
@@ -98,6 +98,15 @@ let
     ''
     else "";
 
+  runSetupSchemaScript =
+    if cfg.setupSchemaScript == null
+    then ''
+      echo "script not provided, skipping."
+    ''
+    else ''
+      ${cfg.setupSchemaScript}
+    '';
+
   toStr = value:
     if true == value
     then "yes"
@@ -158,10 +167,29 @@ let
     fi
     unset POSTGRES_RUN_INITIAL_SCRIPT
   '';
+
+  setupSchemaScript = pkgs.writeShellScriptBin "setup-schema-script" ''
+    echo
+    echo "PostgreSQL is setting up the schema"
+    echo
+    OLDPGHOST="$PGHOST"
+    PGHOST=${q runtimeDir}
+
+    pg_ctl -D "$PGDATA" -w start -o "-c unix_socket_directories=${runtimeDir} -c listen_addresses= -p ${toString cfg.port}"
+    ${runSetupSchemaScript}
+    pg_ctl -D "$PGDATA" -m fast -w stop
+    PGHOST="$OLDPGHOST"
+    unset OLDPGHOST
+    echo
+    echo "PostgreSQL setup schema complete."
+    echo
+  '';
+
   startScript = pkgs.writeShellScriptBin "start-postgres" ''
     set -euo pipefail
     mkdir -p ${q runtimeDir}
     ${setupScript}/bin/setup-postgres
+    ${setupSchemaScript}/bin/setup-schema-script
     exec ${postgresPkg}/bin/postgres
   '';
 in
@@ -327,6 +355,18 @@ in
       '';
     };
 
+    setupSchemaScript = lib.mkOption {
+      type = types.nullOr types.str;
+      default = null;
+      description = ''
+        Path to a script that will set up or update the PostgreSQL database schema. This script must be idempotent, meaning it can be run multiple times without causing unintended side effects. 
+        If your schema changes dynamically, ensure that this script handles such cases gracefully to maintain database integrity.
+      '';
+      example = lib.literalExpression ''
+        "path/to/your/schema/setup/script.sh"
+      '';
+    };
+
     hbaConf = lib.mkOption {
       type = types.nullOr types.str;
       default = null;
diff --git a/tests/postgresql-localhost/devenv.nix b/tests/postgresql-localhost/devenv.nix
index 0b03a291b..fdd824af6 100644
--- a/tests/postgresql-localhost/devenv.nix
+++ b/tests/postgresql-localhost/devenv.nix
@@ -6,5 +6,8 @@
     initialScript = ''
       CREATE USER postgres SUPERUSER;
     '';
+    setupSchemaScript = ''
+      echo "script to run to setup or update database schema. This script must be idempotent."
+    '';
   };
 }