arborist/bench/extract_audit_events.py

"""Extract every audit_events row from every shard to NDJSON.

For #000065 Option-A audit-chain consolidation: the reshard moves
documents to new shards by content hash, so the per-shard audit
chains are broken at boundaries. Option A rebuilds a single canonical
chain on shard 000, sorted by ts. Inputs to that rebuild are the
union of all source shards' audit_events rows — captured here BEFORE
any source shard is touched.

Output: /tmp/audit-events.ndjson
  fields per line:
    src_shard      source shard index (0..3)
    src_seq        source rowid (forensic reference)
    src_event_hash original chain hash (forensic reference)
    prev_event_hash  original chain predecessor (forensic)
    event_type
    subject_root
    body           canonical JSON, preserved unchanged
    ts             unix seconds

Verifies chain integrity on each source shard during extraction so
we know what's captured is intact going in. Read-only; never mutates
source shards.
"""
from __future__ import annotations

import hashlib
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("/tmp/audit-events.ndjson")
SHARD_FILES = sorted(SHARDS_DIR.glob("00[0-9].db"))


def _verify_chain(conn: sqlite3.Connection, shard_idx: int) -> tuple[int, int]:
    """Walk audit_events in seq order, recompute event_hash, count breaks."""
    cur = conn.execute(
        "SELECT seq, event_hash, prev_event_hash, body "
        "FROM audit_events ORDER BY seq ASC"
    )
    prev = None
    total = 0
    breaks = 0
    for seq, event_hash, prev_event_hash, body in cur:
        total += 1
        if prev_event_hash != prev:
            breaks += 1
        h = hashlib.sha256()
        if prev_event_hash is not None:
            h.update(bytes.fromhex(prev_event_hash))
        h.update(body.encode("utf-8", errors="surrogatepass"))
        if h.hexdigest() != event_hash:
            breaks += 1
        prev = event_hash
    return total, breaks


def _extract_shard(shard_idx: int, path: Path, out) -> int:
    conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    try:
        total, breaks = _verify_chain(conn, shard_idx)
        if breaks:
            print(
                f"  WARNING: shard {shard_idx:03d} has {breaks} chain breaks",
                file=sys.stderr,
            )
        cur = conn.execute(
            "SELECT seq, event_hash, prev_event_hash, event_type, "
            "subject_root, body, ts FROM audit_events ORDER BY seq ASC"
        )
        count = 0
        for seq, event_hash, prev_eh, etype, subj, body, ts in cur:
            row = {
                "src_shard": shard_idx,
                "src_seq": seq,
                "src_event_hash": event_hash,
                "prev_event_hash": prev_eh,
                "event_type": etype,
                "subject_root": subj,
                "body": body,
                "ts": ts,
            }
            out.write(json.dumps(row, separators=(",", ":")))
            out.write("\n")
            count += 1
        return count
    finally:
        conn.close()


def main() -> int:
    if not SHARD_FILES:
        print(f"no shards under {SHARDS_DIR}", file=sys.stderr)
        return 1
    t0 = time.time()
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    total = 0
    per_shard: list[int] = []
    with OUT_PATH.open("w") as out:
        for shard in SHARD_FILES:
            idx = int(shard.stem)
            n = _extract_shard(idx, shard, out)
            per_shard.append(n)
            total += n
            print(f"  shard {idx:03d}: {n:,} events")
    elapsed = time.time() - t0
    size_bytes = OUT_PATH.stat().st_size
    print()
    print(f"wrote {OUT_PATH}")
    print(f"  total events:  {total:,}")
    print(f"  per shard:     {per_shard}")
    print(f"  file size:     {size_bytes / 1e6:.1f} MB")
    print(f"  elapsed:       {elapsed:.1f}s")
    return 0


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