From 523885ac25e16848e1dfcc05cb484504b7aa8d33 Mon Sep 17 00:00:00 2001
From: alexbelgium <alexandre.pary@gmail.com>
Date: Thu, 29 May 2025 15:28:53 +0200
Subject: [PATCH] Upgrade logic

---
 postgres/rootfs/etc/cont-init.d/99-run.sh | 117 +++++++++++++---------
 1 file changed, 72 insertions(+), 45 deletions(-)

diff --git a/postgres/rootfs/etc/cont-init.d/99-run.sh b/postgres/rootfs/etc/cont-init.d/99-run.sh
index c02dc99e4..39b700ad6 100755
--- a/postgres/rootfs/etc/cont-init.d/99-run.sh
+++ b/postgres/rootfs/etc/cont-init.d/99-run.sh
@@ -5,8 +5,6 @@ set -e
 CONFIG_HOME="/config"
 PGDATA="${PGDATA:-/config/database}"
 export PGDATA
-PG_VERSION_FILE="$PGDATA/pg_major_version"
-
 PG_MAJOR_VERSION="${PG_MAJOR:-15}"
 
 mkdir -p "$PGDATA"
@@ -31,7 +29,18 @@ fi
 bashio::log.info "Waiting for PostgreSQL to start..."
 
 (
-bashio::net.wait_for 5432 localhost 900
+wait_for_postgres() {
+    local tries=0
+    while ! pg_isready -h "$DB_HOSTNAME" -p "$DB_PORT" -U "$DB_USERNAME" >/dev/null 2>&1; do
+        tries=$((tries+1))
+        if [ "$tries" -ge 60 ]; then
+            bashio::log.error "Postgres did not start after 2 minutes, aborting."
+            exit 1
+        fi
+        echo "PostgreSQL is starting up... ($tries/60)"
+        sleep 2
+    done
+}
 
 DB_PORT=5432
 DB_HOSTNAME=localhost
@@ -43,77 +52,93 @@ if bashio::config.has_value "POSTGRES_USER"; then
 fi
 export DB_PORT DB_HOSTNAME DB_USERNAME DB_PASSWORD
 
-until pg_isready -h "$DB_HOSTNAME" -p "$DB_PORT" -U "$DB_USERNAME" >/dev/null 2>&1; do
-    echo "PostgreSQL is starting up..."
-    sleep 2
-done
+wait_for_postgres
 
-update_postgres() {
-    OLD_PG_VERSION=$(cat "$PG_VERSION_FILE" 2>/dev/null || echo "$PG_MAJOR_VERSION")
-    if [ "$OLD_PG_VERSION" != "$PG_MAJOR_VERSION" ]; then
-        bashio::log.warning "PostgreSQL major version changed ($OLD_PG_VERSION → $PG_MAJOR_VERSION). Running upgrade..."
+# --------- Version detection directly from cluster & databases ----------
 
-        export DATA_DIR="$PGDATA"
-        export BINARIES_DIR="/usr/lib/postgresql/$PG_MAJOR_VERSION/bin"
-        export BACKUP_DIR="/config/backups"
-        export PSQL_VERSION="$PG_MAJOR_VERSION"
-
-        apt-get update &>/dev/null
-        apt-get install -y procps rsync postgresql-$PG_MAJOR_VERSION postgresql-$OLD_PG_VERSION &>/dev/null
-
-        TMP_SCRIPT=$(mktemp)
-        wget https://raw.githubusercontent.com/linkyard/postgres-upgrade/refs/heads/main/upgrade-postgres.sh -O "$TMP_SCRIPT"
-        chmod +x "$TMP_SCRIPT"
-        "$TMP_SCRIPT"
-
-        echo "$PG_MAJOR_VERSION" > "$PG_VERSION_FILE"
-        bashio::log.info "PostgreSQL major version upgrade complete."
-        RESTART_NEEDED=true
+get_pgdata_version() {
+    if [ -f "$PGDATA/PG_VERSION" ]; then
+        cat "$PGDATA/PG_VERSION"
+    else
+        bashio::log.error "FATAL: $PGDATA/PG_VERSION not found; cannot determine cluster version."
+        exit 1
     fi
 }
 
 get_available_extension_version() {
     local extname="$1"
-    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -tAc \
+    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
         "SELECT default_version FROM pg_available_extensions WHERE name = '$extname';" 2>/dev/null | xargs
 }
 
 is_extension_available() {
     local extname="$1"
     local result
-    result=$(psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -tAc \
+    result=$(psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
         "SELECT 1 FROM pg_available_extensions WHERE name = '$extname';" 2>/dev/null | xargs)
     [[ "$result" == "1" ]]
 }
 
 get_user_databases() {
-    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -tAc \
+    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
         "SELECT datname FROM pg_database WHERE datistemplate = false AND datallowconn = true;"
 }
 
 get_installed_extension_version() {
     local extname="$1"
     local dbname="$2"
-    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$dbname" -tAc \
+    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$dbname" -v ON_ERROR_STOP=1 -tAc \
         "SELECT extversion FROM pg_extension WHERE extname = '$extname';" 2>/dev/null | xargs
 }
 
 compare_versions() {
-    # Compares version strings (e.g., 0.3.0 vs 0.4.0)
-    if [ "$1" = "$2" ]; then
-        return 1 # Same
+    # Returns 0 if $1 < $2, 1 otherwise
+    local v1="$1"
+    local v2="$2"
+    if [ "$v1" = "$v2" ]; then return 1; fi
+    if [ "$(printf '%s\n' "$v1" "$v2" | sort -V | head -n1)" = "$v1" ]; then
+        return 0
     fi
-    # Split version by dots and compare segments
-    IFS=. read -r a1 a2 a3 <<<"$1"
-    IFS=. read -r b1 b2 b3 <<<"$2"
-    if [ "$a1" -lt "$b1" ]; then return 0; fi
-    if [ "$a1" -gt "$b1" ]; then return 1; fi
-    if [ "$a2" -lt "$b2" ]; then return 0; fi
-    if [ "$a2" -gt "$b2" ]; then return 1; fi
-    if [ "$a3" -lt "$b3" ]; then return 0; fi
     return 1
 }
 
+# --------- Main logic ----------
+
+upgrade_postgres_if_needed() {
+    CLUSTER_VERSION=$(get_pgdata_version)
+    IMAGE_VERSION="$PG_MAJOR_VERSION"
+
+    if [ "$CLUSTER_VERSION" != "$IMAGE_VERSION" ]; then
+        bashio::log.warning "Postgres data directory version is $CLUSTER_VERSION but image wants $IMAGE_VERSION. Running upgrade..."
+
+        export DATA_DIR="$PGDATA"
+        export BINARIES_DIR="/usr/lib/postgresql/$IMAGE_VERSION/bin"
+        export BACKUP_DIR="/config/backups"
+        export PSQL_VERSION="$IMAGE_VERSION"
+
+        apt-get update &>/dev/null
+        apt-get install -y procps rsync postgresql-$IMAGE_VERSION postgresql-$CLUSTER_VERSION &>/dev/null
+
+        TMP_SCRIPT=$(mktemp)
+        wget https://raw.githubusercontent.com/linkyard/postgres-upgrade/refs/heads/main/upgrade-postgres.sh -O "$TMP_SCRIPT"
+        chmod +x "$TMP_SCRIPT"
+        "$TMP_SCRIPT"
+        UPGRADE_STATUS=$?
+        if [ "$UPGRADE_STATUS" -ne 0 ]; then
+            bashio::log.error "Postgres major version upgrade failed. Aborting startup."
+            exit 1
+        fi
+
+        bashio::log.info "PostgreSQL major version upgrade complete."
+        RESTART_NEEDED=true
+
+        # Wait for server to come back online after upgrade
+        wait_for_postgres
+    else
+        bashio::log.info "PostgreSQL data directory version ($CLUSTER_VERSION) matches image version ($IMAGE_VERSION)."
+    fi
+}
+
 upgrade_extension_if_needed() {
     local extname="$1"
     if ! is_extension_available "$extname"; then
@@ -133,9 +158,9 @@ upgrade_extension_if_needed() {
             compare_versions "$installed_version" "$available_version"
             if [ $? -eq 0 ]; then
                 bashio::log.info "Upgrading $extname in $db from $installed_version to $available_version"
-                psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$db" -c "ALTER EXTENSION $extname UPDATE;" \
+                psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$db" -v ON_ERROR_STOP=1 -c "ALTER EXTENSION $extname UPDATE;" \
                     && RESTART_NEEDED=true \
-                    || bashio::log.warning "Failed to upgrade $extname in $db"
+                    || { bashio::log.error "Failed to upgrade $extname in $db. Aborting startup."; exit 1; }
             else
                 bashio::log.info "$extname in $db already at latest version ($installed_version)"
             fi
@@ -143,7 +168,9 @@ upgrade_extension_if_needed() {
     done
 }
 
-update_postgres
+# --------- Run logic ----------
+
+upgrade_postgres_if_needed
 
 upgrade_extension_if_needed "vectors"
 upgrade_extension_if_needed "vchord"