arborist/bench/pre_migration_snapshot.py

"""Pre-migration row-count snapshot.

Captures per-shard (docs, chunks, edges, audit_events) + global
document_root count for ~/.arborist/shards/00[0-3].db. Writes to
bench/results/pre-migration-snapshot.json so post-#000065 hydration
can verify row counts match.

Read-only; no shard mutation.
"""
from __future__ import annotations

import json
import os
import sqlite3
import sys
import time
from pathlib import Path

SHARDS_DIR = Path(os.path.expanduser("~/.arborist/shards"))
OUT_PATH = Path(__file__).resolve().parent / "results" / "pre-migration-snapshot.json"
SHARD_FILES = sorted(SHARDS_DIR.glob("00[0-9].db"))


def _count(conn: sqlite3.Connection, sql: str) -> int:
    row = conn.execute(sql).fetchone()
    return int(row[0]) if row else 0


def _snapshot_shard(path: Path) -> dict:
    conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    try:
        info = {
            "path": str(path),
            "size_bytes": path.stat().st_size,
            "docs": _count(conn, "SELECT COUNT(*) FROM documents"),
            "chunks": _count(conn, "SELECT COUNT(*) FROM chunks"),
            "edges": _count(conn, "SELECT COUNT(*) FROM edges"),
            "audit_events": _count(conn, "SELECT COUNT(*) FROM audit_events"),
            "providence_cache": _count(
                conn, "SELECT COUNT(*) FROM providence_cache"
            ) if _table_exists(conn, "providence_cache") else 0,
            "schema_version": _scalar(
                conn, "SELECT value FROM meta WHERE key='schema_version'"
            ),
        }
        cur = conn.execute(
            "SELECT source_type, COUNT(*) FROM documents GROUP BY source_type"
        )
        info["docs_by_source_type"] = {row[0] or "": int(row[1]) for row in cur}
        info["audit_first_ts"] = _scalar(
            conn, "SELECT MIN(ts) FROM audit_events"
        )
        info["audit_last_ts"] = _scalar(
            conn, "SELECT MAX(ts) FROM audit_events"
        )
        cur = conn.execute(
            "SELECT event_type, COUNT(*) FROM audit_events GROUP BY event_type"
        )
        info["audit_by_type"] = {row[0]: int(row[1]) for row in cur}
        return info
    finally:
        conn.close()


def _scalar(conn: sqlite3.Connection, sql: str):
    row = conn.execute(sql).fetchone()
    return row[0] if row else None


def _table_exists(conn: sqlite3.Connection, name: str) -> bool:
    row = conn.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)
    ).fetchone()
    return row is not None


def _global_unique_docs(shards: list[Path]) -> int:
    seen: set[bytes] = set()
    for shard in shards:
        conn = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
        try:
            cur = conn.execute("SELECT document_root FROM documents")
            for (root,) in cur:
                seen.add(root)
        finally:
            conn.close()
    return len(seen)


def main() -> int:
    if not SHARD_FILES:
        print(f"no shards under {SHARDS_DIR}", file=sys.stderr)
        return 1
    t0 = time.time()
    out = {
        "captured_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "shards_dir": str(SHARDS_DIR),
        "shards": [_snapshot_shard(p) for p in SHARD_FILES],
    }
    out["totals"] = {
        "size_bytes": sum(s["size_bytes"] for s in out["shards"]),
        "docs": sum(s["docs"] for s in out["shards"]),
        "chunks": sum(s["chunks"] for s in out["shards"]),
        "edges": sum(s["edges"] for s in out["shards"]),
        "audit_events": sum(s["audit_events"] for s in out["shards"]),
    }
    out["totals"]["docs_globally_unique"] = _global_unique_docs(SHARD_FILES)
    out["elapsed_seconds"] = round(time.time() - t0, 2)
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUT_PATH.write_text(json.dumps(out, indent=2, default=str))
    print(json.dumps(out["totals"], indent=2, default=str))
    print(f"\nwrote {OUT_PATH}")
    print(f"elapsed {out['elapsed_seconds']}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())