arborist/tests/test_migrate_executor.py
russell@unturf.com c86d5ac4f6
#000065 follow-up #48: WAL checkpoint between executor phases
Add PRAGMA wal_checkpoint(TRUNCATE) at two points in
_execute_all_at_once so committed WAL pages don't pin disk through
subsequent passes. Production migration on 2026-05-26 hit 7 GB free
disk (down from 89) because SQLite's auto-checkpoint can't run while
a reader cursor is open, and FTS rebuild keeps a SELECT cursor open
through 1.5M chunks per target. Across 4 targets the FTS rebuild
plus audit consolidate held ~37 GB of committed-but-unreclaimed WAL.
Manual sibling-connection wal_checkpoint(TRUNCATE) freed 27 GB
mid-migration.

Checkpoints land at:
  * end of _rebuild_fts_on_target (after the SELECT cursor is
    explicitly cur.close()'d so the TRUNCATE checkpoint can actually
    fire — TRUNCATE/RESTART block on active readers)
  * end of _consolidate_audit_chain (after the 3.47M-row giant
    transaction commits, before the next phase touches the same
    connection)
VACUUM is already implicitly a checkpoint, so the existing per-
target VACUUM pass continues to handle the final checkpoint
naturally.

Helper _checkpoint_truncate(conn) returns the (busy, log_frames,
checkpointed) tuple SQLite emits; for the serial executor, busy=1
is improbable since each phase finishes before moving on.

Regression test
(TestWalCheckpointing.test_no_large_wal_after_migration) asserts
no WAL file exceeds 4 MB after migration completes. Without the
checkpoint calls this would fail on real-sized corpora; with them
the test passes deterministically.

Doesn't affect the running migration (it loaded the module from
memory before this commit). Future reshards run with bounded WAL —
no near-ENOSPC scares.
2026-05-26 15:34:40 -04:00

