chore: update pg_upgrade initiate.sh to support nix-based upgrades (#1057)

Author: pcnc

ansible/files/admin_api_scripts/pg_upgrade_scripts/complete.sh

@@ -161,12 +161,20 @@ function apply_auth_scheme_updates {
 
 function start_vacuum_analyze {
     echo "complete" > /tmp/pg-upgrade-status
-    su -c 'vacuumdb --all --analyze-in-stages' -s "$SHELL" postgres
+    if ! command -v nix &> /dev/null; then
+        su -c 'vacuumdb --all --analyze-in-stages' -s "$SHELL" postgres
+    else
+        su -c '. /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh && vacuumdb --all --analyze-in-stages' -s "$SHELL" postgres
+    fi
     echo "Upgrade job completed"
 }
 
 trap cleanup ERR
 
+echo "C.UTF-8 UTF-8" > /etc/locale.gen
+echo "en_US.UTF-8 UTF-8" >> /etc/locale.gen
+locale-gen
+
 if [ -z "$IS_CI" ]; then
     complete_pg_upgrade >> $LOG_FILE 2>&1 &
 else 

ansible/files/admin_api_scripts/pg_upgrade_scripts/initiate.sh

@@ -29,10 +29,13 @@ SCRIPT_DIR=$(dirname -- "$0";)
 source "$SCRIPT_DIR/common.sh"
 
 IS_CI=${IS_CI:-}
-LOG_FILE="/var/log/pg-upgrade-initiate.log"
+IS_LOCAL_UPGRADE=${IS_LOCAL_UPGRADE:-}
+IS_NIX_UPGRADE=${IS_NIX_UPGRADE:-}
+IS_NIX_BASED_SYSTEM="false"
 
 PGVERSION=$1
 MOUNT_POINT="/data_migration"
+LOG_FILE="/var/log/pg-upgrade-initiate.log"
 
 POST_UPGRADE_EXTENSION_SCRIPT="/tmp/pg_upgrade/pg_upgrade_extensions.sql"
 OLD_PGVERSION=$(run_sql -A -t -c "SHOW server_version;")
@@ -41,6 +44,15 @@ POSTGRES_CONFIG_PATH="/etc/postgresql/postgresql.conf"
 PGBINOLD="/usr/lib/postgresql/bin"
 PGLIBOLD="/usr/lib/postgresql/lib"
 
+PG_UPGRADE_BIN_DIR="/tmp/pg_upgrade_bin/$PGVERSION"
+
+if [ -L "$PGBINOLD/pg_upgrade" ]; then
+    BINARY_PATH=$(readlink -f "$PGBINOLD/pg_upgrade")
+    if [[ "$BINARY_PATH" == *"nix"* ]]; then
+        IS_NIX_BASED_SYSTEM="true"
+    fi
+fi
+
 # If upgrading from older major PG versions, disable specific extensions
 if [[ "$OLD_PGVERSION" =~ ^14.* ]]; then
     EXTENSIONS_TO_DISABLE+=("${PG14_EXTENSIONS_TO_DISABLE[@]}")
@@ -107,7 +119,7 @@ cleanup() {
     echo "Removing SUPERUSER grant from postgres"
     run_sql -c "ALTER USER postgres WITH NOSUPERUSER;"
 
-    if [ -z "$IS_CI" ]; then
+    if [ -z "$IS_CI" ] && [ -z "$IS_LOCAL_UPGRADE" ]; then
         echo "Unmounting data disk from ${MOUNT_POINT}"
         umount $MOUNT_POINT
     fi
@@ -175,18 +187,50 @@ function initiate_upgrade {
     PGDATAOLD=$(cat "$POSTGRES_CONFIG_PATH" | grep data_directory | sed "s/data_directory = '\(.*\)'.*/\1/")
 
     PGDATANEW="$MOUNT_POINT/pgdata"
-    PG_UPGRADE_BIN_DIR="/tmp/pg_upgrade_bin/$PGVERSION"
-    PGBINNEW="$PG_UPGRADE_BIN_DIR/bin"
-    PGLIBNEW="$PG_UPGRADE_BIN_DIR/lib"
-    PGSHARENEW="$PG_UPGRADE_BIN_DIR/share"
 
     # running upgrade using at least 1 cpu core
     WORKERS=$(nproc | awk '{ print ($1 == 1 ? 1 : $1 - 1) }')
+
+    # To make nix-based upgrades work for testing, create a pg binaries tarball with the following contents:
+    #  - nix_flake_version - a7189a68ed4ea78c1e73991b5f271043636cf074
+    # Where the value is the commit hash of the nix flake that contains the binaries
+
+    if [ -n "$IS_LOCAL_UPGRADE" ]; then
+        mkdir -p "$PG_UPGRADE_BIN_DIR"
+        mkdir -p /tmp/persistent/
+        echo "a7189a68ed4ea78c1e73991b5f271043636cf074" > "$PG_UPGRADE_BIN_DIR/nix_flake_version"
+        tar -czf "/tmp/persistent/pg_upgrade_bin.tar.gz" -C "/tmp/pg_upgrade_bin" .
+        rm -rf /tmp/pg_upgrade_bin/
+    fi
 
     echo "1. Extracting pg_upgrade binaries"
     mkdir -p "/tmp/pg_upgrade_bin"
     tar zxf "/tmp/persistent/pg_upgrade_bin.tar.gz" -C "/tmp/pg_upgrade_bin"
 
+    PGSHARENEW="$PG_UPGRADE_BIN_DIR/share"
+
+    if [ -f "$PG_UPGRADE_BIN_DIR/nix_flake_version" ]; then
+        IS_NIX_UPGRADE="true"
+        NIX_FLAKE_VERSION=$(cat "$PG_UPGRADE_BIN_DIR/nix_flake_version")
+
+        if [ "$IS_NIX_BASED_SYSTEM" = "false" ]; then
+            echo "1.1. Nix is not installed; installing."
+
+            curl --proto '=https' --tlsv1.2 -sSf -L https://install.determinate.systems/nix | sh -s -- install --no-confirm \
+            --extra-conf "substituters = https://cache.nixos.org https://nix-postgres-artifacts.s3.amazonaws.com" \
+            --extra-conf "trusted-public-keys = nix-postgres-artifacts:dGZlQOvKcNEjvT7QEAJbcV6b6uk7VF/hWMjhYleiaLI=% cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY="
+        fi
+
+        echo "1.2. Installing flake revision: $NIX_FLAKE_VERSION"
+        # shellcheck disable=SC1091
+        source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
+        PG_UPGRADE_BIN_DIR=$(nix build "github:supabase/postgres/${NIX_FLAKE_VERSION}#psql_15/bin" --no-link --print-out-paths --extra-experimental-features nix-command --extra-experimental-features flakes)
+        PGSHARENEW="$PG_UPGRADE_BIN_DIR/share/postgresql"
+    fi
+
+    PGBINNEW="$PG_UPGRADE_BIN_DIR/bin"
+    PGLIBNEW="$PG_UPGRADE_BIN_DIR/lib"
+
     # copy upgrade-specific pgsodium_getkey script into the share dir
     chmod +x "$SCRIPT_DIR/pgsodium_getkey.sh"
     mkdir -p "$PGSHARENEW/extension"
@@ -220,7 +264,7 @@ function initiate_upgrade {
         locale-gen
     fi
 
-    if [ -z "$IS_CI" ]; then
+    if [ -z "$IS_CI" ] && [ -z "$IS_LOCAL_UPGRADE" ]; then
         # awk NF==3 prints lines with exactly 3 fields, which are the block devices currently not mounted anywhere
         # excluding nvme0 since it is the root disk
         echo "5. Determining block device to mount"
@@ -254,13 +298,17 @@ function initiate_upgrade {
     echo "8. Granting SUPERUSER to postgres user"
     run_sql -c "ALTER USER postgres WITH SUPERUSER;"
 
-    if [ -d "/usr/share/postgresql/${PGVERSION}" ]; then
-        mv "/usr/share/postgresql/${PGVERSION}" "/usr/share/postgresql/${PGVERSION}.bak"
-    fi
-    ln -s "$PGSHARENEW" "/usr/share/postgresql/${PGVERSION}"
+    if [ -z "$IS_NIX_UPGRADE" ]; then
+        if [ -d "/usr/share/postgresql/${PGVERSION}" ]; then
+            mv "/usr/share/postgresql/${PGVERSION}" "/usr/share/postgresql/${PGVERSION}.bak"
+        fi
+
+        ln -s "$PGSHARENEW" "/usr/share/postgresql/${PGVERSION}"
+        cp --remove-destination "$PGLIBNEW"/*.control "$PGSHARENEW/extension/"
+        cp --remove-destination "$PGLIBNEW"/*.sql "$PGSHARENEW/extension/"
 
-    cp --remove-destination "$PGLIBNEW"/*.control "$PGSHARENEW/extension/"
-    cp --remove-destination "$PGLIBNEW"/*.sql "$PGSHARENEW/extension/"
+        export LD_LIBRARY_PATH="${PGLIBNEW}"
+    fi
 
     # This is a workaround for older versions of wrappers which don't have the expected
     #  naming scheme, containing the version in their library's file name
@@ -287,8 +335,6 @@ function initiate_upgrade {
         fi
     fi
 
-    export LD_LIBRARY_PATH="${PGLIBNEW}"
-
     echo "9. Creating new data directory, initializing database"
     chown -R postgres:postgres "$MOUNT_POINT/"
     rm -rf "${PGDATANEW:?}/"
@@ -308,6 +354,10 @@ function initiate_upgrade {
 EOF
     )
 
+    if [ "$IS_NIX_BASED_SYSTEM" = "true" ]; then
+        UPGRADE_COMMAND=". /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh && $UPGRADE_COMMAND"
+    fi
+    
     su -c "$UPGRADE_COMMAND --check" -s "$SHELL" postgres
 
     echo "10. Stopping postgres; running pg_upgrade"
@@ -324,7 +374,11 @@ EOF
         CI_stop_postgres
     fi
 
-    su -c "$UPGRADE_COMMAND" -s "$SHELL" postgres
+    if [ "$IS_NIX_BASED_SYSTEM" = "true" ]; then
+        LC_ALL=en_US.UTF-8 LC_CTYPE=en_US.UTF-8 LANGUAGE=en_US.UTF-8 LANG=en_US.UTF-8 LOCALE_ARCHIVE=/usr/lib/locale/locale-archive su -pc "$UPGRADE_COMMAND" -s "$SHELL" postgres
+    else
+        su -c "$UPGRADE_COMMAND" -s "$SHELL" postgres
+    fi
 
     # copying custom configurations
     echo "11. Copying custom configurations"
@@ -349,7 +403,7 @@ EOF
 trap cleanup ERR
 
 echo "running" > /tmp/pg-upgrade-status
-if [ -z "$IS_CI" ]; then
+if [ -z "$IS_CI" ] && [ -z "$IS_LOCAL_UPGRADE" ]; then
     initiate_upgrade >> "$LOG_FILE" 2>&1 &
     echo "Upgrade initiate job completed"
 else