FabledCurator/backend/app/services/backup_verify.py

"""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 `<dump name>.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.")