From 6404489fbcc190f32e558035d5c1ca107119e2da Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Thalheim?= <joerg@thalheim.io>
Date: Tue, 4 Jun 2024 12:56:39 +0200
Subject: [PATCH] postgresql: add new method to create users and databases

---
 checks/postgresql/default.nix      |  4 +-
 clanModules/postgresql/default.nix | 99 +++++++++++++++++++++++++++---
 2 files changed, 91 insertions(+), 12 deletions(-)

diff --git a/checks/postgresql/default.nix b/checks/postgresql/default.nix
index ced0863b..aec5a649 100644
--- a/checks/postgresql/default.nix
+++ b/checks/postgresql/default.nix
@@ -9,14 +9,14 @@
           self.clanModules.postgresql
           self.clanModules.localbackup
         ];
-        clan.postgresl.databases = [ "test" ];
+        clan.postgresl.databases.test = {};
         clan.localbackup.targets.hdd.directory = "/mnt/external-disk";
       };
     testScript = ''
       start_all()
       machine.succeed("systemctl status postgresql")
       machine.wait_for_unit("postgresql")
-      machine.succeed("localbackup-create")
+      machine.succeed("/run/current-system/sw/bin/localbackup-create >&2")
       machine.succeed("ls -la /var/backups/postgresql")
     '';
   }
diff --git a/clanModules/postgresql/default.nix b/clanModules/postgresql/default.nix
index 5d819477..a6abe9c5 100644
--- a/clanModules/postgresql/default.nix
+++ b/clanModules/postgresql/default.nix
@@ -8,7 +8,7 @@ let
   createDatatbaseState =
     db:
     let
-      folder = "/var/backup/postgresql/${db}";
+      folder = "/var/backup/postgresql/${db.name}";
       curFile = "${folder}/dump.sql.zstd";
       prevFile = "${folder}/dump.sql.prev.zstd";
       inProgressFile = "${folder}/dump.sql.in-progress.zstd";
@@ -21,7 +21,11 @@ let
           if [ -e ${curFile} ]; then
             mv ${curFile} ${prevFile}
           fi
-          pg_dump -C ${db} | \
+          while [[ "$(systemctl is-active postgres)" == activating ]]; then
+            sleep 1
+          done
+          systemctl is-active postgres
+          pg_dump -C ${db.name} | \
             ${pkgs.zstd}/bin/zstd --rsyncable | \
             > ${inProgressFile}
           mv ${inProgressFile} ${curFile}
@@ -29,24 +33,99 @@ let
       '';
       postRestoreCommand = ''
         if [[ -f ${prevFile} ]]; then
-          zstd --decompress --stdout ${prevFile} | psql -d ${db}
+          zstd --decompress --stdout ${prevFile} | psql -d ${db.name}
         fi
       '';
     };
+
+  createDatabase = db: ''
+    CREATE DATABASE ${db.name} ${
+      lib.concatStringsSep " " (
+        lib.mapAttrsToList (name: value: "${name} = ':${value}'") db.createOptions
+      )
+    }
+  '';
+
+  createDatabaseArgs = db: ''
+    ${lib.concatStringsSep " " (
+      lib.mapAttrsToList (name: value: "-v ${name}=${lib.escapeShellArg value}") db.createOptions
+    )}
+  '';
+  cfg = config.clan.postgresl;
+
+  userClauses = lib.mapAttrsToList (
+    _: user: ""
+    ''$PSQL -tAc "SELECT 1 FROM pg_roles WHERE rolname='${user.name}'" | grep -q 1 || $PSQL -tAc 'CREATE USER "${user.name}"' ''
+  ) cfg.users;
+  databaseClauses = lib.mapAttrsToList (
+    name: db:
+    lib.optionalString (db.create) ''$PSQL -d postgres -c "SELECT 1 FROM pg_database WHERE datname = '${name}'" | grep -q 1 || $PSQL -d postgres -c ${lib.escapeShellArg (createDatabase db)} ${createDatabaseArgs db}''
+  ) cfg.databases;
 in
 {
   options.clan.postgresl = {
     databases = lib.mkOption {
-      type = lib.types.listOf lib.types.str;
-      default = [ "clan" ];
+      default = { };
+      type = lib.types.attrsOf (
+        lib.types.submodule (
+          { name, ... }:
+          {
+            options = {
+              name = lib.mkOption {
+                type = lib.types.str;
+                default = name;
+              };
+              # set to false, in case the upstream module uses ensureDatabase option
+              create = lib.mkOption {
+                type = lib.types.bool;
+                default = true;
+              };
+              createOptions = lib.mkOption {
+                type = lib.types.lazyAttrsOf lib.types.str;
+                default = { };
+                example = {
+                  TEMPLATE = "template0";
+                  LC_COLLATE = "C";
+                  LC_CTYPE = "C";
+                  ENCODING = "UTF8";
+                  OWNER = "foo";
+                };
+              };
+            };
+          }
+        )
+      );
+    };
+    users = lib.mkOption {
+      default = { };
+      type = lib.types.attrsOf (
+        lib.types.submodule (
+          { name, ... }:
+          {
+            options.name = lib.mkOption {
+              type = lib.types.str;
+              default = name;
+            };
+          }
+        )
+      );
     };
   };
   config = {
     services.postgresql.enable = true;
-    clanCore.state = lib.listToAttrs (
-      builtins.map (
-        db: lib.nameValuePair "postgresql-${db}" (createDatatbaseState db)
-      ) config.clan.postgresl.databases
-    );
+    # We are duplicating a bit the upstream module but allow to create databases with options
+    systemd.services.postgresql.postStart = ''
+      PSQL="psql --port=${builtins.toString config.services.postgresql.settings.port}"
+
+      while ! $PSQL -d postgres -c "" 2> /dev/null; do
+          if ! kill -0 "$MAINPID"; then exit 1; fi
+          sleep 0.1
+      done
+      ${lib.concatStringsSep "\n" userClauses}
+      ${lib.concatStringsSep "\n" databaseClauses}
+    '';
+    clanCore.state = lib.mapAttrs' (
+      _: db: lib.nameValuePair "postgresql-${db.name}" (createDatatbaseState db)
+    ) config.clan.postgresl.databases;
   };
 }