From 9c869b0ec990e13bc5cc000fad3ff6b1f1334d4b Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavezast@gmail.com>
Date: Mon, 29 Apr 2024 19:55:44 -0500
Subject: [PATCH] nix: add --replica to postgrest-with-postgresql-*

---
 nix/tools/withTools.nix | 43 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/nix/tools/withTools.nix b/nix/tools/withTools.nix
index 82781b7697..cfeb9e4a28 100644
--- a/nix/tools/withTools.nix
+++ b/nix/tools/withTools.nix
@@ -28,6 +28,7 @@ let
             "ARG_USE_ENV([PGRST_DB_SCHEMAS], [test], [Schema to expose])"
             "ARG_USE_ENV([PGTZ], [utc], [Timezone to use])"
             "ARG_USE_ENV([PGOPTIONS], [-c search_path=public,test], [PG options to use])"
+            "ARG_OPTIONAL_BOOLEAN([replica],, [Enable a replica for the database])"
           ];
         positionalCompletion = "_command";
         workingDir = "/";
@@ -61,6 +62,7 @@ let
           HBA_FILE="$tmpdir/pg_hba.conf"
           echo "local $PGDATABASE some_protected_user password" > "$HBA_FILE"
           echo "local $PGDATABASE all trust" >> "$HBA_FILE"
+          echo "local replication all trust" >> "$HBA_FILE"
 
           log "Initializing database cluster..."
           # We try to make the database cluster as independent as possible from the host
@@ -74,19 +76,46 @@ let
           pg_ctl -l "$tmpdir/db.log" -w start -o "-F -c listen_addresses=\"\" -c hba_file=$HBA_FILE -k $PGHOST -c log_statement=\"all\" " \
             >> "$setuplog"
 
+          log "Creating a minimally privileged $PGUSER connection role..."
+          createuser "$PGUSER" -U postgres --host="$tmpdir/socket" --no-createdb --no-inherit --no-superuser --no-createrole --no-replication --login
+
+          >&2 echo "${commandName}: You can connect with: psql 'postgres:///$PGDATABASE?host=$PGHOST' -U postgres"
+          >&2 echo "${commandName}: You can tail the logs with: tail -f $tmpdir/db.log"
+
+          if test "$_arg_replica" = "on"; then
+            replica_slot="replica_$RANDOM"
+            replica_dir="$tmpdir/$replica_slot"
+            replica_host="$tmpdir/socket_$replica_slot"
+
+            mkdir -p "$replica_host"
+
+            replica_setuplog="$tmpdir/setup_$replica_slot.log"
+            replica_dblog="$tmpdir/db_$replica_slot.log"
+
+            log "Running pg_basebackup for $replica_slot"
+
+            pg_basebackup -h "$PGHOST" -U postgres --wal-method=stream --create-slot --slot="$replica_slot" --write-recovery-conf -D "$replica_dir"
+
+            log "Starting replica on $replica_host"
+
+            pg_ctl -D "$replica_dir" -l "$replica_dblog" -w start -o "-F -c listen_addresses=\"\" -c hba_file=$HBA_FILE -k $replica_host -c log_statement=\"all\" " \
+              >> "$replica_setuplog"
+
+            >&2 echo "${commandName}: Replica enabled. You can connect to it with: psql 'postgres:///$PGDATABASE?host=$replica_host' -U postgres"
+            >&2 echo "${commandName}: You can tail the replica logs with: tail -f $replica_dblog"
+          fi
+
           # shellcheck disable=SC2317
           stop () {
             log "Stopping the database cluster..."
-            pg_ctl stop -m i >> "$setuplog"
+            pg_ctl stop --mode=immediate >> "$setuplog"
             rm -rf "$tmpdir/db"
+            if test "$_arg_replica" = "on"; then
+              pg_ctl -D "$replica_dir" stop --mode=immediate >> "$replica_setuplog"
+              rm -rf "$replica_dir"
+            fi
           }
           trap stop EXIT
-
-          log "Creating a minimally privileged $PGUSER connection role..."
-          createuser "$PGUSER" -U postgres --host="$tmpdir/socket" --no-createdb --no-inherit --no-superuser --no-createrole --no-replication --login
-
-          >&2 echo "${commandName}: You can connect with: psql 'postgres:///$PGDATABASE?host=$tmpdir/socket' -U postgres"
-          >&2 echo "${commandName}: You can tail the logs with: tail -f $tmpdir/db.log"
         fi
 
         if test "$_arg_fixtures"; then