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.
103 lines
3.6 KiB
Python
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()
|