#!/usr/bin/env bash
# surinch.com · PostgreSQL 备份可恢复性演练 v1.0
# 做法：pg_dump -Fc 备份 -> 恢复到一个新建的临时库 -> 逐表比对行数 -> 输出 CSV 证据 -> 删除临时库
# 只读源库；会在同一实例创建并删除名为 <源库>_restore_drill_<时间戳> 的临时库，请先确认磁盘空间与权限。
# 用法：PGHOST=... PGUSER=... ./pg-backup-restore-drill-v1.0.sh <源库名> [输出目录]
set -euo pipefail
SRC_DB="${1:?用法: $0 <源库名> [输出目录]}"
OUT_DIR="${2:-./drill-$(date +%Y%m%d-%H%M%S)}"
TS="$(date +%Y%m%d%H%M%S)"
DRILL_DB="${SRC_DB}_restore_drill_${TS}"
mkdir -p "$OUT_DIR"
DUMP="$OUT_DIR/${SRC_DB}.dump"

count_rows() { # 精确行数（count(*)），大表会慢，可在低峰执行
  psql -X -At -d "$1" -c "SELECT format('SELECT %L, count(*) FROM %I.%I', n.nspname||'.'||c.relname, n.nspname, c.relname)
    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%'
    ORDER BY 1" | while read -r q; do psql -X -At -F, -d "$1" -c "$q"; done | sort
}

t0=$(date +%s)
pg_dump -Fc -d "$SRC_DB" -f "$DUMP"
t1=$(date +%s)
createdb "$DRILL_DB"
trap 'dropdb --if-exists "$DRILL_DB" >/dev/null 2>&1 || true' EXIT
pg_restore --no-owner --exit-on-error -d "$DRILL_DB" "$DUMP"
t2=$(date +%s)

count_rows "$SRC_DB"   > "$OUT_DIR/src_counts.csv"
count_rows "$DRILL_DB" > "$OUT_DIR/restored_counts.csv"
echo "table,src_rows,restored_rows,match" > "$OUT_DIR/compare.csv"
join -t, -a1 -a2 -e MISSING -o 0,1.2,2.2 "$OUT_DIR/src_counts.csv" "$OUT_DIR/restored_counts.csv" \
  | awk -F, '{print $0","($2==$3?"OK":"DIFF")}' >> "$OUT_DIR/compare.csv"
DIFFS=$(grep -c ',DIFF$' "$OUT_DIR/compare.csv" || true)
TABLES=$(($(wc -l < "$OUT_DIR/compare.csv") - 1))
{
  echo "item,value"
  echo "source_db,$SRC_DB"
  echo "dump_bytes,$(stat -c %s "$DUMP")"
  echo "dump_seconds,$((t1-t0))"
  echo "restore_seconds,$((t2-t1))"
  echo "tables_compared,$TABLES"
  echo "tables_diff,$DIFFS"
  echo "result,$([ "$DIFFS" = 0 ] && echo PASS || echo FAIL)"
} > "$OUT_DIR/summary.csv"
cat "$OUT_DIR/summary.csv"
[ "$DIFFS" = 0 ]