diff --git a/arborist/cli.py b/arborist/cli.py index 7b5b6b2..d355c3b 100644 --- a/arborist/cli.py +++ b/arborist/cli.py @@ -2857,13 +2857,31 @@ def _cmd_corpus_reshard(args: argparse.Namespace) -> int: ) return 3 + expected_row_counts: dict[str, int] = {} + if args.expected_row_counts_json: + snap_path = Path(args.expected_row_counts_json).expanduser() + if not snap_path.exists(): + print(f"snapshot not found at {snap_path}", file=sys.stderr) + return 5 + snap = json.loads(snap_path.read_text()) + totals = snap.get("totals", {}) + for table in ("documents", "chunks", "edges"): + if table in totals: + expected_row_counts[table] = int(totals[table]) + try: result = execute_plan( - plan, audit_events_ndjson=audit_ndjson, dry_run=args.dry_run + plan, + audit_events_ndjson=audit_ndjson, + expected_row_counts=expected_row_counts, + dry_run=args.dry_run, ) except NotImplementedError as e: print(f"strategy not implemented: {e}", file=sys.stderr) return 4 + except RuntimeError as e: + print(f"validation failed: {e}", file=sys.stderr) + return 6 print( json.dumps( @@ -5930,6 +5948,15 @@ def build_parser() -> argparse.ArgumentParser: help="do everything except actually open the target DBs and move " "rows. Plan still emits, target dir still created.", ) + reshard.add_argument( + "--expected-row-counts", dest="expected_row_counts_json", + default=None, + help="path to a pre-migration snapshot JSON " + "(e.g., bench/results/pre-migration-snapshot.json). If provided, " + "the executor compares post-migration row totals against the " + "snapshot's `totals` block within ±1%% tolerance, and refuses to " + "promote .db.new files into final names on mismatch.", + ) reshard.set_defaults(func=_cmd_corpus_reshard) activity_cmd = sub.add_parser( diff --git a/arborist/migrate.py b/arborist/migrate.py index 48c9db1..ae9226c 100644 --- a/arborist/migrate.py +++ b/arborist/migrate.py @@ -675,8 +675,16 @@ def _execute_all_at_once( strategy="all_at_once", started_at=started, finished_at=started ) - # 1) Open M target connections, each pre-schema-applied. - target_paths = [ + # 1) Build paths carry a `.new` suffix during construction. Final + # paths get the canonical `00X.db` names only after the executor + # validates the output. This lets the migration coexist with the + # source shards in the SAME directory (fox's α-shape preference): + # only one ``~/.arborist/shards/`` directory ever, originals get + # replaced atomically by ``os.replace`` after validation. + build_paths = [ + plan.target_dir / f"{i:03d}.db.new" for i in range(plan.target_M) + ] + final_paths = [ plan.target_dir / f"{i:03d}.db" for i in range(plan.target_M) ] if dry_run: @@ -684,7 +692,14 @@ def _execute_all_at_once( result.finished_at = time.time() return result - targets = [_connect_target(p) for p in target_paths] + # Clear any stale .new files from a prior failed run before opening. + for p in build_paths: + for sidecar_suffix in ("", "-wal", "-shm"): + sc = Path(str(p) + sidecar_suffix) + if sc.exists(): + sc.unlink() + + targets = [_connect_target(p) for p in build_paths] try: # 2) Per-source-shard: stream every table, route or consolidate. for src_path in plan.corpus_facts.shard_files: @@ -760,17 +775,88 @@ def _execute_all_at_once( targets=targets, expected=expected_row_counts ) result.target_row_counts["aggregate"] = actual + if expected_row_counts: + _check_validation_passes( + actual=actual, expected=expected_row_counts, result=result + ) # 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() + # 9) Promote .new → final names. os.replace is POSIX-atomic per + # file and overwrites an existing destination; for in-place + # migrations (source_dir == target_dir) this replaces the original + # shards on disk. Done outside the try/finally so a validation + # failure leaves .new files in place for inspection. + _promote_build_paths(build_paths=build_paths, final_paths=final_paths) + result.notes.append( + f"promoted {len(build_paths)} .db.new files → final names" + ) + return result + + +def _check_validation_passes( + *, + actual: dict[str, int], + expected: dict[str, int], + result: ExecutionResult, + tolerance: float = 0.01, +) -> None: + """Refuse to promote if target counts deviate >1% from expected. + + Cross-shard duplicates from re-ingest history cause exact-match to + drift slightly (INSERT OR IGNORE collapses dupes). 1% tolerance + covers that without masking real data loss. + """ + for table, expected_n in expected.items(): + actual_n = actual.get(table) + if actual_n is None: + continue + if expected_n == 0: + continue + delta = abs(actual_n - expected_n) / expected_n + if delta > tolerance: + raise RuntimeError( + f"validation failed: {table} expected ~{expected_n:,} but " + f"got {actual_n:,} (delta {delta:.1%}, tolerance " + f"{tolerance:.0%}). Leaving .db.new files in place for " + f"inspection; no atomic rename performed." + ) + result.notes.append( + f"{table}: actual {actual_n:,} vs expected {expected_n:,} " + f"(delta {delta:.2%}, within {tolerance:.0%} tolerance)" + ) + + +def _promote_build_paths( + *, + build_paths: list[Path], + final_paths: list[Path], +) -> None: + """Atomic rename each ``.db.new`` to its final ``.db`` name. + + Per-file POSIX atomic (``os.replace``). For an in-place migration + this overwrites the original shard files. Sidecar files + (``-wal``, ``-shm``) are removed first since they belong to the + closed build connection and the new live ``.db`` will produce + fresh sidecars on open. + """ + for build, final in zip(build_paths, final_paths): + for suffix in ("-wal", "-shm"): + sc_build = Path(str(build) + suffix) + if sc_build.exists(): + sc_build.unlink() + sc_final = Path(str(final) + suffix) + if sc_final.exists(): + sc_final.unlink() + os.replace(build, final) + def _table_exists(conn: sqlite3.Connection, name: str) -> bool: row = conn.execute( diff --git a/tests/test_migrate_executor.py b/tests/test_migrate_executor.py index 2d416e4..ff7d640 100644 --- a/tests/test_migrate_executor.py +++ b/tests/test_migrate_executor.py @@ -312,6 +312,100 @@ class TestMigrationAllAtOnce: 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