#!/usr/bin/env bash
set -euo pipefail

if [ "$EUID" -ne 0 ]; then
    echo "Error: This script must be run as root" >&2
    exit 1
fi

# Force a known locale; some OS X terminals set an invalid LC_CTYPE
# which crashes pg_upgradecluster
export LC_ALL=C.UTF-8
export LANG=C.UTF-8
export LANGUAGE=C.UTF-8

LATEST_SUPPORTED_VERSION=17
UPGRADE_TO=${1:-$LATEST_SUPPORTED_VERSION}
UPGRADE_FROM=$(crudini --get /etc/zulip/zulip.conf postgresql version)
ZULIP_PATH="$(dirname "$0")/../.."

if [ "$UPGRADE_TO" = "$UPGRADE_FROM" ]; then
    echo "Already running PostgreSQL $UPGRADE_TO!"
    exit 1
fi

if [[ "$UPGRADE_TO" -lt "$UPGRADE_FROM" ]]; then
    echo "Refusing to downgrade PostgreSQL $UPGRADE_FROM to $UPGRADE_TO!"
    exit 1
fi

if [[ "$UPGRADE_TO" -gt "$LATEST_SUPPORTED_VERSION" ]]; then
    echo "This version of Zulip does not support PostgreSQL $UPGRADE_TO."
    echo "You may need to upgrade Zulip before you can upgrade to PostgreSQL $UPGRADE_TO."
    exit 1
fi

# Verify that the version in /etc/zulip/zulip.conf is the version that
# Django actually stores its data in.  We can only do that if the
# database server is on the same host as the application server.
if [ -d /home/zulip/deployments/current ]; then
    DATA_IS_IN=$(
        su -s /usr/bin/env -- zulip \
            DJANGO_SETTINGS_MODULE=zproject.settings \
            uv run --directory=/home/zulip/deployments/current --no-sync \
            python -c 'from django.db import connection; print(connection.cursor().connection.server_version // 10000)'
    )

    if [ "$UPGRADE_FROM" != "$DATA_IS_IN" ]; then
        cat <<EOF

/etc/zulip/zulip.conf claims that Zulip is running PostgreSQL
$UPGRADE_FROM, but the server is connected to a PostgreSQL running
version $DATA_IS_IN.  Check the output from pg_lsclusters to verify
which clusters are running, and update /etc/zulip/zulip.conf to match.

In general, this results from manually upgrading PostgreSQL; you
should use this tool for all PostgreSQL upgrades.
EOF
        exit 1
    fi
fi

set -x

"$ZULIP_PATH"/scripts/lib/setup-apt-repo
apt-get install -y "postgresql-$UPGRADE_TO" "postgresql-client-$UPGRADE_TO"
drop_main_cluster="$(
    pg_lsclusters --json \
        | jq --arg pg_version "$UPGRADE_TO" -r \
            '.[] | select((.version|tostring == $ARGS.named.pg_version) and (.cluster == "main")) | .cluster'
)"
if [ -n "$drop_main_cluster" ]; then
    pg_dropcluster "$UPGRADE_TO" main --stop
fi

(
    # Two-stage application of Puppet; we apply the bare-bones
    # PostgreSQL configuration first, so that FTS will be configured
    # prior to the pg_upgradecluster.
    TEMP_CONF_DIR=$(mktemp -d)
    cp /etc/zulip/zulip.conf "$TEMP_CONF_DIR"
    ZULIP_CONF="${TEMP_CONF_DIR}/zulip.conf"
    crudini --set "$ZULIP_CONF" postgresql version "$UPGRADE_TO"

    if [ -f "/usr/share/postgresql/$UPGRADE_FROM/pgroonga_setup.sql.applied" ]; then
        # This file is intentionally blank, to trigger always running
        # `ALTER EXTENSION pgroonga UPDATE` and not `CREATE EXTENSION pgroonga`.
        touch "/usr/share/postgresql/$UPGRADE_TO/pgroonga_setup.sql.applied"
    fi

    "$ZULIP_PATH"/scripts/zulip-puppet-apply -f --config "$ZULIP_CONF" --tags postgresql_upgrade
    rm -rf "$TEMP_CONF_DIR"
)

# Capture the output so we know where the path to the post-upgrade scripts is
UPGRADE_LOG=$(mktemp "/var/log/zulip/upgrade-postgresql-$UPGRADE_FROM-$UPGRADE_TO.XXXXXXXXX.log")
pg_upgradecluster -v "$UPGRADE_TO" "$UPGRADE_FROM" main --method=upgrade --link --no-start | tee "$UPGRADE_LOG"
SCRIPTS_PATH=$(grep -o "/var/log/postgresql/pg_upgradecluster-$UPGRADE_FROM-$UPGRADE_TO-main.*" "$UPGRADE_LOG" || true)

# If the upgrade completed successfully, lock in the new version in
# our configuration immediately
crudini --set /etc/zulip/zulip.conf postgresql version "$UPGRADE_TO"

# Make sure the new PostgreSQL is running
pg_ctlcluster "$UPGRADE_TO" main start

# Update the statistics
su postgres -c "/usr/lib/postgresql/$UPGRADE_TO/bin/vacuumdb --all --analyze-only --jobs 10"

# Update extensions
if [ -n "$SCRIPTS_PATH" ] && [ -f "$SCRIPTS_PATH/update_extensions.sql" ]; then
    su postgres -c "psql $SCRIPTS_PATH/update_extensions.sql"
    rm "$SCRIPTS_PATH/update_extensions.sql"
fi

# Start the database up cleanly
"$ZULIP_PATH"/scripts/zulip-puppet-apply -f

# Drop the old data, binaries, and scripts
pg_dropcluster "$UPGRADE_FROM" main
apt remove -y "postgresql-$UPGRADE_FROM"
if [ -n "$SCRIPTS_PATH" ] && [ -f "$SCRIPTS_PATH/delete_old_cluster.sh" ]; then
    su postgres -c "$SCRIPTS_PATH/delete_old_cluster.sh"
    rm "$SCRIPTS_PATH/delete_old_cluster.sh"
fi

if [ -n "$SCRIPTS_PATH" ]; then
    if [ -n "$(ls -A "$SCRIPTS_PATH")" ]; then
        set +x
        echo
        echo
        echo ">>>>> pg_upgradecluster succeeded, but unexpected post-upgrade scripts"
        echo "      were found:"
        ls -A1 "$SCRIPTS_PATH"
    else
        rmdir "$SCRIPTS_PATH"
    fi
else
    set +x
    echo
    echo
    echo ">>>>> pg_upgradecluster succeeded, but post-upgrade scripts path could not"
    echo "      be parsed out!  Please read the pg_upgradecluster output to understand"
    echo "      the current status of your cluster:"
    echo "          $UPGRADE_LOG"
    echo "      and report this bug with the PostgreSQL $UPGRADE_FROM -> $UPGRADE_TO upgrade to:"
    echo "          https://github.com/zulip/zulip/issues"
    echo
    echo
fi
