arborist/arborist/snapshot.py

"""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 `arborist snapshot verify <root>` 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 arborist.merkle import MerkleTree
from arborist.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",
    }