#000067: M-aware cold-pack hydration (route incoming docs by content hash)
Today's hydrate_from_metadata_pack writes every row into one shard. With the corpus in M=4 hash-routed topology (#000065), a fresh peer pulling packs must land each doc on shard_for_document(root, M) — same routing function as the producer — or the consumer's M=4 ATTACH-and-route assumption is just decoration over a single-shard reality. Code (new entry points alongside the existing single-conn ones): arborist/cold_pack_metadata.py + restore_shard_metadata_routed(targets, M, table_dir) + _restore_routed_table (per-document tables) + _restore_edges_fan_out_routed (edges by src_root) + _shard_idx_for_root helper (mirrors arborist.document) + routing rules: _ROUTED_BY_COL / _CONSOLIDATED_TO_SHARD_0 (mirror migrate.py's ROUTED_BY_DOCUMENT_ROOT / CONSOLIDATED_TABLES) arborist/evict.py + hydrate_from_metadata_pack_routed(targets, backend, hash, M=, mode=) + _pull_pack_inner_routed (mirrors _pull_pack_inner; phase 2 chunk-body fill iterates every target shard — leaf_hash lookup naturally hits at most one since each chunk's metadata row landed on exactly one target during phase 1) arborist/cli.py arborist cold unpack + --hydrate-shards-dir DIR (M-aware genesis path) + --hydrate-M N (default 4 matches #000065) legacy --db / --global-shards-dir path unchanged Behaviour notes: - FK enforcement off on target writes (cross-shard refs are valid under hash routing, same fix as #000065 reshard executor) - audit_events lands on target 0 (Option A consolidation) - corpus-wide tables (snapshots / concepts / aliases / providence) consolidate to target 0 - per-document tables route by document_root (documents, document_http_meta, chunks, merkle_nodes) - edges route by src_root (matches migrate.py) - derivations route by core_root (matches migrate.py) - existing single-conn API unchanged — callers that didn't pass a shards-dir get the legacy single-shard behaviour 4 new tests: test_pack_then_hydrate_routed — end-to-end pack → hydrate → assert every doc on its hash-routed target, no doc on the wrong shard test_corpus_shard_count_set_on_routed_hydrate — meta plumbing test_M_mismatch_rejected test_empty_targets_rejected All 59 prior tests still pass. Refactor opportunity (not taken): _ROUTED_BY_COL duplicates migrate.py's ROUTED_BY_DOCUMENT_ROOT. A shared arborist/multi_shard.py module would serve both reshard and graft (#000066). Left as follow-up since the duplication is small and graft is still scaffold. Unblocks #46 genesis on 3090: that's now a single arborist cold unpack --hydrate-shards-dir ~/.arborist/shards --hydrate-M 4 invocation instead of the α two-step kludge (hydrate-then-reshard).
This commit is contained in:
parent
9cfb9c8d01
commit
e1f291a3db
4 changed files with 668 additions and 1 deletions
|
|
@ -2982,9 +2982,41 @@ def _cmd_cold_unpack(args: argparse.Namespace) -> int:
|
|||
packs are anonymous from the consumer's perspective — accessed via
|
||||
the metadata pack's reference list, not by name.
|
||||
"""
|
||||
from arborist.evict import hydrate_from_metadata_pack
|
||||
from arborist.evict import (
|
||||
hydrate_from_metadata_pack,
|
||||
hydrate_from_metadata_pack_routed,
|
||||
)
|
||||
from arborist.store import set_corpus_shard_count
|
||||
|
||||
backend = _make_cold_backend()
|
||||
|
||||
# M-aware genesis path (#000067): --hydrate-shards-dir is the target
|
||||
# directory; --hydrate-M picks the routing modulus. Each incoming
|
||||
# row lands on shard_for_document(root, M).
|
||||
if args.hydrate_shards_dir:
|
||||
target_dir = Path(args.hydrate_shards_dir).expanduser()
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
M = args.hydrate_M
|
||||
target_paths = [target_dir / f"{i:03d}.db" for i in range(M)]
|
||||
targets: list = []
|
||||
for p in target_paths:
|
||||
t = connect(p)
|
||||
t.execute("PRAGMA foreign_keys = OFF")
|
||||
targets.append(t)
|
||||
try:
|
||||
result = hydrate_from_metadata_pack_routed(
|
||||
targets, backend, args.pack_hash, M=M, mode=args.mode,
|
||||
)
|
||||
for t in targets:
|
||||
with t:
|
||||
set_corpus_shard_count(t, M)
|
||||
finally:
|
||||
for t in targets:
|
||||
t.close()
|
||||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
# Single-shard legacy path (no --hydrate-shards-dir).
|
||||
conn = (
|
||||
connect_query(args.db, shards_dir=args.global_shards_dir)
|
||||
if args.global_shards_dir
|
||||
|
|
@ -5873,6 +5905,18 @@ def build_parser() -> argparse.ArgumentParser:
|
|||
const="just-enough",
|
||||
help="(default) metadata only — chunks remain NULL for JIT fetch",
|
||||
)
|
||||
cold_unpack.add_argument(
|
||||
"--hydrate-shards-dir", dest="hydrate_shards_dir", default=None,
|
||||
help="(M-aware genesis, #000067) hydrate into M target shards "
|
||||
"under this directory; each incoming row routes by "
|
||||
"shard_for_document(root, M). Mutually exclusive with --db / "
|
||||
"--global-shards-dir.",
|
||||
)
|
||||
cold_unpack.add_argument(
|
||||
"--hydrate-M", dest="hydrate_M", type=int, default=4,
|
||||
help="number of target shards for --hydrate-shards-dir routing "
|
||||
"(default 4 matches #000065 canonical M).",
|
||||
)
|
||||
cold_unpack.set_defaults(func=_cmd_cold_unpack)
|
||||
|
||||
cold_stats = cold_sub.add_parser(
|
||||
|
|
|
|||
|
|
@ -468,3 +468,177 @@ def _restore_generic_table(
|
|||
if batch and insert_sql is not None:
|
||||
conn.executemany(insert_sql, batch)
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-shard restore (#000067).
|
||||
#
|
||||
# Routes incoming rows from one JSONL file across M target shards by content
|
||||
# hash, mirroring the producer's shard_for_document layout. Mirrors the
|
||||
# routing rules used by arborist.migrate so a packed corpus and a resharded
|
||||
# corpus end up with byte-identical per-shard placement (modulo edge
|
||||
# fan-in/fan-out ordering, which is content-address-equivalent).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Per-table routing decisions. The KEY is the table name; the VALUE is the
|
||||
# column that carries the content-addressed root used to pick a target
|
||||
# shard. Tables not in this dict consolidate to canonical shard 0.
|
||||
_ROUTED_BY_COL: dict[str, str] = {
|
||||
"documents": "document_root",
|
||||
"document_http_meta": "document_root",
|
||||
"chunks": "document_root",
|
||||
"merkle_nodes": "document_root",
|
||||
"derivations": "core_root",
|
||||
# edges handled separately via fan-in/fan-out and routed by src_root.
|
||||
}
|
||||
|
||||
# Tables whose rows ALL go to target shard 0. Audit chain is consolidated
|
||||
# there (Option A from #000065); corpus-wide derived layers (snapshots,
|
||||
# concepts, aliases) follow the same convention so they're queryable from
|
||||
# one well-known location.
|
||||
_CONSOLIDATED_TO_SHARD_0: frozenset[str] = frozenset({
|
||||
"audit_events",
|
||||
"snapshots",
|
||||
"concept_relations",
|
||||
"concept_token_idf",
|
||||
"citation_aliases",
|
||||
"term_aliases",
|
||||
"providence_cache",
|
||||
"falsifications",
|
||||
})
|
||||
|
||||
|
||||
def _shard_idx_for_root(root: str, M: int) -> int:
|
||||
"""Pure-function routing — same formula as
|
||||
``arborist.document.shard_for_document`` but local to avoid pulling the
|
||||
full migrate module into this metadata library."""
|
||||
if M < 1:
|
||||
raise ValueError(f"M must be >= 1, got {M}")
|
||||
if M == 1:
|
||||
return 0
|
||||
return int(root[:8], 16) % M
|
||||
|
||||
|
||||
def _restore_routed_table(
|
||||
targets: list[sqlite3.Connection],
|
||||
M: int,
|
||||
table: str,
|
||||
in_path: Path,
|
||||
route_col: str,
|
||||
) -> int:
|
||||
"""Stream JSONL rows for ``table``, route each by ``route_col`` hash
|
||||
to one of ``M`` target shards, executemany INSERT into the chosen
|
||||
target. Mirrors ``arborist.migrate._route_per_doc_table`` but reads
|
||||
from JSONL instead of a source connection."""
|
||||
BATCH = 5000
|
||||
buckets: list[list[tuple]] = [[] for _ in range(M)]
|
||||
cols: list[str] | None = None
|
||||
insert_sql: str | None = None
|
||||
n = 0
|
||||
for row in read_columnar_jsonl(in_path):
|
||||
if cols is None:
|
||||
cols = list(row.keys())
|
||||
placeholders = ", ".join("?" for _ in cols)
|
||||
col_list = ", ".join(f'"{c}"' for c in cols)
|
||||
insert_sql = (
|
||||
f"INSERT OR IGNORE INTO {table} ({col_list}) VALUES ({placeholders})"
|
||||
)
|
||||
root = row[route_col]
|
||||
if not root:
|
||||
continue
|
||||
idx = _shard_idx_for_root(root, M)
|
||||
buckets[idx].append(tuple(row[c] for c in cols))
|
||||
n += 1
|
||||
# Flush any bucket that crossed the batch threshold so we don't
|
||||
# hold millions of rows in Python heap.
|
||||
for i, bucket in enumerate(buckets):
|
||||
if len(bucket) >= BATCH and insert_sql is not None:
|
||||
targets[i].executemany(insert_sql, bucket)
|
||||
bucket.clear()
|
||||
if insert_sql is not None:
|
||||
for i, bucket in enumerate(buckets):
|
||||
if bucket:
|
||||
targets[i].executemany(insert_sql, bucket)
|
||||
return n
|
||||
|
||||
|
||||
def _restore_edges_fan_out_routed(
|
||||
targets: list[sqlite3.Connection],
|
||||
M: int,
|
||||
in_path: Path,
|
||||
) -> int:
|
||||
"""Expand fan-in rows back into one edge per src_root, routing each
|
||||
edge to the shard its src_root hashes to. Mirrors
|
||||
``_restore_edges_fan_out`` but multi-shard."""
|
||||
BATCH = 5000
|
||||
buckets: list[list[tuple]] = [[] for _ in range(M)]
|
||||
sql = (
|
||||
"INSERT OR IGNORE INTO edges "
|
||||
"(src_root, dst_root, dst_uri, edge_type, anchor) "
|
||||
"VALUES (?, ?, ?, ?, ?)"
|
||||
)
|
||||
n = 0
|
||||
for fan_in in read_columnar_jsonl(in_path):
|
||||
dst_uri = fan_in["dst_uri"]
|
||||
edge_type = fan_in["edge_type"]
|
||||
anchor = fan_in["anchor"]
|
||||
dst_root = fan_in["dst_root"]
|
||||
for src_root in fan_in["src_roots"]:
|
||||
if not src_root:
|
||||
continue
|
||||
idx = _shard_idx_for_root(src_root, M)
|
||||
buckets[idx].append((src_root, dst_root, dst_uri, edge_type, anchor))
|
||||
n += 1
|
||||
if len(buckets[idx]) >= BATCH:
|
||||
targets[idx].executemany(sql, buckets[idx])
|
||||
buckets[idx].clear()
|
||||
for i, bucket in enumerate(buckets):
|
||||
if bucket:
|
||||
targets[i].executemany(sql, bucket)
|
||||
return n
|
||||
|
||||
|
||||
def restore_shard_metadata_routed(
|
||||
targets: list[sqlite3.Connection],
|
||||
M: int,
|
||||
table_dir: Path,
|
||||
) -> dict[str, int]:
|
||||
"""M-aware variant of :func:`restore_shard_metadata` (#000067).
|
||||
|
||||
For each shipped table, route every row to the target shard its
|
||||
content hash addresses. Per-document tables route by their root
|
||||
column (``document_root`` / ``src_root`` / ``core_root``); corpus-wide
|
||||
tables (audit chain, snapshots, concept layer, aliases, providence
|
||||
cache) consolidate to canonical target shard 0.
|
||||
|
||||
Targets must already have v9.8 schema applied. Caller is responsible
|
||||
for opening/closing the connections.
|
||||
|
||||
Returns ``{table: total_rows_routed}``.
|
||||
"""
|
||||
if not targets:
|
||||
raise ValueError("at least one target connection required")
|
||||
if M != len(targets):
|
||||
raise ValueError(
|
||||
f"M={M} but len(targets)={len(targets)} — must match"
|
||||
)
|
||||
result: dict[str, int] = {}
|
||||
for table in SHIPPED_TABLES:
|
||||
in_path = table_dir / f"{table}.jsonl"
|
||||
if not in_path.exists():
|
||||
continue
|
||||
if table == "edges":
|
||||
n = _restore_edges_fan_out_routed(targets, M, in_path)
|
||||
elif table in _ROUTED_BY_COL:
|
||||
n = _restore_routed_table(
|
||||
targets, M, table, in_path, _ROUTED_BY_COL[table]
|
||||
)
|
||||
elif table in _CONSOLIDATED_TO_SHARD_0:
|
||||
n = _restore_generic_table(targets[0], table, in_path)
|
||||
else:
|
||||
# Unknown classification — default to consolidate to 0 (safe).
|
||||
n = _restore_generic_table(targets[0], table, in_path)
|
||||
result[table] = n
|
||||
for t in targets:
|
||||
t.commit()
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -703,6 +703,241 @@ def hydrate_from_metadata_pack(
|
|||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M-aware hydration (#000067).
|
||||
#
|
||||
# Genesis a fresh peer in M=4 hash-routed topology by routing each incoming
|
||||
# row to the shard its content hash addresses. Uses
|
||||
# arborist.cold_pack_metadata.restore_shard_metadata_routed for the metadata
|
||||
# pass; chunk-body fill iterates every target shard (each only holds the
|
||||
# rows whose document_root routes there, so the leaf_hash lookup naturally
|
||||
# hits at most one target).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def hydrate_from_metadata_pack_routed(
|
||||
targets: list[sqlite3.Connection],
|
||||
backend: "ObjectStoreBackend",
|
||||
metadata_pack_hash: str,
|
||||
*,
|
||||
M: int,
|
||||
mode: str = "just-enough",
|
||||
) -> dict:
|
||||
"""M-aware variant of :func:`hydrate_from_metadata_pack`. Targets must
|
||||
already have v9.8 schema applied and ``M == len(targets)``. Each
|
||||
incoming row lands on the shard its content hash addresses; corpus-wide
|
||||
tables (audit chain, snapshots, concepts, aliases, providence cache)
|
||||
consolidate to ``targets[0]`` per Option A from #000065.
|
||||
"""
|
||||
if mode not in ("just-enough", "full"):
|
||||
raise ValueError(f"unknown sync mode: {mode!r}")
|
||||
if not targets:
|
||||
raise ValueError("at least one target connection required")
|
||||
if M != len(targets):
|
||||
raise ValueError(f"M={M} but len(targets)={len(targets)} — must match")
|
||||
|
||||
meta_result = _pull_pack_inner_routed(
|
||||
targets, backend, metadata_pack_hash, M=M, kind="metadata"
|
||||
)
|
||||
chunk_results = []
|
||||
if mode == "full":
|
||||
for ch_hash in meta_result.get("chunk_pack_hashes_referenced", []):
|
||||
chunk_results.append(
|
||||
_pull_pack_inner_routed(
|
||||
targets, backend, ch_hash, M=M, kind="chunks"
|
||||
)
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "hydrated",
|
||||
"mode": mode,
|
||||
"metadata_pack_hash": metadata_pack_hash,
|
||||
"metadata_pull": meta_result,
|
||||
"chunk_pulls": chunk_results,
|
||||
"chunk_packs_pulled": len(chunk_results),
|
||||
"M": M,
|
||||
}
|
||||
|
||||
|
||||
def _pull_pack_inner_routed(
|
||||
targets: list[sqlite3.Connection],
|
||||
backend: "ObjectStoreBackend",
|
||||
pack_hash: str,
|
||||
*,
|
||||
M: int,
|
||||
kind: str,
|
||||
) -> dict:
|
||||
"""M-aware variant of :func:`_pull_pack_inner`. Routes metadata rows
|
||||
via :func:`restore_shard_metadata_routed`; fills chunk bodies on the
|
||||
target shard whose ``document_root`` route the chunk's row landed on.
|
||||
"""
|
||||
import io
|
||||
import shutil
|
||||
import tarfile
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import zstandard
|
||||
|
||||
from arborist.cold_object import hash_file_leaf, parse_manifest
|
||||
from arborist.cold_pack_metadata import restore_shard_metadata_routed
|
||||
from arborist.merkle import hash_leaf
|
||||
|
||||
backend_id = backend.identity.to_audit_body()
|
||||
body = backend.get_pack(pack_hash, kind=kind)
|
||||
|
||||
dctx = zstandard.ZstdDecompressor()
|
||||
raw_tar = dctx.stream_reader(io.BytesIO(body)).read()
|
||||
|
||||
tables_dir = Path(tempfile.mkdtemp(prefix="arborist-restore-routed-"))
|
||||
chunks_to_restore: list[tuple[str, bytes]] = []
|
||||
table_files_seen: list[str] = []
|
||||
manifest_bytes: bytes | None = None
|
||||
try:
|
||||
with tarfile.open(fileobj=io.BytesIO(raw_tar), mode="r") as tar:
|
||||
for member in tar:
|
||||
if member.name in ("manifest.ndjson", "manifest.jsonl"):
|
||||
f = tar.extractfile(member)
|
||||
if f is not None:
|
||||
manifest_bytes = f.read()
|
||||
continue
|
||||
if member.name.startswith("tables/"):
|
||||
f = tar.extractfile(member)
|
||||
if f is None:
|
||||
continue
|
||||
dest = tables_dir / Path(member.name).name
|
||||
with open(dest, "wb") as out:
|
||||
out.write(f.read())
|
||||
table_files_seen.append(member.name)
|
||||
elif member.name.startswith("blobs/"):
|
||||
tail = member.name[len("blobs/"):]
|
||||
if "/" not in tail:
|
||||
continue
|
||||
prefix, rest = tail.split("/", 1)
|
||||
if len(prefix) != 2 or len(rest) != 62:
|
||||
continue
|
||||
expected_hash = prefix + rest
|
||||
f = tar.extractfile(member)
|
||||
if f is None:
|
||||
continue
|
||||
chunk_body = f.read()
|
||||
actual_hash = hash_leaf(chunk_body).hex()
|
||||
if actual_hash != expected_hash:
|
||||
raise ValueError(
|
||||
f"pack chunk hash mismatch: declared "
|
||||
f"{expected_hash}, actual {actual_hash}"
|
||||
)
|
||||
chunks_to_restore.append((expected_hash, chunk_body))
|
||||
|
||||
# Verify table content hashes (v2+ packs).
|
||||
tampered_tables: list[dict] = []
|
||||
if manifest_bytes is not None:
|
||||
parsed = parse_manifest(manifest_bytes)
|
||||
table_ref_by_name = {t.member_name: t for t in parsed.tables}
|
||||
for member_name in table_files_seen:
|
||||
ref = table_ref_by_name.get(member_name)
|
||||
if ref is None:
|
||||
continue
|
||||
local_path = tables_dir / Path(member_name).name
|
||||
actual = hash_file_leaf(local_path)
|
||||
if actual != ref.content_hash:
|
||||
tampered_tables.append({
|
||||
"member": member_name,
|
||||
"expected_hash": ref.content_hash,
|
||||
"actual_hash": actual,
|
||||
})
|
||||
if tampered_tables:
|
||||
raise ValueError(
|
||||
f"pack {pack_hash[:12]} table-hash mismatch: "
|
||||
f"{tampered_tables}"
|
||||
)
|
||||
|
||||
# Phase 1: metadata tables → routed across M targets.
|
||||
tables_restored: dict[str, int] = {}
|
||||
if table_files_seen:
|
||||
tables_restored = restore_shard_metadata_routed(
|
||||
targets, M, tables_dir
|
||||
)
|
||||
finally:
|
||||
shutil.rmtree(tables_dir, ignore_errors=True)
|
||||
|
||||
# Phase 2: fill chunk content. Each chunk's metadata row landed on
|
||||
# ONE target (the one its document_root hashes to). Iterate targets;
|
||||
# only the owning shard finds a row for any given leaf_hash. The
|
||||
# idx_chunks_leaf index makes the per-target lookup ~O(log N).
|
||||
restored = 0
|
||||
skipped_unknown = 0
|
||||
skipped_already_hot = 0
|
||||
for leaf_hash, chunk_body in chunks_to_restore:
|
||||
landed = False
|
||||
for target in targets:
|
||||
rows = target.execute(
|
||||
"SELECT chunk_id, tier FROM chunks WHERE leaf_hash = ?",
|
||||
(leaf_hash,),
|
||||
).fetchall()
|
||||
if not rows:
|
||||
continue
|
||||
landed = True
|
||||
text = chunk_body.decode("utf-8")
|
||||
for row in rows:
|
||||
chunk_id = row["chunk_id"] if isinstance(row, sqlite3.Row) else row[0]
|
||||
tier = row["tier"] if isinstance(row, sqlite3.Row) else row[1]
|
||||
if tier == "hot" and chunk_id is not None:
|
||||
existing = target.execute(
|
||||
"SELECT content IS NOT NULL FROM chunks WHERE chunk_id = ?",
|
||||
(chunk_id,),
|
||||
).fetchone()
|
||||
if existing and existing[0]:
|
||||
skipped_already_hot += 1
|
||||
continue
|
||||
with transaction(target):
|
||||
target.execute(
|
||||
"UPDATE chunks SET content = ?, tier = 'hot' WHERE chunk_id = ?",
|
||||
(pack_chunk(text), chunk_id),
|
||||
)
|
||||
target.execute(
|
||||
"INSERT OR REPLACE INTO chunks_fts (rowid, content) "
|
||||
"VALUES (?, ?)",
|
||||
(chunk_id, text),
|
||||
)
|
||||
restored += 1
|
||||
if not landed:
|
||||
skipped_unknown += 1
|
||||
|
||||
chunk_pack_hashes_referenced: tuple[str, ...] = ()
|
||||
if manifest_bytes is not None:
|
||||
chunk_pack_hashes_referenced = parse_manifest(manifest_bytes).chunk_pack_hashes
|
||||
|
||||
# Audit event lands on canonical target 0 (consolidated chain).
|
||||
with transaction(targets[0]):
|
||||
append_audit(
|
||||
targets[0],
|
||||
event_type="cold_pack_pulled",
|
||||
body={
|
||||
"pack_hash": pack_hash,
|
||||
"pack_kind": kind,
|
||||
"M": M,
|
||||
"chunks_restored": restored,
|
||||
"chunks_skipped_already_hot": skipped_already_hot,
|
||||
"chunks_skipped_unknown_locally": skipped_unknown,
|
||||
"tables_restored": tables_restored,
|
||||
"backend": backend_id,
|
||||
},
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "pulled",
|
||||
"pack_hash": pack_hash,
|
||||
"pack_kind": kind,
|
||||
"M": M,
|
||||
"chunk_pack_hashes_referenced": list(chunk_pack_hashes_referenced),
|
||||
"tables_restored": tables_restored,
|
||||
"chunks_restored": restored,
|
||||
"chunks_skipped_already_hot": skipped_already_hot,
|
||||
"chunks_skipped_unknown_locally": skipped_unknown,
|
||||
}
|
||||
|
||||
|
||||
def _pull_pack_inner(
|
||||
conn: sqlite3.Connection,
|
||||
backend: "ObjectStoreBackend",
|
||||
|
|
|
|||
214
tests/test_cold_unpack_routed.py
Normal file
214
tests/test_cold_unpack_routed.py
Normal file
|
|
@ -0,0 +1,214 @@
|
|||
"""Tests for #000067 M-aware cold-pack hydration.
|
||||
|
||||
Build a tiny corpus → pack into a MemoryBackend → hydrate into M=4
|
||||
fresh target shards via the routed path. Assert every doc lands on
|
||||
``shard_for_document(document_root, M)`` and corpus-wide tables
|
||||
consolidate to target 0.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
|
||||
import pytest
|
||||
|
||||
from arborist.cold_object import MemoryBackend
|
||||
from arborist.document import Document, shard_for_document
|
||||
from arborist.evict import (
|
||||
hydrate_from_metadata_pack_routed,
|
||||
push_pack,
|
||||
)
|
||||
from arborist.ingest import ingest_source
|
||||
from arborist.source import Source
|
||||
from arborist.store import (
|
||||
SCHEMA_SQL,
|
||||
connect,
|
||||
get_corpus_shard_count,
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
content = (
|
||||
f"This is doc {uri} seed {seed}. " * 10
|
||||
+ f"Body text {seed} with tokens enough for chunking. " * 30
|
||||
)
|
||||
return Document(
|
||||
uri=uri, content=content, source_type="html",
|
||||
title=uri.split("/")[-1],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def producer_shard(tmp_path: Path) -> Path:
|
||||
"""A small producer corpus packed into one shard."""
|
||||
db = tmp_path / "producer.db"
|
||||
conn = connect(db)
|
||||
try:
|
||||
docs = [
|
||||
_make_doc(f"https://prod/doc-{i}", seed=f"s0d{i}")
|
||||
for i in range(20)
|
||||
]
|
||||
ingest_source(conn, _FakeSource(docs))
|
||||
finally:
|
||||
conn.close()
|
||||
return db
|
||||
|
||||
|
||||
def _make_target_shards(target_dir: Path, M: int) -> list[sqlite3.Connection]:
|
||||
targets: list[sqlite3.Connection] = []
|
||||
for i in range(M):
|
||||
path = target_dir / f"{i:03d}.db"
|
||||
c = sqlite3.connect(str(path))
|
||||
c.row_factory = sqlite3.Row
|
||||
c.executescript(SCHEMA_SQL)
|
||||
c.execute("PRAGMA foreign_keys = OFF")
|
||||
targets.append(c)
|
||||
return targets
|
||||
|
||||
|
||||
class TestRoutedHydrate:
|
||||
def test_pack_then_hydrate_routed(
|
||||
self, producer_shard: Path, tmp_path: Path
|
||||
):
|
||||
# Pack the producer shard into a MemoryBackend.
|
||||
backend = MemoryBackend()
|
||||
with sqlite3.connect(str(producer_shard)) as src_conn:
|
||||
src_conn.row_factory = sqlite3.Row
|
||||
push_result = push_pack(
|
||||
src_conn, backend,
|
||||
document_root=None,
|
||||
max_chunks=None,
|
||||
max_pack_bytes=10 * 1024 ** 2,
|
||||
allow_license_class="unknown",
|
||||
)
|
||||
# Identify the metadata pack hash.
|
||||
metadata_pack_hash = None
|
||||
for p in push_result.get("packs", []):
|
||||
if p.get("kind") == "metadata":
|
||||
metadata_pack_hash = p["pack_hash"]
|
||||
break
|
||||
if metadata_pack_hash is None:
|
||||
metadata_pack_hash = push_result["packs"][0]["pack_hash"]
|
||||
|
||||
# Hydrate into a fresh 4-shard target.
|
||||
target_dir = tmp_path / "tgt"
|
||||
target_dir.mkdir()
|
||||
M = 4
|
||||
targets = _make_target_shards(target_dir, M)
|
||||
try:
|
||||
hydrate_result = hydrate_from_metadata_pack_routed(
|
||||
targets, backend, metadata_pack_hash, M=M, mode="full",
|
||||
)
|
||||
finally:
|
||||
for t in targets:
|
||||
t.close()
|
||||
assert hydrate_result["status"] == "hydrated"
|
||||
assert hydrate_result["M"] == M
|
||||
|
||||
# Verify: every doc on the producer must land on its hash-routed
|
||||
# target shard.
|
||||
prod_roots: dict[str, int] = {}
|
||||
c = sqlite3.connect(f"file:{producer_shard}?mode=ro", uri=True)
|
||||
try:
|
||||
for (root,) in c.execute("SELECT document_root FROM documents"):
|
||||
prod_roots[root] = shard_for_document(root, M)
|
||||
finally:
|
||||
c.close()
|
||||
|
||||
for tidx in range(M):
|
||||
tdb = target_dir / 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, expected_idx in prod_roots.items():
|
||||
if expected_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 {expected_idx}"
|
||||
)
|
||||
|
||||
def test_corpus_shard_count_set_on_routed_hydrate(
|
||||
self, producer_shard: Path, tmp_path: Path
|
||||
):
|
||||
"""CLI sets corpus_shard_count meta on every target; here the
|
||||
test directly calls the helper to confirm the meta plumbing
|
||||
works from the routed entry point."""
|
||||
backend = MemoryBackend()
|
||||
with sqlite3.connect(str(producer_shard)) as src_conn:
|
||||
src_conn.row_factory = sqlite3.Row
|
||||
push_result = push_pack(
|
||||
src_conn, backend,
|
||||
document_root=None,
|
||||
max_chunks=None,
|
||||
max_pack_bytes=10 * 1024 ** 2,
|
||||
allow_license_class="unknown",
|
||||
)
|
||||
metadata_pack_hash = next(
|
||||
p["pack_hash"] for p in push_result["packs"]
|
||||
if p.get("kind") == "metadata"
|
||||
)
|
||||
|
||||
target_dir = tmp_path / "tgt"
|
||||
target_dir.mkdir()
|
||||
M = 4
|
||||
targets = _make_target_shards(target_dir, M)
|
||||
try:
|
||||
hydrate_from_metadata_pack_routed(
|
||||
targets, backend, metadata_pack_hash, M=M, mode="full",
|
||||
)
|
||||
# Simulate the CLI's post-hydrate meta stamp.
|
||||
from arborist.store import set_corpus_shard_count
|
||||
for t in targets:
|
||||
with t:
|
||||
set_corpus_shard_count(t, M)
|
||||
finally:
|
||||
for t in targets:
|
||||
t.close()
|
||||
for i in range(M):
|
||||
tdb = target_dir / f"{i:03d}.db"
|
||||
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
|
||||
try:
|
||||
assert get_corpus_shard_count(c) == M
|
||||
finally:
|
||||
c.close()
|
||||
|
||||
|
||||
class TestRoutedHydrateValidation:
|
||||
def test_M_mismatch_rejected(self, tmp_path: Path):
|
||||
target_dir = tmp_path / "tgt"
|
||||
target_dir.mkdir()
|
||||
targets = _make_target_shards(target_dir, 4)
|
||||
backend = MemoryBackend()
|
||||
try:
|
||||
with pytest.raises(ValueError, match="must match"):
|
||||
hydrate_from_metadata_pack_routed(
|
||||
targets, backend, "deadbeef" * 8, M=3, mode="just-enough",
|
||||
)
|
||||
finally:
|
||||
for t in targets:
|
||||
t.close()
|
||||
|
||||
def test_empty_targets_rejected(self, tmp_path: Path):
|
||||
backend = MemoryBackend()
|
||||
with pytest.raises(ValueError, match="at least one target"):
|
||||
hydrate_from_metadata_pack_routed(
|
||||
[], backend, "deadbeef" * 8, M=0, mode="just-enough",
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue