#!/usr/bin/env bash
# pg-restore-check.sh: restore the newest dump from the bucket into a throwaway
# Postgres container and check that every table in the archive came back.
# From https://backupdrill.com/guides/pg-dump-to-s3 (MIT licensed).
#
# Needs Docker and the AWS CLI v2. Same variables as pg-dump-to-s3.sh, plus:
#   PG_MAJOR       major version of the pg_dump that made the dumps, which
#                  should be the same as your server's (default: 17)
#   AGE_IDENTITY   path to the age private key, if the dumps are encrypted
#   ROLES          space-separated roles your policies or grants name, e.g.
#                  "app_reader app_writer"; created (NOLOGIN) before restoring
set -euo pipefail

: "${S3_BUCKET:?S3_BUCKET is not set}"
prefix="${S3_PREFIX:-postgres}"
major="${PG_MAJOR:-17}"
container="pg-restore-check-$$"

workdir="$(mktemp -d)"
cleanup() {
  # --volumes: the postgres image keeps its data in an anonymous volume, which
  # would otherwise outlive the container with a plain-text copy of the database
  docker rm --force --volumes "$container" > /dev/null 2>&1 || true
  rm -rf "$workdir"
}
trap cleanup EXIT
umask 077

# 1. Newest dump in the prefix. The timestamped names sort by date.
latest="$(aws s3 ls "s3://$S3_BUCKET/$prefix/" \
  | awk '{print $4}' | grep -E '\.dump(\.age)?$' | sort | tail -n 1 || true)"
[ -n "$latest" ] || { echo "no dumps under s3://$S3_BUCKET/$prefix/"; exit 1; }
aws s3 cp "s3://$S3_BUCKET/$prefix/$latest" "$workdir/$latest" --only-show-errors
aws s3 cp "s3://$S3_BUCKET/$prefix/$latest.sha256" "$workdir/$latest.sha256" --only-show-errors

# 2. Checksum, so a damaged download is not mistaken for a bad backup.
if command -v sha256sum > /dev/null; then
  (cd "$workdir" && sha256sum --check --quiet "$latest.sha256")
else
  (cd "$workdir" && shasum -a 256 --check --quiet "$latest.sha256")
fi

dump="$workdir/latest.dump"
if [ "${latest%.age}" != "$latest" ]; then
  : "${AGE_IDENTITY:?the dump is encrypted: set AGE_IDENTITY}"
  age --decrypt --identity "$AGE_IDENTITY" --output "$dump" "$workdir/$latest"
else
  mv "$workdir/$latest" "$dump"
fi

# 3. A throwaway server of the same major version, with no network at all:
#    it holds a copy of your data, and docker exec is all the checks need.
#    It only listens on TCP once initialisation has finished, so wait for a
#    TCP connection, not pg_isready.
docker run --detach --name "$container" --network none \
  --env POSTGRES_PASSWORD=throwaway "postgres:$major" > /dev/null
for attempt in $(seq 1 120); do
  if docker exec "$container" psql -h 127.0.0.1 -U postgres -c 'select 1' > /dev/null 2>&1; then
    break
  fi
  if [ "$attempt" -eq 120 ]; then
    echo "FAIL: the throwaway Postgres did not start within 2 minutes:"
    docker logs --tail 20 "$container"
    exit 1
  fi
  sleep 1
done
docker cp "$dump" "$container:/tmp/latest.dump" > /dev/null

# --no-owner and --no-privileges drop ownership and grants, but a row-level
# security policy that names a role still needs that role to exist.
for role in ${ROLES:-}; do
  docker exec "$container" psql -h 127.0.0.1 -U postgres -q \
    -c "create role \"$role\" nologin" > /dev/null
done

# 4. Restore. --exit-on-error turns the first failure into a failed check
#    instead of a warning count at the end; --no-tablespaces because the
#    throwaway server has none of the source's custom tablespaces.
docker exec "$container" pg_restore -h 127.0.0.1 -U postgres --dbname=postgres \
  --no-owner --no-privileges --no-tablespaces --exit-on-error /tmp/latest.dump

# 5. Every table that has data in the archive must exist after the restore.
#    A TOC line ends in "TABLE DATA <schema> <table> <owner>"; the table name
#    is everything between schema and owner, so table names with spaces
#    survive. Schema or owner names with spaces are not handled.
docker exec "$container" pg_restore --list /tmp/latest.dump \
  | awk '/ TABLE DATA / {
      sub(/^.* TABLE DATA /, "")
      schema = $1; owner = $NF
      print schema "." substr($0, length(schema) + 2, length($0) - length(schema) - length(owner) - 2)
    }' | sort > "$workdir/expected.txt"
docker exec "$container" psql -h 127.0.0.1 -U postgres -At -c "
  select table_schema || '.' || table_name
  from information_schema.tables
  where table_type = 'BASE TABLE'
    and table_schema not in ('pg_catalog', 'information_schema')" \
  | sort > "$workdir/restored.txt"
missing="$(comm -23 "$workdir/expected.txt" "$workdir/restored.txt")"
if [ -n "$missing" ]; then
  echo "FAIL: tables in the dump but not in the restored database:"
  echo "$missing"
  exit 1
fi

# 6. Row counts, so an empty table stands out.
docker exec "$container" psql -h 127.0.0.1 -U postgres -At -F ' ' -c "
  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 not in ('pg_catalog', 'information_schema')
  order by 1"
echo "PASS: $latest restored, $(wc -l < "$workdir/expected.txt" | tr -d ' ') tables present"