"""Testing a database dump without touching this install. The dump is loaded into a separate, temporary database on the same Postgres server, using the same steps a real restore takes. Every table is counted there and compared, the copy is brought up to this version, and the copy is then dropped. The live database is only read, to show its counts alongside. Every dump records how many rows each table held, counted in the same snapshot pg_dump dumped from, so the comparison is exact. For a dump in the bucket the counts sit beside it as `.counts.json`; for a dump on this server they are in its manifest file. A dump taken before counts were recorded can still be tested: it is then compared with the live database, where a difference is what changed after the dump was taken, not a fault. The temporary database needs about as much free space on the database's disk as the live one takes, and only for the length of the test. """ from __future__ import annotations import json import tempfile from collections.abc import Callable from pathlib import Path from sqlalchemy.engine import make_url from . import backup_restore, backup_service from .backup_s3 import ObjectMissing, S3Target, get_file, put_file SCRATCH_SUFFIX = "_restore_test" def counts_key(dump_key: str) -> str: """Where a dump's row counts sit in the bucket, beside the dump.""" return dump_key.removesuffix(".dump") + ".counts.json" def store_counts(target: S3Target, dump_key: str, counts: dict) -> None: with tempfile.TemporaryDirectory() as d: path = Path(d) / "counts.json" path.write_text(json.dumps(counts)) put_file(target, counts_key(dump_key), path, content_type="application/json") def fetch_counts(target: S3Target, dump_key: str) -> dict | None: """The counts recorded with a dump in the bucket, or None for a dump that has none.""" with tempfile.TemporaryDirectory() as d: path = Path(d) / "counts.json" try: get_file(target, counts_key(dump_key), path) except ObjectMissing: return None return json.loads(path.read_text()) def local_counts(manifest_path: str | None) -> dict | None: """The counts in a dump's manifest file on this server, or None.""" try: return json.loads(Path(manifest_path).read_text()).get("counts") except (TypeError, OSError, ValueError): return None def scratch_name(db_url: str) -> str: """The temporary database's name: the live one's, plus SCRATCH_SUFFIX.""" return f"{make_url(db_url).database}{SCRATCH_SUFFIX}" def scratch_url(db_url: str) -> str: return make_url(db_url).set(database=scratch_name(db_url)).render_as_string( hide_password=False) def _on_server(db_url: str, statement) -> None: """Run a CREATE or DROP DATABASE, which can't run inside a transaction.""" import psycopg with psycopg.connect(backup_service._libpq_url(db_url), autocommit=True) as conn: conn.execute(statement) def drop_scratch(db_url: str) -> None: """Drop the temporary database if it is there, ending any connection to it.""" from psycopg import sql _on_server(db_url, sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format( sql.Identifier(scratch_name(db_url)))) def _create_scratch(db_url: str) -> None: from psycopg import sql _on_server(db_url, sql.SQL("CREATE DATABASE {} TEMPLATE template0").format( sql.Identifier(scratch_name(db_url)))) def trial_restore(*, db_url: str, dump_path: Path, recorded: dict | None, note: Callable[[str], None] = lambda phase: None) -> dict: """Load the dump at `dump_path` into the temporary database, compare it, bring it up to this version, and drop it. `note` is told each phase as it starts: "loading", "comparing", "upgrading". Raises RestoreFailed (or S3Error from the upgrade) when the dump doesn't load or can't be brought up; the temporary database is dropped either way.""" scratch = scratch_url(db_url) if make_url(scratch).database == make_url(db_url).database: raise ValueError("The test database would be the live one.") drop_scratch(db_url) # one left by a test cut off by a restart _create_scratch(db_url) try: note("loading") backup_service.restore_db(db_url=scratch, sql_path=dump_path) note("comparing") restored = backup_service.read_counts(scratch) live = backup_service.read_counts(db_url) note("upgrading") backup_restore.migrate(db_name=scratch_name(db_url)) upgraded = backup_service.read_counts(scratch)["alembic"] finally: drop_scratch(db_url) return {**compare(recorded, restored, live), "alembic_dump": restored["alembic"], "alembic_now": upgraded} def compare(recorded: dict | None, restored: dict, live: dict) -> dict: """Each table's rows in the dump (as recorded), in the restored copy, and in the live database. With counts recorded, every table must come back with exactly its count; `mismatched` names those that didn't.""" rec = (recorded or {}).get("tables") names = sorted(set(restored["tables"]) | set(rec or {}) | set(live["tables"])) tables = [ {"table": n, "dump": rec.get(n) if rec is not None else None, "restored": restored["tables"].get(n), "live": live["tables"].get(n)} for n in names ] mismatched = [t["table"] for t in tables if rec is not None and t["table"] in rec and t["restored"] != t["dump"]] return { "counted": rec is not None, "tables": tables, "mismatched": mismatched, "rows_restored": sum(v for v in restored["tables"].values()), } def summary(result: dict) -> str: """One sentence on how a test went.""" n = len([t for t in result["tables"] if t["restored"] is not None]) rows = f"{result['rows_restored']:,}" if result["mismatched"]: shown = ", ".join( f"{t['table']} ({t['dump']:,} in the dump, " + ("table missing)" if t["restored"] is None else f"{t['restored']:,} restored)") for t in result["tables"] if t["table"] in result["mismatched"][:5] ) more = len(result["mismatched"]) - 5 return (f"{len(result['mismatched'])} table(s) came back different: {shown}" + (f", and {more} more" if more > 0 else "") + ".") if result["counted"]: return f"All {n} tables came back with the {rows} rows the dump recorded." return (f"The dump loaded cleanly: {n} tables, {rows} rows. It was taken before " "dumps recorded their counts, so it is shown beside the live database instead.")