brain-of-reese/scripts/restore_test_kb.py

"""One-shot restore of the controlled tool-calling test KB (the fast loop).

The fixture KB (``tests/fixtures/agent_kb/`` — two sources, eight
hand-written markdown documents) is built once by
:mod:`scripts.load_test_kb`, which embeds the documents and snapshots the
resulting database state into ``tests/fixtures/test_kb.dump.sql`` — a
data-only SQL script (``TRUNCATE`` + one multi-row ``INSERT`` per app
table, generated in-process — the same file runs in psql or psycopg). This
script restores that
snapshot in **one transaction** through the app's own database URL
(``BOR_DATABASE_URL``): no git clone of the homelab repo, no re-embedding,
no ``lite``-model calls — the whole known state (documents, chunks +
embeddings, the source registry rows, the KB overview, the sources
version) lands in a fraction of a second, which is what makes a
tool-calling iteration loop fast (see ``TOOL_CALLING_TESTING.md``):

    uv run python -m scripts.restore_test_kb
    # restore_test_kb: ok in 0.41s (8 docs, 2 sources, 16 chunks)

The gate runs the same restore inline:
``uv run python -m scripts.agent_realmodel_check --restore``.

The dump is data-only on purpose: the schema stays owned by alembic, and
the generated ``chunks.tsv`` tsvector column (``GENERATED ALWAYS AS …
STORED``) is recomputed by Postgres, so the restore is safe against schema
drift limited to additive columns. Restoring into a database whose schema
lacks an app table fails loudly with an actionable line (exit 2).

Exit codes: **0** restored, **2** precondition failure (DB unreachable,
dump missing, schema not applied).
"""
from __future__ import annotations

import argparse
import sys
import time
from dataclasses import dataclass
from pathlib import Path

import psycopg
from dotenv import load_dotenv
from sqlalchemy import text

from app.db import SessionLocal, db_available

#: The app tables the dump covers, TRUNCATE order (one statement —
#: Postgres resolves the inter-table FKs within it). ``chunks`` and
#: ``documents`` are listed first for readability; the order is
#: irrelevant inside a single TRUNCATE.
APP_TABLES: tuple[str, ...] = (
    "chunks",
    "documents",
    "git_sources",
    "kb_overview",
    "sources_meta",
    "steering_notes",
    "query_log",
    "saved_chats",
    "doc_drafts",
)

#: Repo-relative default dump location (the load script writes it there).
DEFAULT_DUMP = Path("tests/fixtures/test_kb.dump.sql")


@dataclass(frozen=True)
class RestoreResult:
    """What :func:`restore_dump` did — one line of the summary output."""

    seconds: float
    docs: int
    sources: tuple[str, ...]
    chunks: int
    dump_bytes: int


def _connect():
    """A raw psycopg connection on the app's DB URL (psycopg3 speaks the
    SQLAlchemy URL's driver scheme — ``postgresql+psycopg`` maps to
    ``postgresql`` for psycopg)."""
    from app.config import get_settings

    url = get_settings().database_url
    if url.startswith("postgresql+psycopg://"):
        url = "postgresql://" + url.split("://", 1)[1]
    return psycopg.connect(url)


def restore_dump(dump: Path) -> RestoreResult:
    """Restore *dump* (the data-only SQL script) into the app database.

    One transaction (TRUNCATE + INSERTs + nothing else — a failed restore
    rolls back and leaves the previous KB intact). Returns the measured
    result; raises :class:`RuntimeError` with an actionable line on
    failure (missing table = schema not applied).
    """
    if not dump.is_file():
        raise RuntimeError(
            f"dump not found: {dump} — build it first: "
            "`uv run python -m scripts.load_test_kb`"
        )
    script = dump.read_text(encoding="utf-8")
    started = time.monotonic()
    conn = _connect()
    try:
        with conn.transaction():
            # A plain multi-statement SQL script (TRUNCATE + INSERTs — no
            # parameters) runs on psycopg's simple-protocol execute; the
            # installed stubs type the query parameter as Template-only
            # (and ``sql.SQL`` wants a LiteralString), hence the ignore.
            conn.execute(script)  # pyright: ignore[reportArgumentType, reportCallIssue]
    except Exception as e:
        message = str(e)
        if "relation" in message and "does not exist" in message:
            raise RuntimeError(
                "schema not applied — the dump needs the alembic-managed "
                f"tables; run `uv run alembic upgrade head` first ({e})"
            ) from None
        raise RuntimeError(f"restore failed: {e}") from None
    seconds = time.monotonic() - started
    with SessionLocal() as db:
        docs = db.execute(text("select count(*) from documents")).scalar_one()
        sources = tuple(
            row[0]
            for row in db.execute(
                text("select distinct source from documents order by source")
            )
        )
        chunks = db.execute(text("select count(*) from chunks")).scalar_one()
    return RestoreResult(
        seconds=seconds,
        docs=docs,
        sources=sources,
        chunks=chunks,
        dump_bytes=dump.stat().st_size,
    )


def main(argv: list[str] | None = None) -> int:
    # CLI-only: pick up .env without side effects on import (the house
    # probe pattern, cf. scripts/llm_probe.py).
    load_dotenv()
    parser = argparse.ArgumentParser(
        description=(
            "Restore the controlled tool-calling test KB from the fixture "
            "dump (one transaction, no git clone, no re-embedding). Exit "
            "0 on success, 2 on precondition failure."
        )
    )
    parser.add_argument(
        "--dump",
        type=Path,
        default=DEFAULT_DUMP,
        help=f"the data-only SQL dump to restore (default: {DEFAULT_DUMP})",
    )
    args = parser.parse_args(argv)

    if not db_available():
        print(
            "restore_test_kb: precondition failed — database unreachable; "
            "start Postgres with `podman compose up -d db`"
        )
        return 2
    try:
        result = restore_dump(args.dump)
    except RuntimeError as e:
        print(f"restore_test_kb: {e}")
        return 2
    print(
        f"restore_test_kb: ok in {result.seconds:.2f}s "
        f"({result.docs} docs, {len(result.sources)} sources, "
        f"{result.chunks} chunks, dump {result.dump_bytes // 1024} KB)"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())