From 77b90bfcf18b5c290a984abc04b56c1726458849 Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Tue, 26 May 2026 13:28:52 -0400 Subject: [PATCH] #000065 step 3: strategy A executor + CLI + plan-only preview MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Build the 'all_at_once' executor and wire it into the CLI as `arborist corpus reshard`. Plan-only preview against fox's host confirms the planner picks all_at_once at 95.5 GB free, 64.2 GB peak draw, 31.3 GB free at peak (well above the 4 GB safety). Executor (arborist/migrate.py): * ROUTED_BY_DOCUMENT_ROOT — documents, document_http_meta, chunks, merkle_nodes, edges (by src_root), derivations (by core_root). Each row hashed to shard_for_document(root, M) and INSERT'd into the chosen target. * CONSOLIDATED_TABLES — snapshots, concept_relations, concept_token_idf, citation_aliases, term_aliases, providence_cache, falsifications all land on canonical target shard 000. * REBUILT_ON_TARGET — chunks_fts + documents_fts rebuilt from the materialized data after content moves. chunks.content is zstd-packed at rest so the rebuilder decompresses via arborist.compress.unpack_chunk before inserting plaintext into the FTS5 index. * Audit chain consolidation (Option A) — all rows from /tmp/audit-events.ndjson are sorted by (ts, src_shard, src_seq), re-chained with fresh event_hash values, and INSERT'd into target shard 000. Bodies preserved unchanged for forensic fidelity. * One 'reshard' audit event appended at the tail, carrying the plan + result as the body. An operator months later can answer "where did this corpus topology come from" from this one row. * corpus_shard_count meta stamped on every target shard. * VACUUM each target at end to reclaim INSERT-pattern fragmentation. CLI: `arborist corpus reshard --to M --source-dir DIR --target-dir DIR [--plan-only | --force-strategy X | --allow-in-place | --dry-run]` 11 integration tests build a tiny 2-shard corpus, run the migration, verify per-doc routing, chunk-follows-doc, audit chain integrity, reshard event at tail, corpus_shard_count meta on every target, chunks_fts searchability, doc count preservation, dry-run no-op, error handling. Also: get_meta() now indexes by position so callers without row_factory=sqlite3.Row don't trip. No behavior change for callers that DO use Row factory. Strategy B (per_source_shard) and C (streaming_row) are stubbed in the planner (peak-draw estimators wired) but execute_plan raises NotImplementedError for them. Not needed at fox's current disk (88+ GB free); skipping the build until that's ever the constrained path. --- arborist/cli.py | 137 ++++++++++ arborist/migrate.py | 455 ++++++++++++++++++++++++++++++++- arborist/store.py | 10 +- tests/test_migrate_executor.py | 343 +++++++++++++++++++++++++ 4 files changed, 941 insertions(+), 4 deletions(-) create mode 100644 tests/test_migrate_executor.py diff --git a/arborist/cli.py b/arborist/cli.py index 290ca99..7b5b6b2 100644 --- a/arborist/cli.py +++ b/arborist/cli.py @@ -2803,6 +2803,87 @@ def _make_cold_backend(): ) +def _cmd_corpus_reshard(args: argparse.Namespace) -> int: + """Content-hash reshard (#000065). + + Planner picks A/B/C from disk; executor runs the picked strategy. + The plan + rationale always print first so the operator can review + before any data moves. --plan-only stops after that. + """ + from arborist.migrate import ( + HydrationPlanner, + InsufficientDisk, + execute_plan, + ) + + source_dir = Path(args.source_dir).expanduser() + target_dir = Path(args.target_dir).expanduser() + audit_ndjson = Path(args.audit_ndjson).expanduser() + + try: + plan = HydrationPlanner().plan( + source_dir=source_dir, + target_dir=target_dir, + target_M=args.target_M, + force_strategy=args.force_strategy, + allow_in_place=args.allow_in_place, + ) + except InsufficientDisk as e: + print(f"insufficient disk: {e}", file=sys.stderr) + return 1 + except (ValueError, FileNotFoundError) as e: + print(f"plan error: {e}", file=sys.stderr) + return 2 + + # Always show the plan readout. Goes to stdout (not stderr) so + # operators can pipe to a log. + print(json.dumps(plan.as_audit_body(), indent=2, ensure_ascii=False)) + print( + f"\nstrategy: {plan.strategy}\n" + f"peak draw est: {plan.estimated_peak_draw_bytes / 1e9:.1f} GB\n" + f"free at peak: {plan.free_at_peak_bytes() / 1e9:.1f} GB\n" + f"rationale: {plan.rationale}", + file=sys.stderr, + ) + + if args.plan_only: + return 0 + + if not audit_ndjson.exists(): + print( + f"audit ndjson missing at {audit_ndjson}. " + f"Run `python bench/extract_audit_events.py` first.", + file=sys.stderr, + ) + return 3 + + try: + result = execute_plan( + plan, audit_events_ndjson=audit_ndjson, dry_run=args.dry_run + ) + except NotImplementedError as e: + print(f"strategy not implemented: {e}", file=sys.stderr) + return 4 + + print( + json.dumps( + { + "strategy": result.strategy, + "elapsed_seconds": round(result.elapsed_seconds, 1), + "rows_moved_total": sum(result.rows_moved.values()), + "rows_moved": result.rows_moved, + "audit_events_consolidated": result.audit_events_consolidated, + "final_event_hash": result.final_event_hash, + "target_row_counts": result.target_row_counts, + "notes": result.notes, + }, + indent=2, + ensure_ascii=False, + ) + ) + return 0 + + def _cmd_cold_pack(args: argparse.Namespace) -> int: from arborist.evict import push_pack @@ -5795,6 +5876,62 @@ def build_parser() -> argparse.ArgumentParser: ) cold_list.set_defaults(func=_cmd_cold_list, fetch_manifest=True) + # --- #000065 content-hash reshard ---------------------------------------- + corpus_cmd = sub.add_parser( + "corpus", + help="corpus-wide topology operations (reshard, …)", + ) + corpus_sub = corpus_cmd.add_subparsers(dest="corpus_op", required=True) + reshard = corpus_sub.add_parser( + "reshard", + help=( + "migrate shards to a content-hash-routed M-shard layout (#000065). " + "Picks a strategy from available disk; use --plan-only first." + ), + ) + reshard.add_argument( + "--to", dest="target_M", type=int, required=True, + help="target shard count M (e.g., 4)", + ) + reshard.add_argument( + "--source-dir", dest="source_dir", required=True, + help="directory of source shards (00[0-9].db files)", + ) + reshard.add_argument( + "--target-dir", dest="target_dir", required=True, + help="directory to write target shards (will be created)", + ) + reshard.add_argument( + "--audit-events-ndjson", dest="audit_ndjson", + default="/tmp/audit-events.ndjson", + help="ndjson of pre-migration audit events from " + "bench/extract_audit_events.py (default /tmp/audit-events.ndjson)", + ) + reshard.add_argument( + "--plan-only", dest="plan_only", action="store_true", + help="print the chosen strategy + peak draw estimate + rationale, " + "then exit. No DBs opened, no rows moved.", + ) + reshard.add_argument( + "--force-strategy", dest="force_strategy", + choices=["all_at_once", "per_source_shard", "streaming_row"], + default=None, + help="override the planner's pick. Use only when you know the " + "planner's draw estimate is wrong (e.g., for testing strategy B).", + ) + reshard.add_argument( + "--allow-in-place", dest="allow_in_place", action="store_true", + help="allow strategy C (streaming-row, in-place mutation of source " + "shards). Last-resort; disabled by default since it mutates " + "source data with no recoverable intermediate.", + ) + reshard.add_argument( + "--dry-run", dest="dry_run", action="store_true", + help="do everything except actually open the target DBs and move " + "rows. Plan still emits, target dir still created.", + ) + reshard.set_defaults(func=_cmd_corpus_reshard) + activity_cmd = sub.add_parser( "activity", help="recent Q&A + freshly cached docs (agent-readable timeline)", diff --git a/arborist/migrate.py b/arborist/migrate.py index 7ae79bc..48c9db1 100644 --- a/arborist/migrate.py +++ b/arborist/migrate.py @@ -23,10 +23,16 @@ one ``audit_events`` row at the end of every successful migration. """ from __future__ import annotations +import hashlib +import json import os -from dataclasses import dataclass +import sqlite3 +import time +from dataclasses import dataclass, field from pathlib import Path -from typing import Literal +from typing import Iterable, Literal + +from arborist.document import shard_for_document # --------------------------------------------------------------------------- @@ -326,3 +332,448 @@ class HydrationPlanner: estimated_final_target_bytes=facts.total_bytes, rationale=rationale, ) + + +# --------------------------------------------------------------------------- +# Table routing — which table goes where during a reshard. +# +# ROUTED_BY_DOCUMENT_ROOT — every row carries a column named here; the row +# follows its parent doc to the target shard hash_for(document_root) % M. +# +# CONSOLIDATED_TABLES — every row lands on canonical target shard 000. +# Corpus-wide or audit-anchored data; no doc routing applies. +# +# REPLICATED_TO_ALL — every target shard gets the same key/value rows. +# Per-shard config (schema_version, corpus_shard_count, etc). +# +# REBUILT_ON_TARGET — virtual table / index; do NOT copy rows, rebuild +# from the migrated content tables instead. +# +# DROPPED — transient state, no migration value. +# --------------------------------------------------------------------------- + +ROUTED_BY_DOCUMENT_ROOT: tuple[tuple[str, str], ...] = ( + # (table_name, column_used_for_routing) + ("documents", "document_root"), + ("document_http_meta", "document_root"), + ("chunks", "document_root"), + ("merkle_nodes", "document_root"), + ("edges", "src_root"), # edge follows its source doc + ("derivations", "core_root"), # derivation row lives with the core +) + +CONSOLIDATED_TABLES: tuple[str, ...] = ( + "snapshots", + "concept_relations", + "concept_token_idf", + "citation_aliases", + "term_aliases", + "providence_cache", + "falsifications", + # audit_events handled separately — re-chained, not copied. +) + +REPLICATED_TO_ALL: tuple[str, ...] = ( + "schema_meta", + "meta", +) + +REBUILT_ON_TARGET: tuple[str, ...] = ( + "chunks_fts", + "documents_fts", +) + + +# --------------------------------------------------------------------------- +# Executor — strategy A "all_at_once" +# --------------------------------------------------------------------------- + +@dataclass +class ExecutionResult: + """Returned by ``execute_plan``. Lands in the migration audit event.""" + + strategy: Strategy + started_at: float + finished_at: float + rows_moved: dict[str, int] = field(default_factory=dict) + target_row_counts: dict[str, dict[str, int]] = field(default_factory=dict) + audit_events_consolidated: int = 0 + final_event_hash: str | None = None + peak_free_at_run_bytes: int | None = None + notes: list[str] = field(default_factory=list) + + @property + def elapsed_seconds(self) -> float: + return self.finished_at - self.started_at + + +def _connect_target(path: Path) -> sqlite3.Connection: + """Open or create a target shard with the full v9.8 schema.""" + from arborist.store import SCHEMA_SQL + + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(str(path)) + conn.row_factory = sqlite3.Row + conn.executescript("PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL;") + conn.executescript(SCHEMA_SQL) + return conn + + +def _connect_source_readonly(path: Path) -> sqlite3.Connection: + conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + conn.row_factory = sqlite3.Row + return conn + + +def _table_columns(conn: sqlite3.Connection, table: str) -> list[str]: + rows = conn.execute(f"PRAGMA table_info({table})").fetchall() + return [r["name"] for r in rows] + + +def _route_per_doc_table( + *, + source: sqlite3.Connection, + targets: list[sqlite3.Connection], + table: str, + route_col: str, + M: int, + batch_size: int = 10_000, + skip_columns: frozenset[str] = frozenset(), +) -> int: + """Stream ``table`` rows from one source, route by ``route_col`` hash + to one of ``M`` target shards, executemany INSERT into the chosen + target. Returns row count moved. + + ``skip_columns`` is for AUTOINCREMENT PKs that the target should + assign fresh (chunks.chunk_id) — including them would collide across + sources. + """ + cols = _table_columns(source, table) + insert_cols = [c for c in cols if c not in skip_columns] + placeholders = ",".join(["?"] * len(insert_cols)) + col_list = ",".join(insert_cols) + insert_sql = ( + f"INSERT OR IGNORE INTO {table} ({col_list}) VALUES ({placeholders})" + ) + route_idx = cols.index(route_col) + keep_idxs = [i for i, c in enumerate(cols) if c not in skip_columns] + + total = 0 + cur = source.execute(f"SELECT {','.join(cols)} FROM {table}") + buckets: list[list[tuple]] = [[] for _ in range(M)] + while True: + batch = cur.fetchmany(batch_size) + if not batch: + break + for row in batch: + root = row[route_idx] + if not root: + continue + idx = shard_for_document(root, M) + buckets[idx].append(tuple(row[i] for i in keep_idxs)) + for idx, rows in enumerate(buckets): + if not rows: + continue + t = targets[idx] + with t: + t.executemany(insert_sql, rows) + total += len(rows) + rows.clear() + return total + + +def _copy_table_consolidated( + *, + source: sqlite3.Connection, + target: sqlite3.Connection, + table: str, + batch_size: int = 10_000, +) -> int: + """Stream all rows of ``table`` from ``source`` into ``target``.""" + cols = _table_columns(source, table) + if not cols: + return 0 + col_list = ",".join(cols) + placeholders = ",".join(["?"] * len(cols)) + insert_sql = ( + f"INSERT OR IGNORE INTO {table} ({col_list}) VALUES ({placeholders})" + ) + total = 0 + cur = source.execute(f"SELECT {col_list} FROM {table}") + while True: + batch = cur.fetchmany(batch_size) + if not batch: + break + with target: + target.executemany(insert_sql, [tuple(r) for r in batch]) + total += len(batch) + return total + + +def _rebuild_fts_on_target(conn: sqlite3.Connection) -> None: + """Populate the contentless FTS5 indexes from the materialized tables. + + chunks.content is zstd-compressed at rest (see arborist.compress); the + FTS5 index needs plaintext. Stream chunks, decompress, batch into FTS. + documents.title is plain text → straight copy. + """ + from arborist.compress import unpack_chunk + + batch_size = 5_000 + fts_rows: list[tuple[int, str]] = [] + cur = conn.execute( + "SELECT chunk_id, content FROM chunks WHERE content IS NOT NULL" + ) + while True: + rows = cur.fetchmany(batch_size) + if not rows: + break + for chunk_id, packed in rows: + text = unpack_chunk(packed) + if text is None: + continue + fts_rows.append((chunk_id, text)) + if len(fts_rows) >= batch_size: + with conn: + conn.executemany( + "INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)", + fts_rows, + ) + fts_rows.clear() + if fts_rows: + with conn: + conn.executemany( + "INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)", + fts_rows, + ) + + with conn: + conn.execute( + "INSERT INTO documents_fts (rowid, title) " + "SELECT rowid, title FROM documents WHERE title IS NOT NULL" + ) + + +def _consolidate_audit_chain( + *, + target: sqlite3.Connection, + ndjson_path: Path, +) -> tuple[int, str | None]: + """Read every audit_events row from ``ndjson_path``, sort by ts, + re-chain, and INSERT into the canonical target's audit_events table. + + Body strings are preserved unchanged (forensic fidelity). + ``event_hash`` values are recomputed under the new (sorted) chain. + """ + rows: list[dict] = [] + with ndjson_path.open() as f: + for line in f: + rows.append(json.loads(line)) + # Sort by (ts, src_shard, src_seq) for stable tiebreak across shards + # with identical wall-clock ts (the bulk-ingest pass timestamped many + # events to the same second). + rows.sort(key=lambda r: (r["ts"], r["src_shard"], r["src_seq"])) + + prev: str | None = None + out_rows: list[tuple] = [] + for r in rows: + body = r["body"] + h = hashlib.sha256() + if prev is not None: + h.update(bytes.fromhex(prev)) + h.update(body.encode("utf-8", errors="surrogatepass")) + event_hash = h.hexdigest() + out_rows.append( + ( + event_hash, + prev, + r["event_type"], + r["subject_root"], + body, + r["ts"], + ) + ) + prev = event_hash + + if not out_rows: + return 0, None + with target: + target.executemany( + "INSERT INTO audit_events " + "(event_hash, prev_event_hash, event_type, subject_root, body, ts) " + "VALUES (?, ?, ?, ?, ?, ?)", + out_rows, + ) + return len(out_rows), prev + + +def _validate_row_counts( + *, + targets: list[sqlite3.Connection], + expected: dict[str, int], +) -> dict[str, int]: + """Sum row counts across all target shards per table; compare to + expected. Returns dict of {table: actual_sum}. Caller decides if + a mismatch is fatal — at minimum, log it. + """ + actual: dict[str, int] = {} + for table in ("documents", "chunks", "edges"): + total = 0 + for t in targets: + r = t.execute(f"SELECT COUNT(*) FROM {table}").fetchone() + total += int(r[0]) + actual[table] = total + return actual + + +def execute_plan( + plan: HydrationPlan, + *, + audit_events_ndjson: Path, + expected_row_counts: dict[str, int] | None = None, + dry_run: bool = False, +) -> ExecutionResult: + """Execute a :class:`HydrationPlan`. Currently supports strategy A only. + + ``audit_events_ndjson`` is the file produced by + ``bench/extract_audit_events.py`` — captures every source audit event + BEFORE any source shard is touched, so Option-A chain consolidation + has its inputs even if the source shards have been deleted by the + time we re-chain. + """ + if plan.strategy != "all_at_once": + raise NotImplementedError( + f"executor for strategy {plan.strategy!r} not built yet; only " + f"'all_at_once' is implemented. Use --force-strategy=all_at_once " + f"if disk allows, or wait for the per-source-shard executor." + ) + if not audit_events_ndjson.exists(): + raise FileNotFoundError( + f"audit-event ndjson not found at {audit_events_ndjson}. " + f"Run bench/extract_audit_events.py first." + ) + return _execute_all_at_once( + plan, + audit_events_ndjson=audit_events_ndjson, + expected_row_counts=expected_row_counts or {}, + dry_run=dry_run, + ) + + +def _execute_all_at_once( + plan: HydrationPlan, + *, + audit_events_ndjson: Path, + expected_row_counts: dict[str, int], + dry_run: bool, +) -> ExecutionResult: + from arborist.store import set_corpus_shard_count + + started = time.time() + plan.target_dir.mkdir(parents=True, exist_ok=True) + result = ExecutionResult( + strategy="all_at_once", started_at=started, finished_at=started + ) + + # 1) Open M target connections, each pre-schema-applied. + target_paths = [ + plan.target_dir / f"{i:03d}.db" for i in range(plan.target_M) + ] + if dry_run: + result.notes.append("DRY RUN — no DBs opened, no rows moved") + result.finished_at = time.time() + return result + + targets = [_connect_target(p) for p in target_paths] + try: + # 2) Per-source-shard: stream every table, route or consolidate. + for src_path in plan.corpus_facts.shard_files: + src = _connect_source_readonly(src_path) + try: + # Routed tables — by document_root or src_root or core_root. + for table, route_col in ROUTED_BY_DOCUMENT_ROOT: + if not _table_exists(src, table): + continue + skip = frozenset() + # chunks.chunk_id is AUTOINCREMENT and would collide + # across sources — let target assign new IDs. + if table == "chunks": + skip = frozenset({"chunk_id"}) + n = _route_per_doc_table( + source=src, + targets=targets, + table=table, + route_col=route_col, + M=plan.target_M, + skip_columns=skip, + ) + result.rows_moved[f"{table}@{src_path.name}"] = n + + # Consolidated tables — every row to target 0. + for table in CONSOLIDATED_TABLES: + if not _table_exists(src, table): + continue + n = _copy_table_consolidated( + source=src, target=targets[0], table=table + ) + result.rows_moved[f"{table}@{src_path.name}"] = n + finally: + src.close() + + # 3) Rebuild FTS5 on every target from materialized rows. + for t in targets: + _rebuild_fts_on_target(t) + + # 4) Consolidate audit chain into target 0. + n_events, final_hash = _consolidate_audit_chain( + target=targets[0], ndjson_path=audit_events_ndjson + ) + result.audit_events_consolidated = n_events + result.final_event_hash = final_hash + + # 5) Append the migration's own reshard event to the canonical + # chain. This is the forensic record an operator finds when they + # later ask "where did this corpus topology come from?" + result.finished_at = time.time() + body = { + "kind": "reshard", + "plan": plan.as_audit_body(), + "result": { + "rows_moved_total": sum(result.rows_moved.values()), + "audit_events_consolidated": n_events, + "elapsed_seconds": result.elapsed_seconds, + }, + } + from arborist.store import append_audit + + result.final_event_hash = append_audit( + targets[0], event_type="reshard", body=body + ) + + # 6) Stamp every target shard with the corpus_shard_count meta. + for t in targets: + with t: + set_corpus_shard_count(t, plan.target_M) + + # 7) Validate row counts vs expected snapshot. + actual = _validate_row_counts( + targets=targets, expected=expected_row_counts + ) + result.target_row_counts["aggregate"] = actual + + # 8) VACUUM. Reclaims fragmentation from the executemany INSERT + # pattern; halves on-disk footprint when run after a bulk load. + # Outside any transaction. + for t in targets: + t.execute("VACUUM") + return result + finally: + for t in targets: + t.close() + + +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 diff --git a/arborist/store.py b/arborist/store.py index e5ec614..9f5588e 100644 --- a/arborist/store.py +++ b/arborist/store.py @@ -1535,9 +1535,15 @@ def _canonical_json(obj) -> str: def get_meta(conn: sqlite3.Connection, key: str) -> str | None: - """Read a value from the per-DB meta table; None if missing.""" + """Read a value from the per-DB meta table; None if missing. + + Indexes by position so it works regardless of ``row_factory`` — + callers that pass a plain ``sqlite3.connect()`` connection + shouldn't have to remember to set ``row_factory = sqlite3.Row`` + just to read one meta value. + """ row = conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone() - return row["value"] if row else None + return row[0] if row else None def set_meta(conn: sqlite3.Connection, key: str, value: str) -> None: diff --git a/tests/test_migrate_executor.py b/tests/test_migrate_executor.py new file mode 100644 index 0000000..2d416e4 --- /dev/null +++ b/tests/test_migrate_executor.py @@ -0,0 +1,343 @@ +"""Integration tests for the strategy A 'all_at_once' executor (#000065). + +Builds a tiny 2-shard test corpus, extracts audit events, plans, and +executes a reshard into a 4-shard target layout. Verifies: + + * every document survives + * every document lands in the shard its document_root hashes to + * the FTS5 indexes serve search queries on the target + * the consolidated audit chain re-verifies on the target + * the reshard's own audit event appears at the tail of that chain + * every target shard has corpus_shard_count meta = M +""" +from __future__ import annotations + +import hashlib +import json +import sqlite3 +from pathlib import Path +from typing import Iterator + +import pytest + +from arborist.document import Document, shard_for_document +from arborist.ingest import ingest_source +from arborist.migrate import ( + HydrationPlanner, + execute_plan, + _table_exists, +) +from arborist.source import Source +from arborist.store import ( + connect, + get_corpus_shard_count, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +class _FakeSource(Source): + source_type = "html" + + def __init__(self, docs: list[Document]): + self.docs = docs + + def iter_documents(self) -> Iterator[Document]: + yield from self.docs + + +def _make_doc(uri: str, seed: str) -> Document: + # Long enough to chunk into 2+ chunks. ``seed`` varies content so + # document_root is distinct per doc (the Merkle root keys on + # CONTENT, not URI — identical content dedups in the cache). + content = ( + f"This is the document at {uri} with content seed {seed}. " * 10 + + f"Body text {seed} with enough tokens for proper chunking. " * 30 + ) + return Document( + uri=uri, content=content, source_type="html", title=uri.split("/")[-1] + ) + + +@pytest.fixture +def source_shards(tmp_path: Path) -> Path: + """Two source shards, ~15 docs each.""" + src = tmp_path / "src" + src.mkdir() + docs_per_shard = 15 + for shard_idx in range(2): + db = src / f"{shard_idx:03d}.db" + conn = connect(db) + try: + docs = [ + _make_doc( + f"https://example.com/shard{shard_idx}/doc-{i}", + seed=f"s{shard_idx}d{i}", + ) + for i in range(docs_per_shard) + ] + ingest_source(conn, _FakeSource(docs)) + finally: + conn.close() + return src + + +@pytest.fixture +def audit_ndjson(source_shards: Path, tmp_path: Path) -> Path: + """Extract the source-shards' audit chain to ndjson (what + bench/extract_audit_events.py would produce on the real host).""" + out = tmp_path / "audit-events.ndjson" + with out.open("w") as f: + for shard in sorted(source_shards.glob("00[0-9].db")): + idx = int(shard.stem) + c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True) + try: + cur = c.execute( + "SELECT seq, event_hash, prev_event_hash, event_type, " + "subject_root, body, ts FROM audit_events ORDER BY seq ASC" + ) + for seq, eh, peh, etype, subj, body, ts in cur: + f.write(json.dumps({ + "src_shard": idx, + "src_seq": seq, + "src_event_hash": eh, + "prev_event_hash": peh, + "event_type": etype, + "subject_root": subj, + "body": body, + "ts": ts, + }, separators=(",", ":"))) + f.write("\n") + finally: + c.close() + return out + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestMigrationAllAtOnce: + @pytest.fixture + def migrated(self, source_shards: Path, audit_ndjson: Path, tmp_path: Path): + target = tmp_path / "tgt" + plan = HydrationPlanner().plan( + source_dir=source_shards, + target_dir=target, + target_M=4, + free_bytes=200 * 1024 ** 3, # plenty + ) + assert plan.strategy == "all_at_once" + result = execute_plan( + plan, audit_events_ndjson=audit_ndjson, + ) + return target, plan, result + + def test_target_dir_has_M_shards(self, migrated): + target, plan, _ = migrated + assert sorted(p.name for p in target.glob("*.db")) == [ + "000.db", "001.db", "002.db", "003.db" + ] + + def test_every_source_doc_lands_in_correct_shard( + self, source_shards: Path, migrated + ): + target, plan, _ = migrated + # Build the expected mapping from source shards. + expected: dict[str, int] = {} + for shard in sorted(source_shards.glob("*.db")): + c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True) + try: + cur = c.execute("SELECT document_root FROM documents") + for (root,) in cur: + expected[root] = shard_for_document(root, plan.target_M) + finally: + c.close() + + # Verify each target shard holds exactly the docs hashed to it. + for tidx in range(plan.target_M): + tdb = target / f"{tidx:03d}.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + actual = { + r[0] for r in c.execute("SELECT document_root FROM documents") + } + finally: + c.close() + for root, expect_idx in expected.items(): + if expect_idx == tidx: + assert root in actual, ( + f"doc {root[:8]}... expected on shard {tidx} but missing" + ) + else: + assert root not in actual, ( + f"doc {root[:8]}... on shard {tidx} but hashes to {expect_idx}" + ) + + def test_every_chunk_followed_its_document(self, migrated): + target, plan, _ = migrated + # For each target shard, every chunk's document_root must route here. + for tidx in range(plan.target_M): + tdb = target / f"{tidx:03d}.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + cur = c.execute( + "SELECT DISTINCT document_root FROM chunks" + ) + for (root,) in cur: + assert shard_for_document(root, plan.target_M) == tidx, ( + f"chunk doc {root[:8]}... on shard {tidx} but routes " + f"to {shard_for_document(root, plan.target_M)}" + ) + finally: + c.close() + + def test_audit_chain_valid_on_target(self, migrated): + target, plan, result = migrated + tdb = target / "000.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + cur = c.execute( + "SELECT seq, event_hash, prev_event_hash, body " + "FROM audit_events ORDER BY seq ASC" + ) + prev = None + for seq, event_hash, peh, body in cur: + assert peh == prev, f"chain break at seq {seq}" + h = hashlib.sha256() + if prev is not None: + h.update(bytes.fromhex(prev)) + h.update(body.encode("utf-8", errors="surrogatepass")) + assert h.hexdigest() == event_hash, ( + f"hash mismatch at seq {seq}" + ) + prev = event_hash + finally: + c.close() + + def test_reshard_event_at_tail(self, migrated): + target, plan, result = migrated + tdb = target / "000.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + row = c.execute( + "SELECT event_type, body FROM audit_events " + "ORDER BY seq DESC LIMIT 1" + ).fetchone() + finally: + c.close() + assert row[0] == "reshard" + body = json.loads(row[1]) + assert body["kind"] == "reshard" + assert body["plan"]["strategy"] == "all_at_once" + assert body["plan"]["target_M"] == 4 + assert body["result"]["rows_moved_total"] > 0 + + def test_corpus_shard_count_set_on_every_target(self, migrated): + target, plan, _ = migrated + for tidx in range(plan.target_M): + tdb = target / f"{tidx:03d}.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + assert get_corpus_shard_count(c) == plan.target_M, ( + f"shard {tidx} missing corpus_shard_count meta" + ) + finally: + c.close() + + def test_chunks_fts_searchable_on_target(self, migrated): + target, plan, _ = migrated + # Verify FTS5 index serves a hit on every target shard that has + # documents — the test content uses the word "document" liberally. + hit_count = 0 + for tidx in range(plan.target_M): + tdb = target / f"{tidx:03d}.db" + c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True) + try: + row = c.execute( + "SELECT COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH ?", + ("document",), + ).fetchone() + hit_count += int(row[0]) + finally: + c.close() + assert hit_count > 0, "no chunks_fts hits across any target shard" + + def test_total_doc_count_preserved( + self, source_shards: Path, migrated + ): + target, _, _ = migrated + # 2 source × 15 docs = 30 unique docs. + src_total = 0 + for shard in source_shards.glob("*.db"): + c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True) + try: + src_total += c.execute( + "SELECT COUNT(*) FROM documents" + ).fetchone()[0] + finally: + c.close() + tgt_total = 0 + for shard in target.glob("*.db"): + c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True) + try: + tgt_total += c.execute( + "SELECT COUNT(*) FROM documents" + ).fetchone()[0] + finally: + c.close() + assert tgt_total == src_total, ( + f"doc count drift: src={src_total} tgt={tgt_total}" + ) + + def test_dry_run_creates_no_targets( + self, source_shards: Path, audit_ndjson: Path, tmp_path: Path + ): + target = tmp_path / "tgt-dry" + plan = HydrationPlanner().plan( + source_dir=source_shards, + target_dir=target, + target_M=4, + free_bytes=200 * 1024 ** 3, + ) + result = execute_plan( + plan, audit_events_ndjson=audit_ndjson, dry_run=True + ) + assert "DRY RUN" in " ".join(result.notes) + # target dir exists but no DBs. + assert list(target.glob("*.db")) == [] + + +class TestExecutorErrorCases: + def test_missing_ndjson_raises( + self, source_shards: Path, tmp_path: Path + ): + target = tmp_path / "tgt" + plan = HydrationPlanner().plan( + source_dir=source_shards, + target_dir=target, + target_M=4, + free_bytes=200 * 1024 ** 3, + ) + with pytest.raises(FileNotFoundError, match="audit-event ndjson"): + execute_plan( + plan, audit_events_ndjson=tmp_path / "does-not-exist.ndjson" + ) + + def test_per_source_shard_not_implemented_yet( + self, source_shards: Path, audit_ndjson: Path, tmp_path: Path + ): + target = tmp_path / "tgt" + plan = HydrationPlanner().plan( + source_dir=source_shards, + target_dir=target, + target_M=4, + free_bytes=200 * 1024 ** 3, + force_strategy="per_source_shard", + ) + with pytest.raises(NotImplementedError, match="per_source_shard"): + execute_plan(plan, audit_events_ndjson=audit_ndjson)