#!/bin/bash
# Weekly restore drill for the postgresql01 cluster (rustinion, forgejodb,
# ssherpa and every other database on 10.104.130.203).
#
# Deployed INSIDE the postgresql01 LXD container as
# /usr/sbin/ctg-pg-restore-drill and run by the host timer
# pg-restore-drill.timer (see ../README.md). It never touches the live
# cluster: it restores the latest pgBackRest backup from the off-site
# Backblaze B2 repo (the same encrypted repo the nightly cron writes to)
# into a scratch data dir, replays archived WAL to the end of the
# archive, starts a throwaway postmaster on a unix socket only (no TCP,
# archive_mode=off so it can never push anything back into the repo),
# runs integrity checks, stops it and deletes the scratch dir.
#
# The live cluster is only ever READ (catalog queries for comparison).
#
# Prints human-readable progress on stderr and exactly one JSON result
# object as the last line of stdout. Exit status 0 = drill passed.
set -uo pipefail

DRILL_DIR=/var/tmp/pg-restore-drill
DATA="$DRILL_DIR/data"
SOCK="$DRILL_DIR/sock"
PGBIN=/usr/lib/postgresql/18/bin
PORT=55432
# Restored data must be at most this old relative to the drill start
# (measured by the canary device, which checks in every 15 minutes, so its
# last_seen_at is a direct measure of how far WAL replay got).
MAX_RESTORE_LAG_S=${MAX_RESTORE_LAG_S:-21600}
# Database the recency check reads (devices.last_seen_at).
RECENCY_DB=rustinion
# Backblaze B2 sometimes serves one particular object truncated for
# minutes at a time: HTTP 200 with the right Content-Length, then the
# connection closes early (pgbackrest: "unexpected EOF reading HTTP
# content"). Root-caused 2026-10-16: base/16718/4154.gz of 20261011-030003F
# failed 15/15 GETs from postgresql01 AND from a different network (curl,
# 0 of 112 bytes), then read fine again ~15 minutes later; a full read of
# all 119k backup objects found nothing else. The data is intact, B2 just
# can't serve it for a while. pgbackrest's default (2 retries, 15 s) gives
# up after ~15 s, so retry each file 10 times 60 s apart (~9 min per file),
# and count the retries (debug file log) so the drill result says when B2
# misbehaved instead of hiding it. See README "Restoring from B2".
PGBACKREST_RESTORE_OPTS=${PGBACKREST_RESTORE_OPTS:---job-retry=10 --job-retry-interval=60}

STARTED_AT=$(date -u +%Y-%m-%dT%H:%M:%SZ)
START_EPOCH=$(date +%s)
CHECKS=()
FAILED=0
ERROR=""
BACKUP_LABEL=""
LAG_S=""

