russell@unturf.com
c26d03956d
Reshard no longer creates a parallel "shards.v2/" directory. Builds
write to "<target_dir>/00X.db.new" alongside originals; when the
executor's validation passes, each .db.new is atomically renamed to
its final 00X.db name via os.replace (POSIX-atomic per file).
For an in-place migration (target_dir == source_dir, which is the
canonical use case), the rename REPLACES the original shard files.
~/.arborist/shards/ never contains a parallel "v2" or "next" or
"backup" directory — it always holds exactly one corpus, just with
the new topology after promote completes.
Validation gate (refuses to promote on mismatch):
* --expected-row-counts <snapshot.json> threads the pre-migration
snapshot's documents/chunks/edges totals through to the executor.
* Tolerance: ±1% to absorb the dupe-collapse from INSERT OR IGNORE
on cross-shard duplicate document_roots (166 dupes measured
pre-migration; <0.005% drift).
* Mismatch → RuntimeError, .db.new files left in place for
inspection, no rename performed.
CLI:
arborist corpus reshard --to 4 \
--source-dir ~/.arborist/shards \
--target-dir ~/.arborist/shards \
--audit-events-ndjson /tmp/audit-events.ndjson \
--expected-row-counts bench/results/pre-migration-snapshot.json
3 new tests:
* test_target_dir_only_has_db_files_after_completion — no .db.new
sidecars survive a successful run
* test_inplace_reshard_overwrites_originals — target_dir ==
source_dir works end-to-end; final files are the new hash-routed
shards
* test_validation_failure_leaves_new_files — bad expected counts
trip the validation guard; .db.new files survive; no .db files
promoted
Stale .db.new files from a prior failed run are unlinked before
opening fresh targets, so a partial-fail-and-retry is idempotent.
WAL/SHM sidecars removed at promote time so the new live .db
produces fresh sidecars on next open.
437 lines
16 KiB
Python
437 lines
16 KiB
Python
"""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 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)
|