arborist/tests/test_resume.py
russell@unturf.com db3e3c00ab
resumable ingest + per-shard audit chain integrity
Resume (rsync-style)
  Each shard DB gains a `meta` table. After every successful batch
  flush, ingest_source persists `source_high_water:<source_type>` ->
  the largest row id seen (cur_id or old_id). On --resume, the source
  reads it back and skips rows whose id is <= the mark, so an ingest
  killed at any point can be re-run cheaply: already-cached docs are
  fast-forwarded past without re-hashing or DB writes.

  WikipediaSqlDump now exposes `start_id` (skip threshold) and
  `last_id` (running max). cur_id is surfaced in Document.extra
  alongside old_id so both tables behave the same.

  CLI: aborist ingest --resume

  Demo on shared DB (cur):
    Round 1: --limit 3000 --resume   high_water = 5714
    Round 2: --limit 5000 --resume   skips 1..5714, picks up at 5715
    Round 3: --limit 100  --resume   skips 1..15362, ingests 100 more
  Always idempotent on re-run; no dups, no missing docs.

Per-shard audit chain integrity for sharded analyze
  Cross-shard analyze previously reported nonsense breaks counts —
  each shard owns its own audit chain (genesis -> ... -> latest), and
  the UNION view interleaves them so cross-shard transitions look
  like break events. The fix: when --shards-dir is set, open each
  shard's DB directly and run _check_audit_chain on it, then
  aggregate.

  Output now reads:
    "audit_chain": {
      "events": <total>,
      "breaks": 0,
      "shards": [{"shard": "000.db", "events": N, "breaks": 0}, ...]
    }

53 tests passing. Tests cover: high-water write, skip-on-resume,
idempotency across two resume runs, kill-and-resume continuity.
2026-04-27 11:29:27 -04:00

103 lines
3.6 KiB
Python

"""Resumable ingest: high-water mark in meta lets a stopped ingest rsync forward."""
from __future__ import annotations
from typing import Iterator
from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.source import Source
from aborist.store import connect, get_meta, set_meta
class IndexedSource(Source):
"""Test source that exposes an `id` per doc and a start_id filter.
Imitates the WikipediaSqlDump resume contract.
"""
source_type = "indexed_test"
def __init__(self, n_docs: int = 10):
self.n_docs = n_docs
self.start_id = 0
self.last_id = 0
def iter_documents(self) -> Iterator[Document]:
for i in range(1, self.n_docs + 1):
if i <= self.start_id:
continue
self.last_id = i
yield Document(
uri=f"test://doc/{i}",
content=f"document number {i} content " * 10,
source_type=self.source_type,
title=f"Doc {i}",
extra={"indexed_test_id": str(i)},
)
def test_resume_writes_high_water(tmp_path):
db = tmp_path / "rh.db"
conn = connect(db)
try:
src = IndexedSource(n_docs=10)
ingest_source(conn, src, batch_size=5, resume=True)
# high-water now equals the last id seen
assert get_meta(conn, "source_high_water:indexed_test") == "10"
finally:
conn.close()
def test_resume_skips_processed(tmp_path):
db = tmp_path / "skip.db"
conn = connect(db)
try:
# Pre-seed a high-water as if a prior run got through doc 4.
with __import__("aborist").store.transaction(conn):
set_meta(conn, "source_high_water:indexed_test", "4")
src = IndexedSource(n_docs=10)
result = ingest_source(conn, src, batch_size=5, resume=True)
# Source should have skipped docs 1..4 entirely (start_id=4, then >4 emitted).
# Note: `seen` counts what the source yielded post-skip.
assert result.seen == 6
assert result.inserted == 6
# high-water now updated to 10
assert get_meta(conn, "source_high_water:indexed_test") == "10"
finally:
conn.close()
def test_resume_idempotent(tmp_path):
"""Running ingest twice with resume yields no new docs the second time."""
db = tmp_path / "idem.db"
conn = connect(db)
try:
ingest_source(conn, IndexedSource(n_docs=10), batch_size=3, resume=True)
n1 = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
ingest_source(conn, IndexedSource(n_docs=10), batch_size=3, resume=True)
n2 = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
assert n1 == n2 == 10
finally:
conn.close()
def test_resume_continues_after_kill(tmp_path):
"""Simulate a kill mid-batch: high-water is at the last successfully
flushed batch, restart picks up from there without dups or loss."""
db = tmp_path / "kill.db"
conn = connect(db)
try:
# First run hits 7 of 10, then "killed" — but with batch_size=5 the
# high-water should be 5 (one full batch flushed).
ingest_source(conn, IndexedSource(n_docs=5), batch_size=5, resume=True)
assert get_meta(conn, "source_high_water:indexed_test") == "5"
# Resume with the full source — should pick up at 6..10.
result = ingest_source(conn, IndexedSource(n_docs=10), batch_size=5, resume=True)
assert result.seen == 5
assert result.inserted == 5
n_total = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
assert n_total == 10
finally:
conn.close()