diff --git a/aborist/cli.py b/aborist/cli.py index 68d8829..4e45e1d 100644 --- a/aborist/cli.py +++ b/aborist/cli.py @@ -13,7 +13,14 @@ from aborist.ingest import ingest_source, verify_random_sample from aborist.progress import Progress from aborist.search import FTS5Backend from aborist.sources import WikipediaCurDump -from aborist.store import DEFAULT_DB_PATH, connect, connect_query, stats +from aborist.store import ( + DEFAULT_DB_PATH, + append_audit, + connect, + connect_query, + stats, + transaction, +) def _cmd_ingest(args: argparse.Namespace) -> int: @@ -877,6 +884,143 @@ def _cmd_analyze(args: argparse.Namespace) -> int: return 0 +def _cmd_snapshot_create(args: argparse.Namespace) -> int: + """Compute snapshot_root over the read scope, persist into args.db. + + Single-DB mode (--db only): read + write are the same connection; + delegate to the snapshot module's create_snapshot(). + + Cross-shard mode (--shards-dir + --db): read against the in-memory + UNION view to get the cluster-level Merkle root, then persist into + args.db (a dedicated snapshots store, conventionally + `~/.aborist/shards/snapshots.db`). The writer's own documents table + is irrelevant to the snapshot value — only the union scope counts. + """ + import time as _time + + from aborist.snapshot import compute_snapshot_root, create_snapshot + + if args.global_shards_dir is None: + conn = connect(args.db) + try: + result = create_snapshot( + conn, reason=args.reason, parent_snapshot=args.parent, + ) + finally: + conn.close() + print(json.dumps(result, indent=2)) + return 0 + + # Cross-shard: compute against UNION, write to args.db. + read_conn = connect_query(args.db, shards_dir=args.global_shards_dir) + try: + snapshot_root, doc_count = compute_snapshot_root(read_conn) + finally: + read_conn.close() + + write_conn = connect(args.db) + try: + parent = args.parent + if parent is None: + row = write_conn.execute( + "SELECT snapshot_root FROM snapshots ORDER BY taken_at DESC LIMIT 1" + ).fetchone() + if row is not None: + parent = row["snapshot_root"] + + now = int(_time.time()) + body = { + "snapshot_root": snapshot_root, + "doc_count": doc_count, + "parent_snapshot": parent, + "reason": args.reason, + "scope": "shards-union", + } + audit_event_hash = append_audit( + write_conn, + event_type="snapshot_create", + body=body, + subject_root=snapshot_root, + ts=now, + ) + with transaction(write_conn): + write_conn.execute( + "INSERT OR IGNORE INTO snapshots " + "(snapshot_root, taken_at, audit_event_hash, doc_count, " + " parent_snapshot, reason) VALUES (?, ?, ?, ?, ?, ?)", + ( + snapshot_root, + now, + audit_event_hash, + doc_count, + parent, + args.reason, + ), + ) + finally: + write_conn.close() + + print( + json.dumps( + { + "snapshot_root": snapshot_root, + "doc_count": doc_count, + "parent_snapshot": parent, + "audit_event_hash": audit_event_hash, + "taken_at": now, + "reason": args.reason, + "scope": "shards-union", + }, + indent=2, + ) + ) + return 0 + + +def _cmd_snapshot_list(args: argparse.Namespace) -> int: + from aborist.snapshot import list_snapshots + + conn = connect(args.db) + try: + rows = list_snapshots(conn, limit=args.limit) + finally: + conn.close() + print(json.dumps(rows, indent=2)) + return 0 + + +def _cmd_snapshot_verify(args: argparse.Namespace) -> int: + from aborist.snapshot import verify_snapshot + + conn = ( + connect_query(args.db, shards_dir=args.global_shards_dir) + if args.global_shards_dir + else connect(args.db) + ) + try: + result = verify_snapshot(conn, args.snapshot_root) + finally: + conn.close() + print(json.dumps(result, indent=2)) + return 0 if result["matches"] else 1 + + +def _cmd_snapshot_diff(args: argparse.Namespace) -> int: + from aborist.snapshot import diff_against_current + + conn = ( + connect_query(args.db, shards_dir=args.global_shards_dir) + if args.global_shards_dir + else connect(args.db) + ) + try: + result = diff_against_current(conn, args.snapshot_root) + finally: + conn.close() + print(json.dumps(result, indent=2)) + return 0 + + def _cmd_mesh_status(args: argparse.Namespace) -> int: from aborist.mesh import current_epoch, is_enabled, load_identity from aborist.mesh.state import roster_at @@ -1427,6 +1571,42 @@ def build_parser() -> argparse.ArgumentParser: ) analyze_cmd.set_defaults(func=_cmd_analyze) + # ----- snapshot subcommands ---------------------------------------------- + snap_cmd = sub.add_parser( + "snapshot", + help="corpus-level Merkle snapshots: pin a forest state by single root", + ) + snap_sub = snap_cmd.add_subparsers(dest="snap_op", required=True) + + snap_create = snap_sub.add_parser( + "create", help="compute snapshot_root from current corpus, persist + audit" + ) + snap_create.add_argument("--reason", default="manual") + snap_create.add_argument( + "--parent", + default=None, + help="explicit parent_snapshot hex (default: auto-link to latest prior snapshot)", + ) + snap_create.set_defaults(func=_cmd_snapshot_create) + + snap_list = snap_sub.add_parser("list", help="recent snapshots, newest first") + snap_list.add_argument("--limit", type=int, default=20) + snap_list.set_defaults(func=_cmd_snapshot_list) + + snap_verify = snap_sub.add_parser( + "verify", + help="recompute root from current corpus; matches=True iff nothing has changed", + ) + snap_verify.add_argument("snapshot_root", help="hex snapshot_root to verify") + snap_verify.set_defaults(func=_cmd_snapshot_verify) + + snap_diff = snap_sub.add_parser( + "diff", + help="coarse drift signal between a snapshot and the current corpus", + ) + snap_diff.add_argument("snapshot_root", help="hex snapshot_root to diff against current") + snap_diff.set_defaults(func=_cmd_snapshot_diff) + # ----- mesh subcommands (off by default) --------------------------------- mesh_cmd = sub.add_parser( "mesh", diff --git a/aborist/snapshot.py b/aborist/snapshot.py new file mode 100644 index 0000000..bd515b3 --- /dev/null +++ b/aborist/snapshot.py @@ -0,0 +1,215 @@ +"""Corpus-level snapshots: Merkle-rooted pin of the document forest. + +A `snapshot_root` is `MerkleTree.build([sorted document_roots]).root`. +That single 32-byte hash is the canonical identity of the corpus state +at one point in time. Two peers that ingested the same dump compute +bit-identical snapshot_roots — so peers can answer "are we synced?" +with a single hash comparison, derive the TF-IDF scope_root for free +(snapshot_root *is* a scope), and pin Q&A answers to a verifiable +corpus state. + +Storage of the snapshots themselves is local: a `snapshots` row records +the root, audit-chain pin, doc_count, optional parent (for chains), and +human-readable reason. The root is reproducible from the documents +table at any time, so `aborist snapshot verify ` re-runs the +computation and compares. + +In sharded mode (`--shards-dir`) the snapshot reads the UNION view +across shards (cluster-level corpus), but writes go to whichever DB +the writer connection points at — typically a dedicated snapshots +shard alongside the ingest shards (see CLI defaults). +""" + +from __future__ import annotations + +import sqlite3 +import time +from typing import Iterable + +from aborist.merkle import MerkleTree +from aborist.store import append_audit, transaction + + +# Hash bytes of an empty corpus. Distinct from the genesis Merkle root over +# zero leaves (which is undefined); we name "no docs ingested yet" with +# a constant so the API never returns None for a real query. +EMPTY_SNAPSHOT_ROOT = "00" * 32 + + +def _all_document_roots(conn: sqlite3.Connection) -> list[str]: + """Sorted, *deduplicated* list of every document_root visible here. + + Cross-shard `--shards-dir` connections see the UNION view, which can + yield the same document_root from multiple shards if (rare but real) + two shards happen to compute the same content hash for content that + canonicalizes identically — e.g. near-empty redirect stubs whose + bodies are byte-identical after `canonicalize()`. The snapshot is a + membership root, not a multiset root, so we dedup here. Without this, + `verify` and `diff` could disagree on the same DB. + + Sorting is required for the Merkle root to be order-independent + across machines that ingested the same dump in different orders. + """ + rows = conn.execute( + "SELECT DISTINCT document_root FROM documents " + "ORDER BY document_root ASC" + ).fetchall() + return [r["document_root"] for r in rows] + + +def compute_snapshot_root( + conn: sqlite3.Connection | None = None, + *, + document_roots: Iterable[str] | None = None, +) -> tuple[str, int]: + """Return (snapshot_root_hex, doc_count) for the given document set. + + Pass `document_roots` to compute against an explicit set (for `verify` + and `diff`); otherwise reads from the connection's documents table. + """ + if document_roots is None: + if conn is None: + raise ValueError("either conn or document_roots must be provided") + roots = _all_document_roots(conn) + else: + roots = sorted(document_roots) + if not roots: + return EMPTY_SNAPSHOT_ROOT, 0 + if len(roots) == 1: + # Degenerate single-leaf tree: the leaf hash IS the root. + return roots[0], 1 + leaves = [bytes.fromhex(r) for r in roots] + return MerkleTree.build(leaves).root.hex(), len(roots) + + +def create_snapshot( + conn: sqlite3.Connection, + *, + reason: str = "manual", + parent_snapshot: str | None = None, +) -> dict: + """Compute the current snapshot_root, persist it, and pin it into the + audit chain. Idempotent on re-run when the corpus hasn't changed — + same root collides on the PK and we update only `last_seen_at` would + require a column we don't have, so the second insert is a no-op. + + `parent_snapshot` is optional. If omitted, we auto-link to the most + recent prior snapshot (gives you a snapshot chain for free). + """ + snapshot_root, doc_count = compute_snapshot_root(conn) + now = int(time.time()) + + if parent_snapshot is None: + prior = conn.execute( + "SELECT snapshot_root FROM snapshots ORDER BY taken_at DESC LIMIT 1" + ).fetchone() + if prior is not None: + parent_snapshot = prior["snapshot_root"] + + body = { + "snapshot_root": snapshot_root, + "doc_count": doc_count, + "parent_snapshot": parent_snapshot, + "reason": reason, + } + audit_event_hash = append_audit( + conn, + event_type="snapshot_create", + body=body, + subject_root=snapshot_root, + ts=now, + ) + + with transaction(conn): + conn.execute( + "INSERT OR IGNORE INTO snapshots " + "(snapshot_root, taken_at, audit_event_hash, doc_count, parent_snapshot, reason) " + "VALUES (?, ?, ?, ?, ?, ?)", + ( + snapshot_root, + now, + audit_event_hash, + doc_count, + parent_snapshot, + reason, + ), + ) + + return { + "snapshot_root": snapshot_root, + "doc_count": doc_count, + "parent_snapshot": parent_snapshot, + "audit_event_hash": audit_event_hash, + "taken_at": now, + "reason": reason, + } + + +def list_snapshots(conn: sqlite3.Connection, *, limit: int = 20) -> list[dict]: + """Recent snapshots, newest first.""" + rows = conn.execute( + "SELECT snapshot_root, taken_at, audit_event_hash, doc_count, " + " parent_snapshot, reason " + "FROM snapshots ORDER BY taken_at DESC LIMIT ?", + (limit,), + ).fetchall() + return [dict(r) for r in rows] + + +def verify_snapshot(conn: sqlite3.Connection, snapshot_root: str) -> dict: + """Re-derive the snapshot_root from the *current* corpus and compare. + + Returns a result with `matches: True` only if the current document + set produces the same root — i.e. nothing has been added, removed, + or superseded since the snapshot was taken. A False result is the + drift signal; pair with `diff_snapshots` to see what changed. + """ + current_root, current_count = compute_snapshot_root(conn) + return { + "snapshot_root": snapshot_root, + "current_root": current_root, + "current_doc_count": current_count, + "matches": snapshot_root == current_root, + } + + +def diff_against_current( + conn: sqlite3.Connection, + snapshot_root: str, +) -> dict: + """Set-diff the recorded snapshot's roots against the current corpus. + + `added` = roots in current corpus but not in the snapshot + `removed` = roots in the snapshot but not in current + `unchanged_count` = intersection size + + Note: the recorded snapshot stores only its top-level Merkle root, not + the full leaf set. To diff, we need to re-derive the leaves. This + function expects the snapshot to be re-computable from the current + DB if it still matches; otherwise it can only say "different" without + enumerating. For full forensic diffing across drifted snapshots, + persist the leaf set separately (future work). + """ + current_roots = set(_all_document_roots(conn)) + current_root, _ = compute_snapshot_root( + conn, document_roots=current_roots + ) + if snapshot_root == current_root: + return { + "snapshot_root": snapshot_root, + "current_root": current_root, + "added_count": 0, + "removed_count": 0, + "unchanged_count": len(current_roots), + "diff": "identical", + } + # We can't enumerate WHAT changed without the snapshot's leaf set on + # hand. Surface the count of each side as a coarse signal. + return { + "snapshot_root": snapshot_root, + "current_root": current_root, + "current_doc_count": len(current_roots), + "diff": "drifted", + "note": "exact membership diff requires the snapshot's leaf set; " + "store leaf sets alongside snapshot rows for full forensic diffs", + } diff --git a/aborist/store.py b/aborist/store.py index 517f97a..6e05067 100644 --- a/aborist/store.py +++ b/aborist/store.py @@ -179,6 +179,21 @@ CREATE TABLE IF NOT EXISTS falsifications ( PRIMARY KEY (cache_key, at) ); +-- Snapshots: corpus-level Merkle root pinning a forest state at a point in +-- time. snapshot_root = MerkleTree.build([sorted document_roots]). Audit- +-- chain-linked so peers can verify a claimed snapshot was actually witnessed +-- by this instance. parent_snapshot lets snapshots chain (A -> B -> C) for +-- diff/replay. doc_count is informational; the root is the canonical id. +CREATE TABLE IF NOT EXISTS snapshots ( + snapshot_root TEXT PRIMARY KEY, + taken_at INTEGER NOT NULL, + audit_event_hash TEXT NOT NULL, + doc_count INTEGER NOT NULL, + parent_snapshot TEXT, + reason TEXT +); +CREATE INDEX IF NOT EXISTS idx_snapshots_taken_at ON snapshots(taken_at); + -- Mesh layer tables. Off by default — populated only when the user runs -- `aborist mesh init`. Never accessed by ingest / query / distill paths; -- mesh state is opt-in plumbing for federated peers (see aborist.mesh). diff --git a/tests/test_snapshot.py b/tests/test_snapshot.py new file mode 100644 index 0000000..bb742fa --- /dev/null +++ b/tests/test_snapshot.py @@ -0,0 +1,217 @@ +"""Tests for corpus-level snapshots. + +Snapshots pin a forest of document_roots into a single Merkle root +that's bit-identical across machines that ingested the same dump. +""" + +from __future__ import annotations + +import time + +import pytest + +from aborist.document import Document +from aborist.ingest import ingest_source +from aborist.snapshot import ( + EMPTY_SNAPSHOT_ROOT, + compute_snapshot_root, + create_snapshot, + diff_against_current, + list_snapshots, + verify_snapshot, +) +from aborist.source import Source +from aborist.store import connect + + +class _FakeSource(Source): + source_type = "test" + + def __init__(self, docs): + self._docs = docs + + def iter_documents(self): + for d in self._docs: + yield d + + +def _doc(uri: str, content: str) -> Document: + return Document(uri=uri, content=content, source_type="test", title=uri) + + +# --------------------------------------------------------------------------- +# Pure compute +# --------------------------------------------------------------------------- + + +def test_compute_snapshot_root_empty_corpus(tmp_path): + conn = connect(tmp_path / "a.db") + try: + root, count = compute_snapshot_root(conn) + assert root == EMPTY_SNAPSHOT_ROOT + assert count == 0 + finally: + conn.close() + + +def test_compute_snapshot_root_single_doc_uses_doc_root(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + root, count = compute_snapshot_root(conn) + # Degenerate single-leaf: snapshot_root == document_root. + doc_row = conn.execute( + "SELECT document_root FROM documents" + ).fetchone() + assert root == doc_row["document_root"] + assert count == 1 + finally: + conn.close() + + +def test_compute_snapshot_root_is_order_independent(tmp_path): + """Two corpora with the same documents in different ingest order + must produce the same snapshot_root. That's the cross-machine + convergence property.""" + a = tmp_path / "a.db" + b = tmp_path / "b.db" + + docs1 = [_doc("test://x", "x" * 100), _doc("test://y", "y" * 100)] + docs2 = list(reversed(docs1)) + + conn_a = connect(a) + conn_b = connect(b) + try: + ingest_source(conn_a, _FakeSource(docs1)) + ingest_source(conn_b, _FakeSource(docs2)) + root_a, _ = compute_snapshot_root(conn_a) + root_b, _ = compute_snapshot_root(conn_b) + assert root_a == root_b + finally: + conn_a.close() + conn_b.close() + + +def test_compute_snapshot_root_changes_when_doc_added(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "alpha alpha alpha alpha")])) + root1, _ = compute_snapshot_root(conn) + ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")])) + root2, _ = compute_snapshot_root(conn) + assert root1 != root2 + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Persistence + audit chain +# --------------------------------------------------------------------------- + + +def test_create_snapshot_persists_and_audits(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + result = create_snapshot(conn, reason="post-ingest") + assert result["doc_count"] == 1 + assert len(result["snapshot_root"]) == 64 + assert len(result["audit_event_hash"]) == 64 + + rows = list_snapshots(conn) + assert len(rows) == 1 + assert rows[0]["snapshot_root"] == result["snapshot_root"] + assert rows[0]["reason"] == "post-ingest" + + # Audit chain has a snapshot_create event linked to the snapshot_root. + audit = conn.execute( + "SELECT event_type, subject_root FROM audit_events " + "WHERE event_type = 'snapshot_create' ORDER BY seq DESC LIMIT 1" + ).fetchone() + assert audit["event_type"] == "snapshot_create" + assert audit["subject_root"] == result["snapshot_root"] + finally: + conn.close() + + +def test_create_snapshot_auto_links_to_prior(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + first = create_snapshot(conn, reason="first") + # Add a doc; second snapshot's parent should auto-link to first. + ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")])) + # Sleep briefly so taken_at differs even on fast machines. + time.sleep(1.1) + second = create_snapshot(conn, reason="second") + assert second["snapshot_root"] != first["snapshot_root"] + assert second["parent_snapshot"] == first["snapshot_root"] + finally: + conn.close() + + +def test_create_snapshot_idempotent_on_unchanged_corpus(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + a = create_snapshot(conn, reason="first") + b = create_snapshot(conn, reason="second-no-change") + # Same corpus -> same snapshot_root. Second insert is a no-op + # against the PK; we don't get a duplicate row. + assert a["snapshot_root"] == b["snapshot_root"] + assert len(list_snapshots(conn)) == 1 + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Verify + diff +# --------------------------------------------------------------------------- + + +def test_verify_snapshot_matches_when_corpus_unchanged(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + s = create_snapshot(conn) + result = verify_snapshot(conn, s["snapshot_root"]) + assert result["matches"] + finally: + conn.close() + + +def test_verify_snapshot_rejects_after_drift(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + s = create_snapshot(conn) + ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")])) + result = verify_snapshot(conn, s["snapshot_root"]) + assert not result["matches"] + assert result["current_doc_count"] == 2 + finally: + conn.close() + + +def test_diff_returns_identical_for_matching_root(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + s = create_snapshot(conn) + result = diff_against_current(conn, s["snapshot_root"]) + assert result["diff"] == "identical" + finally: + conn.close() + + +def test_diff_reports_drift_after_adds(tmp_path): + conn = connect(tmp_path / "a.db") + try: + ingest_source(conn, _FakeSource([_doc("test://a", "hello world a")])) + s = create_snapshot(conn) + ingest_source(conn, _FakeSource([_doc("test://b", "beta beta beta beta")])) + result = diff_against_current(conn, s["snapshot_root"]) + assert result["diff"] == "drifted" + assert result["snapshot_root"] != result["current_root"] + finally: + conn.close()