log() { echo "[$(date -u +%H:%M:%S)] $*" >&2; }
check() { # name ok detail
    local esc
    esc=$(printf '%s' "$3" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')
    CHECKS+=("{\"name\":\"$1\",\"ok\":$2,\"detail\":$esc}")
    if [ "$2" != "true" ]; then FAILED=1; log "CHECK FAILED: $1: $3"; else log "check ok: $1: $3"; fi
}
live_sql() { sudo -u postgres psql -X -q -At -p 5432 -d "$1" -c "$2"; }
drill_sql() { sudo -u postgres psql -X -q -At -h "$SOCK" -p "$PORT" -d "$1" -c "$2"; }

stop_drill_server() {
    if [ -f "$DATA/postmaster.pid" ]; then
        sudo -u postgres "$PGBIN/pg_ctl" -D "$DATA" -m immediate -w -t 120 stop >/dev/null 2>&1 || true
    fi
}
cleanup() {
    stop_drill_server
    rm -rf "$DRILL_DIR"
}

finish() {
    local finished duration ok
    finished=$(date -u +%Y-%m-%dT%H:%M:%SZ)
    duration=$(( $(date +%s) - START_EPOCH ))
    ok=true
    [ "$FAILED" = 0 ] || ok=false
    local checks_json
    checks_json=$(IFS=,; echo "${CHECKS[*]:-}")
    local err_json
    err_json=$(printf '%s' "$ERROR" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')
    cleanup
    echo "{\"ok\":$ok,\"started_at\":\"$STARTED_AT\",\"finished_at\":\"$finished\",\"duration_s\":$duration,\"backup_label\":\"$BACKUP_LABEL\",\"restore_lag_s\":${LAG_S:-null},\"error\":$err_json,\"checks\":[${checks_json}]}"
    [ "$ok" = true ]
    exit $?
}
fail() { ERROR="$*"; FAILED=1; log "FATAL: $*"; finish; }

trap 'fail "interrupted"' INT TERM

# A previous drill killed mid-run leaves its scratch cluster behind.
cleanup
mkdir -p "$DATA" "$SOCK"
chown -R postgres:postgres "$DRILL_DIR"
chmod 700 "$DATA"

BACKUP_LABEL=$(sudo -u postgres pgbackrest --stanza=main --output=json info 2>/dev/null \
    | python3 -c 'import json,sys; s=json.load(sys.stdin)[0]; print(s["backup"][-1]["label"] if s.get("backup") else "")') \
    || fail "pgbackrest info failed"
[ -n "$BACKUP_LABEL" ] || fail "repository has no backups"
log "restoring latest backup set $BACKUP_LABEL from the repo into $DATA"

# --archive-mode=off: the restored cluster must never archive-push into
# the production repo. Default recovery type = replay all archived WAL.
# Full pgbackrest output goes to the drill log (stderr); the error lines
# themselves go into the result so the alert says what actually broke.
RESTORE_LOG="$DRILL_DIR/pgbackrest-restore.log"
if ! sudo -u postgres pgbackrest --stanza=main --pg1-path="$DATA" --archive-mode=off \
        --process-max=2 $PGBACKREST_RESTORE_OPTS --log-level-console=detail \
        --log-level-file=debug --log-subprocess --log-path="$DRILL_DIR/pgbackrest-log" \
        restore >"$RESTORE_LOG" 2>&1; then
    cat "$RESTORE_LOG" >&2
    grep -h 'retry FileReadError\|protocolServerError' "$DRILL_DIR"/pgbackrest-log/*.log 2>/dev/null | cut -c1-300 >&2
    fail "pgbackrest restore failed: $(grep -A3 ' ERROR: ' "$RESTORE_LOG" | tr -s ' \n' ' ' | cut -c1-600)"
fi
grep -v ' restore file ' "$RESTORE_LOG" >&2
# Transient B2 read failures that a retry absorbed: reported, not fatal.
RETRIES=$(cat "$DRILL_DIR"/pgbackrest-log/*.log 2>/dev/null | grep -c 'retry FileReadError' || true)
RETRY_FILES=$(grep -h -B30 'retry FileReadError' "$DRILL_DIR"/pgbackrest-log/*.log 2>/dev/null \
    | grep -o 'repoFile: {"[^"]*"' | sed 's/.*BACKUP>\///; s/"$//' | sort -u | head -5 | tr '\n' ' ')
check repo_read_retries true "${RETRIES:-0} B2 read retries during restore${RETRY_FILES:+ (files: $RETRY_FILES)}"
rm -rf "$DRILL_DIR/pgbackrest-log"
check restore true "pgbackrest restore of $BACKUP_LABEL completed ($(du -sh "$DATA" | cut -f1))"

# The Debian layout keeps postgresql.conf/pg_hba.conf in /etc, so the
# restored data dir has none: give the scratch server its own minimal set.
echo "local all postgres peer" > "$DRILL_DIR/pg_hba.conf"
: > "$DRILL_DIR/pg_ident.conf"
: > "$DRILL_DIR/postgresql.conf"
chown postgres:postgres "$DRILL_DIR"/*.conf

# hot_standby=off: no connections until recovery has finished and the
# cluster promoted, and no "setting lower than on the primary" checks.
OPTS="-c config_file=$DRILL_DIR/postgresql.conf -c hba_file=$DRILL_DIR/pg_hba.conf"
OPTS="$OPTS -c ident_file=$DRILL_DIR/pg_ident.conf -c port=$PORT -c listen_addresses="
OPTS="$OPTS -c unix_socket_directories=$SOCK -c archive_mode=off -c hot_standby=off"
OPTS="$OPTS -c max_connections=20 -c shared_buffers=256MB -c max_wal_senders=0"
OPTS="$OPTS -c wal_level=replica -c max_logical_replication_workers=0 -c track_commit_timestamp=on"
OPTS="$OPTS -c cluster_name=restore-drill -c logging_collector=off -c autovacuum=off"
if ! sudo -u postgres "$PGBIN/pg_ctl" -D "$DATA" -l "$DRILL_DIR/server.log" -W -o "$OPTS" start >&2; then
    fail "scratch postmaster failed to start: $(tail -5 "$DRILL_DIR/server.log" 2>/dev/null)"
fi

log "waiting for WAL replay + promotion"
deadline=$(( $(date +%s) + 14400 ))
while :; do
    if [ "$(drill_sql postgres 'select pg_is_in_recovery()' 2>/dev/null)" = "f" ]; then break; fi
    if ! sudo -u postgres "$PGBIN/pg_ctl" -D "$DATA" status >/dev/null 2>&1; then
        fail "scratch postmaster exited during recovery: $(tail -5 "$DRILL_DIR/server.log" 2>/dev/null)"
    fi
    [ "$(date +%s)" -lt "$deadline" ] || fail "recovery did not finish within 4h"
    sleep 10
done
check recovery true "WAL replay finished and cluster promoted ($(grep -c 'restored log file' "$DRILL_DIR/server.log" 2>/dev/null) WAL segments replayed)"

# Same set of databases as live.
LIVE_DBS=$(live_sql postgres "select datname from pg_database where not datistemplate order by 1" | tr '\n' ' ')
DRILL_DBS=$(drill_sql postgres "select datname from pg_database where not datistemplate order by 1" | tr '\n' ' ')
if [ "$LIVE_DBS" = "$DRILL_DBS" ]; then
    check databases true "all $(echo $LIVE_DBS | wc -w) databases present: $LIVE_DBS"
else
    check databases false "live=[$LIVE_DBS] restored=[$DRILL_DBS]"
fi

# Every database: user-table count matches live (a migration deployed
# after the backup can legitimately add a few) and every table is readable.
TABLES_Q="select count(*) from pg_class c join pg_namespace n on n.oid=c.relnamespace where c.relkind in ('r','p') and n.nspname not in ('pg_catalog','information_schema') and n.nspname not like 'pg_toast%'"
for db in $DRILL_DBS; do
    live_n=$(live_sql "$db" "$TABLES_Q" 2>/dev/null || echo "?")
    drill_n=$(drill_sql "$db" "$TABLES_Q" 2>/dev/null || echo "?")
    if [ "$drill_n" = "?" ] || [ "$live_n" = "?" ]; then
        check "tables:$db" false "could not count tables (live=$live_n restored=$drill_n)"
    elif [ "$drill_n" -le 0 ] && [ "$live_n" -gt 0 ]; then
        check "tables:$db" false "restored copy has no tables, live has $live_n"
    elif [ $(( live_n - drill_n )) -gt 5 ] || [ $(( drill_n - live_n )) -gt 5 ]; then
        check "tables:$db" false "table count differs: live=$live_n restored=$drill_n"
    else
        check "tables:$db" true "live=$live_n restored=$drill_n"
    fi
done

# Row-level spot checks on the three application databases this drill
# exists for.
n=$(drill_sql rustinion "select count(*) from devices" 2>/dev/null || echo -1)
[ "$n" -gt 0 ] && check rows:rustinion.devices true "$n devices" || check rows:rustinion.devices false "devices count=$n"
n=$(drill_sql rustinion "select count(*) from tenants" 2>/dev/null || echo -1)
[ "$n" -gt 0 ] && check rows:rustinion.tenants true "$n tenants" || check rows:rustinion.tenants false "tenants count=$n"
lm=$(live_sql rustinion "select max(version) from _sqlx_migrations" 2>/dev/null || echo "?")
dm=$(drill_sql rustinion "select max(version) from _sqlx_migrations" 2>/dev/null || echo "?")
if [ "$dm" != "?" ] && [ -n "$dm" ] && [ "$lm" != "?" ] && [ $(( lm - dm )) -le 3 ]; then
    check rustinion.migrations true "restored schema at migration $dm (live $lm)"
else
    check rustinion.migrations false "restored migration=$dm live=$lm"
fi
n=$(drill_sql forgejodb "select count(*) from repository" 2>/dev/null || echo -1)
[ "$n" -gt 0 ] && check rows:forgejodb.repository true "$n repositories" || check rows:forgejodb.repository false "repository count=$n"
n=$(drill_sql forgejodb "select count(*) from \"user\"" 2>/dev/null || echo -1)
[ "$n" -gt 0 ] && check rows:forgejodb.user true "$n users" || check rows:forgejodb.user false "user count=$n"
n=$(drill_sql ssherpa "$TABLES_Q" 2>/dev/null || echo -1)
[ "$n" -gt 0 ] && check rows:ssherpa true "$n tables readable" || check rows:ssherpa false "table count=$n"

# How far WAL replay got: the production canary device checks in every
# 15 minutes, so the newest devices.last_seen_at is the restore point.
newest=$(drill_sql "$RECENCY_DB" "select coalesce(extract(epoch from max(last_seen_at))::bigint, 0) from devices" 2>/dev/null || echo 0)
if [ "$newest" -gt 0 ]; then
    LAG_S=$(( START_EPOCH - newest ))
    [ "$LAG_S" -lt 0 ] && LAG_S=0
    if [ "$LAG_S" -le "$MAX_RESTORE_LAG_S" ]; then
        check restore_point true "restored data is ${LAG_S}s behind the drill start"
    else
        check restore_point false "restored data is ${LAG_S}s behind the drill start (max ${MAX_RESTORE_LAG_S}s): WAL archive gap?"
    fi
else
    check restore_point false "could not read devices.last_seen_at from the restored copy"
fi

# Physical integrity of the application databases' heap + btree indexes.
for db in rustinion ssherpa forgejodb; do
    if out=$(sudo -u postgres "$PGBIN/pg_amcheck" -h "$SOCK" -p "$PORT" --install-missing --heapallindexed -d "$db" 2>&1); then
        check "amcheck:$db" true "heap and index checks clean"
    else
        check "amcheck:$db" false "$(echo "$out" | tail -5)"
    fi
done

finish