586 lines
22 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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 TestInPlacePromotion:
"""Verify the .db.new → .db atomic-promote dance (α-shape)."""
def test_target_dir_only_has_db_files_after_completion(
self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
):
"""No .new sidecars left after a successful run."""
target = tmp_path / "tgt"
plan = HydrationPlanner().plan(
source_dir=source_shards,
target_dir=target,
target_M=4,
free_bytes=200 * 1024 ** 3,
)
execute_plan(plan, audit_events_ndjson=audit_ndjson)
files = sorted(p.name for p in target.glob("*"))
# 4 final .db files; -wal/-shm sidecars after close are
# permissible but no .db.new files.
assert all(not f.endswith(".db.new") for f in files), files
assert sum(1 for f in files if f.endswith(".db")) == 4
def test_inplace_reshard_overwrites_originals(
self, source_shards: Path, audit_ndjson: Path
):
"""Plan with target_dir == source_dir overwrites originals."""
# Capture source roots so we can later detect topology change.
source_roots: set[str] = set()
for shard in source_shards.glob("*.db"):
c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
try:
for (root,) in c.execute("SELECT document_root FROM documents"):
source_roots.add(root)
finally:
c.close()
plan = HydrationPlanner().plan(
source_dir=source_shards,
target_dir=source_shards, # in-place
target_M=4,
free_bytes=200 * 1024 ** 3,
)
execute_plan(plan, audit_events_ndjson=audit_ndjson)
# After in-place: target_dir has exactly M files (no .new, no
# parallel dir), each carrying corpus_shard_count = 4, and
# each doc lands on its hashed shard.
db_files = sorted(p.name for p in source_shards.glob("*.db"))
assert db_files == ["000.db", "001.db", "002.db", "003.db"]
assert not list(source_shards.glob("*.db.new"))
for tidx in range(4):
tdb = source_shards / f"{tidx:03d}.db"
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
try:
# corpus_shard_count meta set.
row = c.execute(
"SELECT value FROM meta WHERE key='corpus_shard_count'"
).fetchone()
assert row and int(row[0]) == 4
# All docs route to this shard.
for (root,) in c.execute("SELECT document_root FROM documents"):
assert shard_for_document(root, 4) == tidx
finally:
c.close()
def test_validation_failure_leaves_new_files(
self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
):
"""When row-count validation fails, .db.new files survive for
inspection and no atomic rename runs."""
from arborist.migrate import execute_plan as run
target = tmp_path / "tgt"
plan = HydrationPlanner().plan(
source_dir=source_shards,
target_dir=target,
target_M=4,
free_bytes=200 * 1024 ** 3,
)
# Pass a deliberately wrong expected row count to trip the
# >1% tolerance guard.
with pytest.raises(RuntimeError, match="validation failed"):
run(
plan,
audit_events_ndjson=audit_ndjson,
expected_row_counts={"documents": 9_999_999},
)
# Inspection artefacts must remain.
assert sorted(p.name for p in target.glob("*.db.new")) == [
"000.db.new", "001.db.new", "002.db.new", "003.db.new",
]
# No promoted .db files yet.
assert not list(target.glob("00[0-9].db"))
class TestWalCheckpointing:
"""Regression test: WAL files must be reclaimed between executor
phases so they don't pin disk through subsequent passes. The
2026-05-26 production migration came within 7 GB of ENOSPC because
SQLite auto-checkpoint can't run while a reader cursor is open and
the FTS rebuild keeps a cursor open through every chunk. Manual
PRAGMA wal_checkpoint(TRUNCATE) freed 27 GB."""
def test_no_large_wal_after_migration(
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,
)
execute_plan(plan, audit_events_ndjson=audit_ndjson)
# After completion, no .db-wal file should retain more than a
# few MB. The checkpoint(TRUNCATE) calls inside the executor
# collapse the WAL between phases; final VACUUM produces a
# zero-WAL state.
for i in range(plan.target_M):
for suffix in (".db-wal", ".db.new-wal"):
wal = target / f"{i:03d}{suffix}"
if wal.exists():
assert wal.stat().st_size < 4 * 1024 * 1024, (
f"WAL retained {wal.stat().st_size / 1e6:.1f} MB "
f"at {wal} — checkpoint between phases didn't fire"
)
class TestCrossShardForeignKeys:
"""Regression test: derivations.src_root can legitimately reference
a surface doc that hash-routes to a different target shard. FK
enforcement on the migration writer connection must be OFF or the
INSERT fails (the 2026-05-26 cutover-crash reproducer)."""
def test_cross_shard_derivation_succeeds(
self, tmp_path: Path
):
"""Synthesize a derivation row whose core_root and src_root
hash to different target shards, run migration, verify it
lands without FK violation."""
# Build a source shard with two docs A and B that we KNOW hash
# to different target shards under M=4. Then INSERT a
# derivation linking them, run migration, expect success.
src_dir = tmp_path / "src"
src_dir.mkdir()
src_db = src_dir / "000.db"
conn = connect(src_db)
try:
# Two docs with content chosen so their roots split across
# M=4 targets. Find them by trial.
roots: dict[int, str] = {}
for i in range(200):
content = f"crosshard fixture seed {i} body text " * 30
ds = _FakeSource([
Document(
uri=f"https://test/cs-{i}",
content=content,
source_type="html",
title=f"cs-{i}",
)
])
ingest_source(conn, ds)
row = conn.execute(
"SELECT document_root FROM documents WHERE document_uri = ?",
(f"https://test/cs-{i}",),
).fetchone()
root = row[0]
idx = shard_for_document(root, 4)
roots.setdefault(idx, root)
if len(roots) >= 2:
break
assert len(roots) >= 2, "couldn't find two docs hashing to different shards"
two_idxs = sorted(roots.keys())[:2]
core_root = roots[two_idxs[0]]
src_root = roots[two_idxs[1]]
# Synthesize a derivation linking core_root (target A) to
# src_root (target B != A).
conn.execute(
"INSERT INTO derivations "
"(core_root, src_root, proof_blob, process_id, distilled_at) "
"VALUES (?, ?, ?, ?, ?)",
(core_root, src_root, "{}", "test_cross_shard", 0),
)
conn.commit()
finally:
conn.close()
# Extract audit events.
ndjson = tmp_path / "a.ndjson"
with ndjson.open("w") as f:
c = sqlite3.connect(f"file:{src_db}?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"
)
for seq, eh, peh, et, subj, body, ts in cur:
f.write(json.dumps({
"src_shard": 0, "src_seq": seq, "src_event_hash": eh,
"prev_event_hash": peh, "event_type": et,
"subject_root": subj, "body": body, "ts": ts,
}, separators=(",", ":")) + "\n")
finally:
c.close()
target = tmp_path / "tgt"
plan = HydrationPlanner().plan(
source_dir=src_dir, target_dir=target, target_M=4,
free_bytes=200 * 1024 ** 3,
)
# The reproducer: pre-fix this raised sqlite3.IntegrityError
# ("FOREIGN KEY constraint failed") inside _route_per_doc_table.
from arborist.migrate import execute_plan as run
run(plan, audit_events_ndjson=ndjson)
# Verify: derivation row lands on the core_root's target shard
# (where the row is routed by `core_root`); src_root reference
# points cross-shard.
tdb = target / f"{shard_for_document(core_root, 4):03d}.db"
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
try:
row = c.execute(
"SELECT core_root, src_root FROM derivations "
"WHERE process_id = 'test_cross_shard'"
).fetchone()
finally:
c.close()
assert row is not None
assert row[0] == core_root
assert row[1] == src_root
# src_root is NOT in this same target shard — that's the
# cross-shard semantic this test exists to protect.
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
try:
row = c.execute(
"SELECT document_root FROM documents WHERE document_root = ?",
(src_root,),
).fetchone()
finally:
c.close()
assert row is None, "src_root unexpectedly local; fixture didn't cross-shard"
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)