#!/usr/bin/env bash
# self-hosted-supabase-restore-drill.sh: restore the newest off-site backup into a
# throwaway Postgres container and check it against the live stack.
# From https://backupdrill.com/guides/self-hosted-supabase-backup (MIT licensed).
#
# Run it from the directory with docker-compose.yml; the live stack is only read.
# Needs rclone (remote "offsite") and getfattr (Ubuntu: apt install attr).
# Exits with 1 when a check fails.
set -euo pipefail
IMAGE=supabase/postgres:17.6.1.136  # the db image tag from your docker-compose.yml

umask 077
DRILL=$(mktemp -d)
# Names unique to this run, so cleanup never touches another run's container
NAME="supabase-drill-$(basename "$DRILL" | tr -c 'A-Za-z0-9\n' '-')"
cleanup() {
  docker rm -f "$NAME" > /dev/null 2>&1 || true
  docker volume rm "$NAME-config" > /dev/null 2>&1 || true
  rm -rf "$DRILL"
}
trap cleanup EXIT

LATEST=$(rclone lsf --dirs-only offsite: | sort | tail -n 1)
LATEST=${LATEST%/}
echo "drilling backup $LATEST"
rclone copy "offsite:$LATEST" "$DRILL"
if ! (cd "$DRILL" && sha256sum --check --quiet SHA256SUMS); then echo "DRILL FAILED: checksums"; exit 1; fi

# A config volume with the backed-up Vault key and conf.d, owned by the image's postgres user
docker volume create "$NAME-config" > /dev/null
docker run --rm -i -v "$NAME-config:/etc/postgresql-custom" --entrypoint sh "$IMAGE" -c \
  'cd /etc/postgresql-custom && tar -xf - && chown -R postgres:postgres pgsodium_root.key conf.d && chmod 640 pgsodium_root.key' \
  < "$DRILL/db-config.tar"

# Same image, empty init directory: the dump brings every Supabase schema with it.
# No network and no pg_cron jobs, so restored jobs and pg_net calls can't reach anything.
docker run -d --name "$NAME" --network none --tmpfs /docker-entrypoint-initdb.d \
  -v "$NAME-config:/etc/postgresql-custom" -e POSTGRES_PASSWORD=drill-only \
  "$IMAGE" postgres -c config_file=/etc/postgresql/postgresql.conf \
  -c cron.launch_active_jobs=off > /dev/null
for _ in $(seq 90); do
  docker exec "$NAME" pg_isready -q -h localhost -U supabase_admin -d postgres 2> /dev/null && break
  if [ "$(docker inspect -f '{{.State.Running}}' "$NAME")" != true ]; then break; fi
  sleep 2
done
if ! docker exec "$NAME" pg_isready -q -h localhost -U supabase_admin -d postgres 2> /dev/null; then
  docker logs --tail 30 "$NAME"; echo "DRILL FAILED: the drill database did not start"; exit 1
fi

FAILED=0
# Roles first. One error is expected: the image already created supabase_admin.
if ! docker exec -i "$NAME" psql -q -U supabase_admin -d template1 \
    < "$DRILL/globals.sql" 2> "$DRILL/globals.log"; then
  echo "restoring roles failed:"; cat "$DRILL/globals.log"; FAILED=1
else
  UNEXPECTED=$(grep 'ERROR' "$DRILL/globals.log" | grep -v 'role "supabase_admin" already exists' || true)
  if [ -n "$UNEXPECTED" ]; then echo "unexpected errors restoring roles:"; echo "$UNEXPECTED"; FAILED=1; fi
fi

# Then each database. postgres.dump may fail exactly once, on CREATE DATABASE,
# because the image created that database too; pg_restore then exits with 1.
for dump in "$DRILL"/*.dump; do
  log="$dump.log"
  status=0
  docker exec -i "$NAME" pg_restore -U supabase_admin -d template1 --create < "$dump" 2> "$log" || status=$?
  errors=$(grep -c '^pg_restore: error:' "$log" || true)
  if [ "$status" -eq 0 ] && [ "$errors" -eq 0 ]; then
    continue
  elif [ "$status" -eq 1 ] && [ "$errors" -eq 1 ] && [ "$(basename "$dump")" = postgres.dump ] \
      && grep -q 'database "postgres" already exists' "$log"; then
    continue
  fi
  echo "restoring $(basename "$dump") failed (exit $status):"; cat "$log"; FAILED=1
done
echo "restore: $([ "$FAILED" -eq 0 ] && echo "only the 2 expected errors" || echo "FAILED")"

# Exact row counts in public, auth and storage, live versus restored
COUNTS="select table_schema || '.' || table_name, (xpath('/row/n/text()', query_to_xml(
  format('select count(*) as n from %I.%I', table_schema, table_name), false, true, '')))[1]::text
  from information_schema.tables
  where table_type = 'BASE TABLE' and table_schema in ('public', 'auth', 'storage') order by 1"
docker compose exec -T db psql -U supabase_admin -d postgres -AtF ' ' -c "$COUNTS" > "$DRILL/counts-live.txt"
docker exec "$NAME" psql -U supabase_admin -d postgres -AtF ' ' -c "$COUNTS" > "$DRILL/counts-restored.txt"
if diff "$DRILL/counts-live.txt" "$DRILL/counts-restored.txt"; then
  echo "row counts match: $(wc -l < "$DRILL/counts-restored.txt") tables"
else
  echo "row counts differ (< live, > restored): check each difference against writes since the backup"
fi

# Vault secrets decrypt only with the original key
if docker exec "$NAME" psql -U supabase_admin -d postgres -Atc \
    "select count(*) from vault.decrypted_secrets where decrypted_secret is not null" > "$DRILL/vault.txt"; then
  echo "vault secrets readable: $(cat "$DRILL/vault.txt")"
else
  echo "vault secrets do not decrypt: wrong or missing pgsodium_root.key in db-config.tar"; FAILED=1
fi

# Every object in storage.objects has its file, with the right size and content type
tar --xattrs --xattrs-include='user.*' -xzf "$DRILL/storage.tar.gz" -C "$DRILL"
ROOT="$DRILL/storage/$(cat "$DRILL/storage-path.txt")"
docker exec "$NAME" psql -U supabase_admin -d postgres -AtF $'\t' -c \
  "select bucket_id || '/' || name || coalesce('/' || version, ''), metadata->>'size', metadata->>'mimetype' from storage.objects" \
  > "$DRILL/objects.tsv"
while IFS=$'\t' read -r key size type; do
  file="$ROOT/$key"
  if [ ! -f "$file" ]; then echo "missing file: $key"
  elif [ "$(stat -c %s "$file")" != "$size" ]; then echo "wrong size: $key"
  elif [ "$(getfattr --only-values -n user.supabase.content-type "$file" 2> /dev/null)" != "$type" ]; then
    echo "content type lost: $key"
  fi
done < "$DRILL/objects.tsv" > "$DRILL/file-problems.txt"
if [ -s "$DRILL/file-problems.txt" ]; then cat "$DRILL/file-problems.txt"; FAILED=1; fi
echo "storage objects checked: $(wc -l < "$DRILL/objects.tsv"), problems: $(wc -l < "$DRILL/file-problems.txt")"

if [ "$FAILED" -eq 0 ]; then echo "DRILL PASSED"; else echo "DRILL FAILED"; fi
exit "$FAILED"