# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Restore test and time to restore.

A PostgreSQL database with WAL archiving and about 1 GB of data: a base
backup, 100,000 more rows, a checksum, then DROP TABLE. A point-in-time
restore to just before the drop: rows, checksum, restore time and GB/s.
Then a restore from the same backup with the WAL archive missing.
Runs its own containers with the docker CLI (needs Docker).
"""
import argparse, json, subprocess, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=12_000_000, help="Rows in the bulk table (about 1 GB at the default).")
p.add_argument("--image", default="postgres:17.6", help="PostgreSQL image.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
SRC, DST = "lab-pitr-src", "lab-pitr-restore"
VOLS = ["lab-pitr-archive", "lab-pitr-backup", "lab-pitr-empty"]

def sh(*a, check=True): return subprocess.run(["docker", *a], capture_output=True, text=True, check=check)
def cleanup():
    sh("rm", "-f", SRC, DST, check=False)
    for v in VOLS: sh("volume", "rm", "-f", v, check=False)
def wait_ready(port, timeout=600):
    t = time.time()
    while time.time() - t < timeout:
        try:
            c = psycopg.connect(f"postgresql://postgres:lab@localhost:{port}/postgres", autocommit=True, connect_timeout=2)
            if not c.execute("SELECT pg_is_in_recovery()").fetchone()[0]: return c
            c.close()
        except Exception: pass
        time.sleep(1)
    raise TimeoutError("database not ready")

cleanup()
for v in VOLS: sh("volume", "create", v)
sh("run", "-d", "--name", SRC, "-e", "POSTGRES_PASSWORD=lab", "-p", "55510:5432",
   "-v", "lab-pitr-archive:/archive", "-v", "lab-pitr-backup:/backup", args.image,
   "-c", "wal_level=replica", "-c", "archive_mode=on", "-c", "archive_command=cp %p /archive/%f", "-c", "max_wal_size=4GB")
time.sleep(3); sh("exec", "-u", "root", SRC, "chown", "-R", "postgres", "/archive", "/backup")
con = wait_ready(55510)
con.execute(f"CREATE TABLE bulk AS SELECT g AS id, md5(g::text) AS a, md5((g * 7)::text) AS b FROM generate_series(1, {args.rows}) g")
con.execute("CREATE TABLE orders (id int PRIMARY KEY, amount numeric)")
con.execute("INSERT INTO orders SELECT g, g % 997 FROM generate_series(1, 200000) g")
db_gb = con.execute("SELECT pg_database_size('postgres')").fetchone()[0] / 1e9
sh("exec", "-u", "postgres", SRC, "pg_basebackup", "-D", "/backup/base", "-X", "stream", "-c", "fast")
con.execute("INSERT INTO orders SELECT g, g % 997 FROM generate_series(200001, 300000) g")       # 100,000 more after the backup
CHECK = "SELECT count(*), md5(string_agg(o::text, '' ORDER BY id)) FROM orders o"
before_rows, before_md5 = con.execute(CHECK).fetchone()
target = con.execute("SELECT now()").fetchone()[0]
time.sleep(1.5)
con.execute("DROP TABLE orders")
last = con.execute("SELECT pg_walfile_name(pg_switch_wal())").fetchone()[0]
t0 = time.time()                                                       # wait until the archive really holds that segment
while (con.execute("SELECT coalesce(last_archived_wal, '') FROM pg_stat_archiver").fetchone()[0] or "") < last:
    if time.time() - t0 > 300: raise TimeoutError("WAL archiving did not catch up")
    time.sleep(1)
con.close(); sh("stop", "-t", "60", SRC)

def restore(archive_vol):
    sh("rm", "-f", DST, check=False)
    script = ("cp -r /backup/base /tmp/pgdata && chmod 700 /tmp/pgdata && "
              f"echo \"restore_command = 'cp /archive/%f %p'\" >> /tmp/pgdata/postgresql.auto.conf && "
              f"echo \"recovery_target_time = '{target.isoformat()}'\" >> /tmp/pgdata/postgresql.auto.conf && "
              "echo \"recovery_target_action = 'promote'\" >> /tmp/pgdata/postgresql.auto.conf && "
              "touch /tmp/pgdata/recovery.signal && exec postgres -D /tmp/pgdata")
    t = time.time()
    sh("run", "-d", "--name", DST, "-u", "postgres", "-p", "55511:5432", "-v", "lab-pitr-backup:/backup", "-v", f"{archive_vol}:/archive",
       "--entrypoint", "bash", args.image, "-c", script)
    try:
        c = wait_ready(55511, timeout=300); secs = time.time() - t
        rows_, md5_ = c.execute(CHECK).fetchone(); c.close()
        return {"rows": rows_, "md5_match": "yes" if md5_ == before_md5 else "no", "seconds": secs, "error": ""}
    except Exception:
        log = sh("logs", DST, check=False).stderr.strip().splitlines()
        fatal = next((l for l in reversed(log) if "FATAL" in l or "recovery ended" in l), log[-1] if log else "")
        return {"rows": None, "md5_match": "no", "seconds": time.time() - t, "error": fatal.split("FATAL:")[-1].strip()[:160]}

ok = restore("lab-pitr-archive")
gbps = db_gb / ok["seconds"] if ok["seconds"] else None
print(f"database {db_gb:.2f} GB; orders before the drop: {before_rows:,} rows")
print(f"point-in-time restore: {ok['rows']} rows, checksum matches: {ok['md5_match']}, {ok['seconds']:.0f} s ({gbps:.3f} GB/s)")
hours_10tb = 10_000 / gbps / 3600 if gbps else None
print(f"at this rate 10 TB would take {hours_10tb:.1f} hours (500 MB/s would take 5.6)")
nowal = restore("lab-pitr-empty")
print(f"restore with the WAL archive missing: {nowal['error'] or str(nowal['rows']) + ' rows'}")
cleanup()

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"restores": [{"restore": "with the WAL archive", **ok}, {"restore": "WAL archive missing", **nowal}],
                   "summary": [{"db_gb": db_gb, "rows_before": before_rows, "rows_after": ok["rows"], "md5_match": ok["md5_match"],
                                "restore_s": ok["seconds"], "gbps": gbps, "hours_10tb": hours_10tb, "nowal_error": nowal["error"] or "none"}]}}, f)