engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

#000065 step 3: strategy A executor + CLI + plan-only preview

Browse source 
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.
This commit is contained in:
russell@unturf.com 2026-05-26 13:28:52 -04:00
parent 8fe389d45e
commit 77b90bfcf1
No known key found for this signature in database
4 changed files with 941 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

137
arborist/cli.py
View file

@ -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)",

455
arborist/migrate.py
View file

@ -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

10
arborist/store.py
View file

@ -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:

343
tests/test_migrate_executor.py Normal file
View file

@ -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)