Merge pull request #452 from ggozad/feat/run-batch-dry-run
Add dry-run manifests for ingestor's run-batch
This commit is contained in:
commit
0be4f7e1bd
13 changed files with 1304 additions and 34 deletions
|
|
@ -1,6 +1,11 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- `haiku-ingester run-batch --dry-run` writes a YAML manifest of planned upserts/deletes without mutating queue jobs or `sync_state`; `run-batch --manifest <path>` replays that frozen changeset without another discovery sweep.
|
||||
- `haiku-ingester run-batch` and `run-batch --manifest` show an interactive progress bar with ETA while draining queued jobs.
|
||||
|
||||
### Fixed
|
||||
|
||||
- CPU-bound ingest steps now run off the event loop: Docling document serialization, docling-serve zip parsing, split-PDF concatenation, fetched-body temp writes, and filesystem read/hash work no longer stall concurrent ingester workers on large image-bearing documents.
|
||||
|
|
|
|||
|
|
@ -524,6 +524,31 @@ haiku-ingester run-batch
|
|||
haiku-ingester run-batch --db rag.lancedb
|
||||
```
|
||||
|
||||
To review a batch before it mutates the document store, use `--dry-run`.
|
||||
Dry-run performs the same discovery checks but writes no queue jobs and does
|
||||
not update `sync_state`. It writes a YAML manifest named
|
||||
`manifest-<datestamp>.yaml` by default:
|
||||
|
||||
```bash
|
||||
haiku-ingester run-batch --dry-run
|
||||
haiku-ingester run-batch --dry-run --output manifest-20260622.yaml
|
||||
```
|
||||
|
||||
The manifest records the `upsert` and `delete` changes discovered for each
|
||||
source. Replay it later to ingest exactly that changeset, without another
|
||||
discovery sweep:
|
||||
|
||||
```bash
|
||||
haiku-ingester run-batch --manifest manifest-20260622.yaml
|
||||
```
|
||||
|
||||
Manifest replay rejects sources with queued or claimed work, preserving the
|
||||
one-active-changeset-per-source pattern. Revisioned upserts are checked
|
||||
against the current upstream revision before fetch; if the resource changed
|
||||
after dry-run, that job dead-letters and the newer version waits for the next
|
||||
dry-run. Sources that provide no revision can freeze URI discovery but cannot
|
||||
prove byte identity at replay time.
|
||||
|
||||
Orphan deletion compares each source against `sync_state` in the queue DB,
|
||||
so persist `ingester.db` between runs for deletions to be detected. It exits
|
||||
non-zero if any job dead-letters or a source's discovery sweep does not
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import logging
|
||||
import signal
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
|
@ -9,9 +10,11 @@ from typing import TYPE_CHECKING
|
|||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.config import AppConfig
|
||||
from haiku.rag.ingester.batch import BatchChange, BatchDryRunReport, BatchManifest
|
||||
from haiku.rag.ingester.metadata import build_providers, load_metadata_providers
|
||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||
from haiku.rag.ingester.queue.migrations import open_queue
|
||||
from haiku.rag.ingester.queue.models import Job, JobStatus
|
||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||
from haiku.rag.ingester.workers.pool import WorkerPool
|
||||
from haiku.rag.ingester.workers.retry import RetryPolicy
|
||||
|
|
@ -21,6 +24,8 @@ if TYPE_CHECKING:
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MANIFEST_EXTRA_KEY = "_manifest"
|
||||
|
||||
|
||||
def _api_access_log_enabled() -> bool:
|
||||
"""Per-request access logging only when the haiku.rag logger is at DEBUG.
|
||||
|
|
@ -38,6 +43,28 @@ class BatchReport(BaseModel):
|
|||
failed_sweeps: list[str] = []
|
||||
|
||||
|
||||
class BatchProgress(BaseModel):
|
||||
"""Snapshot emitted while a one-shot batch drains queued work."""
|
||||
|
||||
total: int = 0
|
||||
completed: int = 0
|
||||
succeeded: int = 0
|
||||
dead: int = 0
|
||||
queued: int = 0
|
||||
claimed: int = 0
|
||||
|
||||
|
||||
BatchProgressCallback = Callable[[BatchProgress], None]
|
||||
|
||||
|
||||
def _manifest_change_key(change: BatchChange) -> tuple[str, str, str, str | None]:
|
||||
return (change.source_id, change.uri, change.op.value, change.revision)
|
||||
|
||||
|
||||
def _manifest_job_key(job: Job) -> tuple[str, str, str, str | None]:
|
||||
return (job.source_id, job.uri, job.op.value, job.revision)
|
||||
|
||||
|
||||
class IngesterApp:
|
||||
"""Top-level lifecycle for the production ingester.
|
||||
|
||||
|
|
@ -129,6 +156,35 @@ class IngesterApp:
|
|||
await self._engine.dispose()
|
||||
self._engine = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _discovery_resources(self):
|
||||
"""Open only the queue and source pollers needed for discovery.
|
||||
Dry-runs must not create/open the LanceDB document store or worker
|
||||
pool because they are upstream checks only."""
|
||||
from haiku.rag.converters import get_converter
|
||||
|
||||
ingester_cfg = self._config.ingester
|
||||
self._engine = await open_queue(ingester_cfg.queue)
|
||||
try:
|
||||
self._jobs = JobRepo(self._engine)
|
||||
self._sync = SyncStateRepo(self._engine)
|
||||
supported_extensions = get_converter(self._config).supported_extensions
|
||||
self._pollers = PollerManager(
|
||||
configs=ingester_cfg.sources,
|
||||
job_repo=self._jobs,
|
||||
sync_repo=self._sync,
|
||||
supported_extensions=supported_extensions,
|
||||
default_max_attempts=ingester_cfg.workers.retry.max_attempts,
|
||||
)
|
||||
yield
|
||||
finally:
|
||||
if self._pollers is not None:
|
||||
await self._pollers.close_sources()
|
||||
self._pollers = None
|
||||
if self._engine is not None:
|
||||
await self._engine.dispose()
|
||||
self._engine = None
|
||||
|
||||
async def _stop_pool(self) -> None:
|
||||
"""Stop the worker pool, honouring the shutdown grace, then drain any
|
||||
cancel-cleanup release tasks before the queue connection closes."""
|
||||
|
|
@ -149,6 +205,51 @@ class IngesterApp:
|
|||
if landed:
|
||||
logger.info("Drained %d cancel-cleanup release(s) before close", landed)
|
||||
|
||||
async def _drain_batch(
|
||||
self,
|
||||
started_at: datetime,
|
||||
*,
|
||||
progress_callback: BatchProgressCallback | None = None,
|
||||
) -> BatchReport:
|
||||
assert self._pool is not None and self._jobs is not None
|
||||
total = 0
|
||||
while True:
|
||||
counts = await self._jobs.batch_progress_counts_since(started_at)
|
||||
queued = counts.get("queued", 0)
|
||||
claimed = counts.get("claimed", 0)
|
||||
outstanding = queued + claimed
|
||||
succeeded = counts.get("succeeded", 0)
|
||||
dead = counts.get("dead", 0)
|
||||
completed_count = succeeded + dead
|
||||
total = max(total, outstanding + completed_count)
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
BatchProgress(
|
||||
total=total,
|
||||
completed=min(completed_count, total),
|
||||
succeeded=succeeded,
|
||||
dead=dead,
|
||||
queued=queued,
|
||||
claimed=claimed,
|
||||
)
|
||||
)
|
||||
if not outstanding:
|
||||
break
|
||||
if self._pool.live_workers == 0:
|
||||
logger.error(
|
||||
"All workers have died with %d outstanding job(s) "
|
||||
"— aborting batch; stranded jobs will be reaped "
|
||||
"on next start",
|
||||
outstanding,
|
||||
)
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
completed = await self._jobs.counts_by_status_since(started_at)
|
||||
return BatchReport(
|
||||
succeeded=completed.get("succeeded", 0),
|
||||
dead=completed.get("dead", 0),
|
||||
)
|
||||
|
||||
async def serve(self, *, api: bool = True) -> None:
|
||||
"""Run pollers + workers (and the HTTP API when enabled) until a
|
||||
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
|
||||
|
|
@ -164,7 +265,7 @@ class IngesterApp:
|
|||
uses_docling_serve = (
|
||||
proc.converter == "docling-serve" or proc.chunker == "docling-serve"
|
||||
)
|
||||
if uses_docling_serve:
|
||||
if uses_docling_serve: # pragma: no cover
|
||||
logger.info(
|
||||
"Ingester running: %d worker(s), %d source(s), "
|
||||
"%d docling-serve instance(s)",
|
||||
|
|
@ -202,7 +303,9 @@ class IngesterApp:
|
|||
await self._stop_pool()
|
||||
await self._pollers.close_sources()
|
||||
|
||||
async def run_batch(self) -> BatchReport:
|
||||
async def run_batch(
|
||||
self, *, progress_callback: BatchProgressCallback | None = None
|
||||
) -> BatchReport:
|
||||
"""Run one discover() sweep across every configured source, drain the
|
||||
queue to completion, then stop. Unlike `serve`, the periodic poller
|
||||
loops never start — discovery is driven explicitly, so the run is
|
||||
|
|
@ -220,25 +323,128 @@ class IngesterApp:
|
|||
await self._pool.start()
|
||||
try:
|
||||
failed_sweeps = await self._pollers.sweep_all()
|
||||
while True:
|
||||
counts = await self._jobs.counts_by_status()
|
||||
if not counts.get("queued") and not counts.get("claimed"):
|
||||
break
|
||||
if self._pool.live_workers == 0:
|
||||
outstanding = counts.get("queued", 0) + counts.get("claimed", 0)
|
||||
logger.error(
|
||||
"All workers have died with %d outstanding job(s) "
|
||||
"— aborting batch; stranded jobs will be reaped "
|
||||
"on next start",
|
||||
outstanding,
|
||||
)
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
completed = await self._jobs.counts_by_status_since(started_at)
|
||||
return BatchReport(
|
||||
succeeded=completed.get("succeeded", 0),
|
||||
dead=completed.get("dead", 0),
|
||||
failed_sweeps=failed_sweeps,
|
||||
report = await self._drain_batch(
|
||||
started_at, progress_callback=progress_callback
|
||||
)
|
||||
report.failed_sweeps = failed_sweeps
|
||||
return report
|
||||
finally:
|
||||
await self._stop_pool()
|
||||
await self._pollers.close_sources()
|
||||
|
||||
async def run_batch_dry_run(self) -> BatchDryRunReport:
|
||||
"""Run one discover() sweep across every configured source and return
|
||||
the jobs that would be enqueued, without mutating jobs or sync_state."""
|
||||
async with self._discovery_resources():
|
||||
assert self._pollers is not None
|
||||
manifest, failed_sweeps = await self._pollers.dry_run_manifest()
|
||||
return BatchDryRunReport(manifest=manifest, failed_sweeps=failed_sweeps)
|
||||
|
||||
async def run_batch_from_manifest(
|
||||
self,
|
||||
manifest: BatchManifest,
|
||||
*,
|
||||
progress_callback: BatchProgressCallback | None = None,
|
||||
) -> BatchReport:
|
||||
"""Enqueue and drain a dry-run manifest without running a fresh
|
||||
discovery sweep."""
|
||||
if manifest.version != 1: # pragma: no cover
|
||||
raise ValueError(f"Unsupported manifest version: {manifest.version}")
|
||||
async with self._resources():
|
||||
assert (
|
||||
self._pollers is not None
|
||||
and self._pool is not None
|
||||
and self._jobs is not None
|
||||
)
|
||||
configured = {source.source_id for source in self._pollers.sources}
|
||||
manifest_sources = {change.source_id for change in manifest.changes}
|
||||
missing = sorted(manifest_sources - configured)
|
||||
if missing: # pragma: no cover
|
||||
await self._pollers.close_sources()
|
||||
raise ValueError(
|
||||
"Manifest references unconfigured source(s): " + ", ".join(missing)
|
||||
)
|
||||
|
||||
seen: set[tuple[str, str]] = set()
|
||||
duplicates: set[tuple[str, str]] = set()
|
||||
for change in manifest.changes:
|
||||
key = (change.source_id, change.uri)
|
||||
if key in seen:
|
||||
duplicates.add(key)
|
||||
seen.add(key)
|
||||
if duplicates:
|
||||
await self._pollers.close_sources()
|
||||
rendered = ", ".join(
|
||||
f"{source_id}:{uri}" for source_id, uri in duplicates
|
||||
)
|
||||
raise ValueError(f"Manifest contains duplicate change(s): {rendered}")
|
||||
|
||||
manifest_key = manifest.generated_at.isoformat()
|
||||
manifest_change_keys = {
|
||||
_manifest_change_key(change) for change in manifest.changes
|
||||
}
|
||||
live_jobs = [
|
||||
*await self._jobs.list_jobs(status=JobStatus.QUEUED, limit=10_000),
|
||||
*await self._jobs.list_jobs(status=JobStatus.CLAIMED, limit=10_000),
|
||||
]
|
||||
stale_jobs: list[Job] = []
|
||||
live_manifest_keys: set[tuple[str, str, str, str | None]] = set()
|
||||
for job in live_jobs:
|
||||
extra = job.extra or {}
|
||||
job_manifest = extra.get(_MANIFEST_EXTRA_KEY) or {}
|
||||
key = _manifest_job_key(job)
|
||||
if (
|
||||
job_manifest.get("generated_at") != manifest_key
|
||||
or key not in manifest_change_keys
|
||||
):
|
||||
stale_jobs.append(job)
|
||||
continue
|
||||
live_manifest_keys.add(key)
|
||||
if stale_jobs:
|
||||
await self._pollers.close_sources()
|
||||
raise ValueError(
|
||||
"Cannot replay manifest while the queue has non-manifest "
|
||||
f"pending work: {len(stale_jobs)} queued/claimed job(s)"
|
||||
)
|
||||
|
||||
default_max_attempts = self._config.ingester.workers.retry.max_attempts
|
||||
max_attempts_by_source = {
|
||||
poller.source_id: (
|
||||
poller.config.retry.max_attempts
|
||||
if poller.config.retry is not None
|
||||
else default_max_attempts
|
||||
)
|
||||
for poller in self._pollers.pollers
|
||||
}
|
||||
for change in manifest.changes:
|
||||
if _manifest_change_key(change) in live_manifest_keys:
|
||||
continue
|
||||
job = await self._jobs.enqueue(
|
||||
change.source_id,
|
||||
change.uri,
|
||||
op=change.op,
|
||||
revision=change.revision,
|
||||
max_attempts=max_attempts_by_source[change.source_id],
|
||||
extra={
|
||||
_MANIFEST_EXTRA_KEY: {
|
||||
"version": manifest.version,
|
||||
"generated_at": manifest_key,
|
||||
"discovered_at": change.discovered_at.isoformat(),
|
||||
}
|
||||
},
|
||||
)
|
||||
if job is None: # pragma: no cover
|
||||
await self._pollers.close_sources()
|
||||
raise ValueError(
|
||||
"Cannot replay manifest because a live job already exists "
|
||||
f"for {change.source_id}:{change.uri}"
|
||||
)
|
||||
|
||||
started_at = datetime.now(UTC)
|
||||
await self._pool.start()
|
||||
try:
|
||||
return await self._drain_batch(
|
||||
started_at, progress_callback=progress_callback
|
||||
)
|
||||
finally:
|
||||
await self._stop_pool()
|
||||
|
|
|
|||
33
haiku_rag_slim/haiku/rag/ingester/batch.py
Normal file
33
haiku_rag_slim/haiku/rag/ingester/batch.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.ingester.queue.models import JobOp
|
||||
|
||||
|
||||
class BatchChange(BaseModel):
|
||||
op: JobOp
|
||||
source_id: str
|
||||
uri: str
|
||||
revision: str | None = None
|
||||
discovered_at: datetime
|
||||
|
||||
|
||||
class BatchSourceSummary(BaseModel):
|
||||
source_id: str
|
||||
upsert_count: int = 0
|
||||
delete_count: int = 0
|
||||
unchanged_count: int = 0
|
||||
ignored_delete_count: int = 0
|
||||
|
||||
|
||||
class BatchManifest(BaseModel):
|
||||
version: int = 1
|
||||
generated_at: datetime
|
||||
sources: list[BatchSourceSummary] = []
|
||||
changes: list[BatchChange] = []
|
||||
|
||||
|
||||
class BatchDryRunReport(BaseModel):
|
||||
manifest: BatchManifest
|
||||
failed_sweeps: list[str] = []
|
||||
|
|
@ -1,9 +1,21 @@
|
|||
import asyncio
|
||||
import sys
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
import yaml
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
from rich.console import Console
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
Progress,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
TimeRemainingColumn,
|
||||
)
|
||||
from sqlalchemy import make_url
|
||||
|
||||
load_dotenv(find_dotenv(usecwd=True))
|
||||
|
|
@ -16,7 +28,12 @@ from haiku.rag.config import ( # noqa: E402
|
|||
load_yaml_config,
|
||||
set_config,
|
||||
)
|
||||
from haiku.rag.ingester.app import IngesterApp # noqa: E402
|
||||
from haiku.rag.ingester.app import ( # noqa: E402
|
||||
BatchProgress,
|
||||
BatchProgressCallback,
|
||||
IngesterApp,
|
||||
)
|
||||
from haiku.rag.ingester.batch import BatchManifest # noqa: E402
|
||||
from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402
|
||||
from haiku.rag.logging import configure_cli_logging # noqa: E402
|
||||
from haiku.rag.store.exceptions import ( # noqa: E402
|
||||
|
|
@ -137,6 +154,64 @@ def _resolve_db_path(config: AppConfig, override: Path | None) -> Path:
|
|||
return override or (config.storage.data_dir / "haiku.rag.lancedb")
|
||||
|
||||
|
||||
def _default_manifest_path() -> Path:
|
||||
datestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%SZ")
|
||||
return Path(f"manifest-{datestamp}.yaml")
|
||||
|
||||
|
||||
def _write_manifest(manifest: BatchManifest, path: Path) -> None:
|
||||
data = manifest.model_dump(mode="json")
|
||||
path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _batch_progress(
|
||||
description: str,
|
||||
) -> Iterator[BatchProgressCallback | None]: # pragma: no cover
|
||||
console = Console(file=sys.stdout)
|
||||
if not console.is_terminal:
|
||||
yield None
|
||||
return
|
||||
|
||||
progress = Progress(
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TextColumn("{task.completed}/{task.total}"),
|
||||
TimeRemainingColumn(),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
)
|
||||
task_id = None
|
||||
|
||||
def _update(snapshot: BatchProgress) -> None:
|
||||
nonlocal task_id
|
||||
task_description = (
|
||||
f"{description} ({snapshot.succeeded} ok, {snapshot.dead} dead)"
|
||||
)
|
||||
if task_id is None:
|
||||
task_id = progress.add_task(
|
||||
task_description,
|
||||
total=snapshot.total,
|
||||
completed=snapshot.completed,
|
||||
)
|
||||
return
|
||||
progress.update(
|
||||
task_id,
|
||||
description=task_description,
|
||||
total=snapshot.total,
|
||||
completed=snapshot.completed,
|
||||
)
|
||||
|
||||
with progress:
|
||||
yield _update
|
||||
|
||||
|
||||
def _load_manifest(path: Path) -> BatchManifest:
|
||||
data = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
return BatchManifest.model_validate(data)
|
||||
|
||||
|
||||
@_cli.command("serve")
|
||||
def serve(
|
||||
db: Path | None = typer.Option(
|
||||
|
|
@ -188,18 +263,98 @@ def run_batch(
|
|||
"--db",
|
||||
help="LanceDB path (overrides config.storage.data_dir).",
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help="Discover planned changes and write a YAML manifest without ingesting.",
|
||||
),
|
||||
output: Path | None = typer.Option(
|
||||
None,
|
||||
"--output",
|
||||
"-o",
|
||||
help="Dry-run manifest path (defaults to manifest-<datestamp>.yaml).",
|
||||
),
|
||||
manifest: Path | None = typer.Option(
|
||||
None,
|
||||
"--manifest",
|
||||
help="Replay a dry-run manifest instead of running discovery.",
|
||||
),
|
||||
) -> None:
|
||||
"""Run one discover sweep across every configured source, drain the queue,
|
||||
then exit. New and changed resources are ingested, resources that vanished
|
||||
from a source are deleted. Exits non-zero if any job dead-letters or a
|
||||
source's sweep does not complete."""
|
||||
asyncio.run(_run_batch(get_config(), db))
|
||||
if manifest is not None and dry_run:
|
||||
typer.echo("Error: --manifest cannot be combined with --dry-run")
|
||||
raise typer.Exit(2)
|
||||
if output is not None and not dry_run:
|
||||
typer.echo("Error: --output is only valid with --dry-run")
|
||||
raise typer.Exit(2)
|
||||
asyncio.run(
|
||||
_run_batch(
|
||||
get_config(),
|
||||
db,
|
||||
dry_run=dry_run,
|
||||
output=output,
|
||||
manifest_path=manifest,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_batch(app_config: AppConfig, db_path: Path | None) -> None:
|
||||
async def _run_batch(
|
||||
app_config: AppConfig,
|
||||
db_path: Path | None,
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
output: Path | None = None,
|
||||
manifest_path: Path | None = None,
|
||||
) -> None:
|
||||
db = _resolve_db_path(app_config, db_path)
|
||||
app = IngesterApp(config=app_config, db_path=db)
|
||||
report = await app.run_batch()
|
||||
if dry_run:
|
||||
report = await app.run_batch_dry_run()
|
||||
if report.failed_sweeps:
|
||||
typer.echo(
|
||||
f"Sources that failed to sweep: {', '.join(report.failed_sweeps)}"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
manifest_path = output or _default_manifest_path()
|
||||
_write_manifest(report.manifest, manifest_path)
|
||||
upserts = sum(source.upsert_count for source in report.manifest.sources)
|
||||
deletes = sum(source.delete_count for source in report.manifest.sources)
|
||||
unchanged = sum(source.unchanged_count for source in report.manifest.sources)
|
||||
typer.echo(
|
||||
"Dry run complete: "
|
||||
f"{upserts} upsert, {deletes} delete, {unchanged} unchanged "
|
||||
f"-> {manifest_path}"
|
||||
)
|
||||
return
|
||||
|
||||
if manifest_path is not None:
|
||||
try:
|
||||
manifest = _load_manifest(manifest_path)
|
||||
with _batch_progress("Replaying manifest") as progress_callback:
|
||||
if progress_callback is None:
|
||||
report = await app.run_batch_from_manifest(manifest)
|
||||
else:
|
||||
report = await app.run_batch_from_manifest(
|
||||
manifest, progress_callback=progress_callback
|
||||
)
|
||||
except ValueError as exc:
|
||||
typer.echo(f"Error: {exc}")
|
||||
raise typer.Exit(1) from exc
|
||||
typer.echo(
|
||||
f"Manifest batch complete: {report.succeeded} succeeded, {report.dead} dead"
|
||||
)
|
||||
if report.dead:
|
||||
raise typer.Exit(1)
|
||||
return
|
||||
|
||||
with _batch_progress("Running batch") as progress_callback:
|
||||
if progress_callback is None:
|
||||
report = await app.run_batch()
|
||||
else:
|
||||
report = await app.run_batch(progress_callback=progress_callback)
|
||||
typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
|
||||
if report.failed_sweeps:
|
||||
typer.echo(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import random
|
|||
from datetime import UTC, datetime
|
||||
|
||||
from haiku.rag.config import SourceConfig
|
||||
from haiku.rag.ingester.batch import BatchChange, BatchSourceSummary
|
||||
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
||||
from haiku.rag.ingester.queue.models import JobOp, SyncRow
|
||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||
|
|
@ -165,6 +166,82 @@ class BasePoller:
|
|||
)
|
||||
return False
|
||||
|
||||
async def _dry_run_once(self) -> tuple[bool, BatchSourceSummary, list[BatchChange]]:
|
||||
"""Collect what one discover() sweep would enqueue without writing
|
||||
jobs or sync_state."""
|
||||
summary = BatchSourceSummary(source_id=self.source_id)
|
||||
changes: list[BatchChange] = []
|
||||
if self._breaker.is_open:
|
||||
self._last_skip_reason = "circuit_open"
|
||||
logger.debug(
|
||||
"Skipping dry-run discover() — circuit breaker open for %s",
|
||||
self.source_id,
|
||||
)
|
||||
return False, summary, changes
|
||||
with logfire.span("ingester.poller.dry_run", source_id=self.source_id) as span:
|
||||
if await self._jobs.has_pending(self.source_id):
|
||||
self._last_skip_reason = "pending_work"
|
||||
span.set_attribute("skipped", True)
|
||||
span.set_attribute("skip_reason", "pending_work")
|
||||
logger.debug(
|
||||
"Skipping dry-run discover() — %s has pending work in the queue",
|
||||
self.source_id,
|
||||
)
|
||||
return False, summary, changes
|
||||
try:
|
||||
revisions = await self._sync.get_revision_snapshot(self.source_id)
|
||||
known = await self._sync.list_known_uris(self.source_id)
|
||||
async for event in self.source.discover(
|
||||
since=revisions, known_uris=known
|
||||
):
|
||||
if event.kind is SourceEventKind.UPSERT:
|
||||
summary.upsert_count += 1
|
||||
changes.append(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id=event.source_id,
|
||||
uri=event.uri,
|
||||
revision=event.revision,
|
||||
discovered_at=event.discovered_at,
|
||||
)
|
||||
)
|
||||
elif event.kind is SourceEventKind.UNCHANGED:
|
||||
summary.unchanged_count += 1
|
||||
elif event.kind is SourceEventKind.DELETE:
|
||||
if self.config.delete_orphans:
|
||||
summary.delete_count += 1
|
||||
changes.append(
|
||||
BatchChange(
|
||||
op=JobOp.DELETE,
|
||||
source_id=event.source_id,
|
||||
uri=event.uri,
|
||||
revision=None,
|
||||
discovered_at=event.discovered_at,
|
||||
)
|
||||
)
|
||||
else:
|
||||
summary.ignored_delete_count += 1
|
||||
self._breaker.record_success()
|
||||
self._last_polled_at = datetime.now(UTC)
|
||||
self._last_skip_reason = None
|
||||
span.set_attribute("upsert", summary.upsert_count)
|
||||
span.set_attribute("delete", summary.delete_count)
|
||||
span.set_attribute("unchanged", summary.unchanged_count)
|
||||
return True, summary, changes
|
||||
except Exception as exc:
|
||||
self._breaker.record_failure()
|
||||
span.set_attribute(
|
||||
"consecutive_failures", self._breaker.consecutive_failures
|
||||
)
|
||||
span.record_exception(exc)
|
||||
logger.exception(
|
||||
"dry-run discover() failed for %s (consecutive=%d): %s",
|
||||
self.source_id,
|
||||
self._breaker.consecutive_failures,
|
||||
exc,
|
||||
)
|
||||
return False, summary, changes
|
||||
|
||||
async def _handle_event(
|
||||
self,
|
||||
event: SourceEvent,
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from haiku.rag.config import FSSourceConfig, SourceConfig
|
||||
from haiku.rag.ingester.batch import BatchManifest
|
||||
from haiku.rag.ingester.pollers.base import BasePoller
|
||||
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
||||
from haiku.rag.ingester.pollers.factory import build_source
|
||||
|
|
@ -90,6 +92,27 @@ class PollerManager:
|
|||
failed.append(poller.source_id)
|
||||
return failed
|
||||
|
||||
async def dry_run_manifest(self) -> tuple[BatchManifest, list[str]]:
|
||||
"""Collect what one sweep across every source would enqueue without
|
||||
mutating queue jobs or sync_state."""
|
||||
failed: list[str] = []
|
||||
summaries = []
|
||||
changes = []
|
||||
for poller in self._pollers:
|
||||
ok, summary, source_changes = await poller._dry_run_once()
|
||||
summaries.append(summary)
|
||||
changes.extend(source_changes)
|
||||
if not ok:
|
||||
failed.append(poller.source_id)
|
||||
return (
|
||||
BatchManifest(
|
||||
generated_at=datetime.now(UTC),
|
||||
sources=summaries,
|
||||
changes=changes,
|
||||
),
|
||||
failed,
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
for poller in self._pollers:
|
||||
await poller.stop()
|
||||
|
|
|
|||
|
|
@ -371,6 +371,27 @@ class JobRepo:
|
|||
rows = (await conn.execute(query)).all()
|
||||
return {status: n for status, n in rows}
|
||||
|
||||
async def batch_progress_counts_since(self, since: datetime) -> dict[str, int]:
|
||||
"""Counts for a one-shot batch progress snapshot.
|
||||
|
||||
Live rows are counted regardless of enqueue time because run-batch
|
||||
drains the whole pending queue. Terminal rows are counted only when
|
||||
this run completed them, matching BatchReport semantics.
|
||||
"""
|
||||
query = (
|
||||
sa.select(jobs.c.status, sa.func.count().label("n"))
|
||||
.where(
|
||||
sa.or_(
|
||||
jobs.c.status.in_(["queued", "claimed"]),
|
||||
jobs.c.completed_at >= since.isoformat(),
|
||||
)
|
||||
)
|
||||
.group_by(jobs.c.status)
|
||||
)
|
||||
async with self._engine.connect() as conn:
|
||||
rows = (await conn.execute(query)).all()
|
||||
return {status: n for status, n in rows}
|
||||
|
||||
async def count_succeeded_since(self, seconds: int) -> int:
|
||||
"""How many jobs reached `succeeded` in the last `seconds` seconds.
|
||||
Drives the dashboard's rolling-throughput chips."""
|
||||
|
|
|
|||
|
|
@ -110,23 +110,38 @@ async def run_job(
|
|||
),
|
||||
):
|
||||
try:
|
||||
manifest_context = extra.get("_manifest")
|
||||
if job.op is JobOp.DELETE:
|
||||
# An atomic-rename save can let a spurious DELETE win the
|
||||
# enqueue race while the file is mid-rewrite. If the resource
|
||||
# is already back, skip the delete (it would blackhole a live
|
||||
# document) and let the next sweep re-ingest it.
|
||||
try:
|
||||
source = resolve_configured_source(job.uri, job.source_id, sources)
|
||||
restored = await source.head(job.uri) is not None
|
||||
except Exception:
|
||||
restored = False
|
||||
if restored:
|
||||
return JobResult(deleted=False)
|
||||
# document) and let the next sweep re-ingest it. Manifest
|
||||
# replay intentionally follows the frozen dry-run changeset.
|
||||
if manifest_context is None:
|
||||
try:
|
||||
source = resolve_configured_source(
|
||||
job.uri, job.source_id, sources
|
||||
)
|
||||
restored = await source.head(job.uri) is not None
|
||||
except Exception:
|
||||
restored = False
|
||||
if restored:
|
||||
return JobResult(deleted=False)
|
||||
doc = await client.get_document_by_uri(job.uri)
|
||||
if doc is not None and doc.id is not None:
|
||||
await client.delete_document(doc.id)
|
||||
return JobResult(deleted=True)
|
||||
|
||||
if manifest_context is not None and job.revision is not None:
|
||||
source = resolve_configured_source(job.uri, job.source_id, sources)
|
||||
current_revision = await source.head(job.uri)
|
||||
if current_revision != job.revision:
|
||||
raise PermanentError(
|
||||
"manifest revision is stale for "
|
||||
f"{job.uri}: expected {job.revision!r}, "
|
||||
f"current {current_revision!r}"
|
||||
)
|
||||
|
||||
result = await client.create_document_from_source(
|
||||
job.uri,
|
||||
sources=sources,
|
||||
|
|
|
|||
|
|
@ -1,13 +1,25 @@
|
|||
"""haiku-ingester CLI: exercises every subcommand via CliRunner with
|
||||
IngesterApp / open_queue patched out so no real ingestion runs."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from haiku.rag.ingester.app import BatchReport
|
||||
from haiku.rag.config import QueueConfig
|
||||
from haiku.rag.ingester.app import BatchProgress, BatchReport
|
||||
from haiku.rag.ingester.batch import (
|
||||
BatchChange,
|
||||
BatchDryRunReport,
|
||||
BatchManifest,
|
||||
BatchSourceSummary,
|
||||
)
|
||||
from haiku.rag.ingester.cli import _cli as cli
|
||||
from haiku.rag.ingester.cli import _resolve_queue_config
|
||||
from haiku.rag.ingester.queue.models import JobOp
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
|
@ -15,6 +27,17 @@ runner = CliRunner()
|
|||
# --- helpers ---
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _progress_context(callback):
|
||||
yield callback
|
||||
|
||||
|
||||
def _config_with_queue(queue: QueueConfig):
|
||||
config = MagicMock()
|
||||
config.ingester.queue = queue
|
||||
return config
|
||||
|
||||
|
||||
def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
||||
fake = AsyncMock()
|
||||
fake.run_batch.return_value = report
|
||||
|
|
@ -22,6 +45,54 @@ def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
|||
return fake
|
||||
|
||||
|
||||
def _manifest() -> BatchManifest:
|
||||
now = datetime(2026, 6, 22, 10, 30, tzinfo=UTC)
|
||||
return BatchManifest(
|
||||
generated_at=now,
|
||||
sources=[
|
||||
BatchSourceSummary(
|
||||
source_id="docs",
|
||||
upsert_count=1,
|
||||
delete_count=1,
|
||||
unchanged_count=2,
|
||||
)
|
||||
],
|
||||
changes=[
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="docs",
|
||||
uri="file:///a.md",
|
||||
revision="r1",
|
||||
discovered_at=now,
|
||||
),
|
||||
BatchChange(
|
||||
op=JobOp.DELETE,
|
||||
source_id="docs",
|
||||
uri="file:///gone.md",
|
||||
discovered_at=now,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _fake_dry_run_app(report: BatchDryRunReport, monkeypatch) -> AsyncMock:
|
||||
fake = AsyncMock()
|
||||
fake.run_batch_dry_run.return_value = report
|
||||
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||
return fake
|
||||
|
||||
|
||||
def _fake_manifest_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
||||
fake = AsyncMock()
|
||||
fake.run_batch_from_manifest.return_value = report
|
||||
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||
return fake
|
||||
|
||||
|
||||
def _write_manifest(path) -> None:
|
||||
path.write_text(yaml.safe_dump(_manifest().model_dump(mode="json")))
|
||||
|
||||
|
||||
def test_run_batch_reports_and_exits_zero(monkeypatch):
|
||||
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
|
||||
|
||||
|
|
@ -32,6 +103,27 @@ def test_run_batch_reports_and_exits_zero(monkeypatch):
|
|||
fake.run_batch.assert_awaited_once()
|
||||
|
||||
|
||||
def test_run_batch_passes_progress_callback_when_enabled(monkeypatch):
|
||||
fake = AsyncMock()
|
||||
|
||||
async def _run_batch(*, progress_callback):
|
||||
progress_callback(BatchProgress(total=1, completed=1, succeeded=1))
|
||||
return BatchReport(succeeded=1, dead=0)
|
||||
|
||||
fake.run_batch.side_effect = _run_batch
|
||||
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.ingester.cli._batch_progress",
|
||||
lambda _: _progress_context(lambda snapshot: None),
|
||||
)
|
||||
|
||||
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
fake.run_batch.assert_awaited_once()
|
||||
assert "1 succeeded, 0 dead" in result.output
|
||||
|
||||
|
||||
def test_run_batch_exits_nonzero_when_dead(monkeypatch):
|
||||
_fake_app(BatchReport(succeeded=1, dead=2), monkeypatch)
|
||||
|
||||
|
|
@ -50,6 +142,181 @@ def test_run_batch_exits_nonzero_when_sweep_fails(monkeypatch):
|
|||
assert "failed to sweep: docs" in result.output
|
||||
|
||||
|
||||
def test_run_batch_dry_run_writes_default_manifest(monkeypatch, tmp_path):
|
||||
fake = _fake_dry_run_app(
|
||||
BatchDryRunReport(manifest=_manifest()),
|
||||
monkeypatch,
|
||||
)
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=tmp_path):
|
||||
result = runner.invoke(cli, ["run-batch", "--dry-run", "--db", "x.lancedb"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Dry run complete: 1 upsert, 1 delete, 2 unchanged -> manifest-" in (
|
||||
result.output
|
||||
)
|
||||
written = list(tmp_path.glob("*/manifest-*.yaml"))
|
||||
assert len(written) == 1
|
||||
data = yaml.safe_load(written[0].read_text())
|
||||
|
||||
assert data["version"] == 1
|
||||
assert data["sources"][0]["source_id"] == "docs"
|
||||
assert [change["op"] for change in data["changes"]] == ["upsert", "delete"]
|
||||
fake.run_batch_dry_run.assert_awaited_once()
|
||||
fake.run_batch.assert_not_awaited()
|
||||
|
||||
|
||||
def test_run_batch_dry_run_writes_explicit_output(monkeypatch, tmp_path):
|
||||
output = tmp_path / "custom.yaml"
|
||||
_fake_dry_run_app(BatchDryRunReport(manifest=_manifest()), monkeypatch)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
["run-batch", "--dry-run", "--output", str(output), "--db", "x.lancedb"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert f"-> {output}" in result.output
|
||||
data = yaml.safe_load(output.read_text())
|
||||
assert data["changes"][0]["uri"] == "file:///a.md"
|
||||
|
||||
|
||||
def test_run_batch_dry_run_exits_nonzero_when_sweep_fails(monkeypatch, tmp_path):
|
||||
output = tmp_path / "failed.yaml"
|
||||
_fake_dry_run_app(
|
||||
BatchDryRunReport(manifest=_manifest(), failed_sweeps=["docs"]),
|
||||
monkeypatch,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
["run-batch", "--dry-run", "--output", str(output), "--db", "x.lancedb"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "failed to sweep: docs" in result.output
|
||||
assert not output.exists()
|
||||
|
||||
|
||||
def test_run_batch_manifest_replays_manifest(monkeypatch, tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
fake = _fake_manifest_app(BatchReport(succeeded=2, dead=0), monkeypatch)
|
||||
|
||||
result = runner.invoke(
|
||||
cli, ["run-batch", "--manifest", str(manifest_path), "--db", "x.lancedb"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Manifest batch complete: 2 succeeded, 0 dead" in result.output
|
||||
fake.run_batch_from_manifest.assert_awaited_once()
|
||||
loaded = fake.run_batch_from_manifest.await_args.args[0]
|
||||
assert isinstance(loaded, BatchManifest)
|
||||
assert loaded.changes[0].uri == "file:///a.md"
|
||||
fake.run_batch.assert_not_awaited()
|
||||
fake.run_batch_dry_run.assert_not_awaited()
|
||||
|
||||
|
||||
def test_run_batch_manifest_passes_progress_callback(monkeypatch, tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
fake = AsyncMock()
|
||||
|
||||
async def _run_manifest(_manifest, *, progress_callback):
|
||||
progress_callback(BatchProgress(total=1, completed=1, succeeded=1))
|
||||
return BatchReport(succeeded=1, dead=0)
|
||||
|
||||
fake.run_batch_from_manifest.side_effect = _run_manifest
|
||||
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.ingester.cli._batch_progress",
|
||||
lambda _: _progress_context(lambda snapshot: None),
|
||||
)
|
||||
|
||||
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||
|
||||
assert result.exit_code == 0
|
||||
fake.run_batch_from_manifest.assert_awaited_once()
|
||||
assert "1 succeeded, 0 dead" in result.output
|
||||
|
||||
|
||||
def test_run_batch_manifest_exits_nonzero_when_dead(monkeypatch, tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
_fake_manifest_app(BatchReport(succeeded=1, dead=1), monkeypatch)
|
||||
|
||||
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "1 dead" in result.output
|
||||
|
||||
|
||||
def test_run_batch_manifest_reports_validation_error(monkeypatch, tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
fake = AsyncMock()
|
||||
fake.run_batch_from_manifest.side_effect = ValueError("bad manifest")
|
||||
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||
|
||||
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "Error: bad manifest" in result.output
|
||||
|
||||
|
||||
def test_run_batch_manifest_conflicts_with_dry_run(tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
|
||||
result = runner.invoke(
|
||||
cli, ["run-batch", "--manifest", str(manifest_path), "--dry-run"]
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "--manifest cannot be combined with --dry-run" in result.output
|
||||
|
||||
|
||||
def test_run_batch_manifest_conflicts_with_output(tmp_path):
|
||||
manifest_path = tmp_path / "manifest.yaml"
|
||||
_write_manifest(manifest_path)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"run-batch",
|
||||
"--manifest",
|
||||
str(manifest_path),
|
||||
"--output",
|
||||
str(tmp_path / "out.yaml"),
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "--output is only valid with --dry-run" in result.output
|
||||
|
||||
|
||||
def test_run_batch_output_requires_dry_run(tmp_path):
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
["run-batch", "--output", str(tmp_path / "out.yaml")],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "--output is only valid with --dry-run" in result.output
|
||||
|
||||
|
||||
def test_resolve_queue_config_keeps_dburi_when_path_override_present(tmp_path):
|
||||
queue = QueueConfig(
|
||||
dburi="postgresql+asyncpg://user:pass@example.test/db",
|
||||
path=tmp_path / "configured.db",
|
||||
)
|
||||
config = _config_with_queue(queue)
|
||||
|
||||
resolved = _resolve_queue_config(config, tmp_path / "override.db")
|
||||
|
||||
assert resolved is queue
|
||||
|
||||
|
||||
# --- serve ---
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ def _job(
|
|||
*,
|
||||
op: JobOp = JobOp.UPSERT,
|
||||
uri: str = "https://example.com/a.pdf",
|
||||
revision: str | None = None,
|
||||
extra: dict | None = None,
|
||||
attempts: int = 0,
|
||||
) -> Job:
|
||||
|
|
@ -25,6 +26,7 @@ def _job(
|
|||
source_id="src",
|
||||
uri=uri,
|
||||
op=op,
|
||||
revision=revision,
|
||||
status=JobStatus.CLAIMED,
|
||||
attempts=attempts,
|
||||
max_attempts=5,
|
||||
|
|
@ -255,6 +257,22 @@ async def test_delete_skipped_when_resource_restored_on_source():
|
|||
client.delete_document.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manifest_delete_proceeds_when_resource_restored_on_source():
|
||||
client = _mock_client()
|
||||
client.get_document_by_uri.return_value = Document(id="doc-9", content="", uri="u")
|
||||
sources: list[Source] = [_StubSource("src", "12345")]
|
||||
|
||||
result = await run_job(
|
||||
client,
|
||||
_job(op=JobOp.DELETE, extra={"_manifest": {"version": 1}}),
|
||||
sources=sources,
|
||||
)
|
||||
|
||||
assert result.deleted is True
|
||||
client.delete_document.assert_awaited_once_with("doc-9")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_proceeds_when_resource_absent_on_source():
|
||||
client = _mock_client()
|
||||
|
|
@ -267,6 +285,42 @@ async def test_delete_proceeds_when_resource_absent_on_source():
|
|||
client.delete_document.assert_awaited_once_with("doc-9")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manifest_upsert_rejects_stale_revision_before_fetch():
|
||||
client = _mock_client()
|
||||
sources: list[Source] = [_StubSource("src", "r2")]
|
||||
|
||||
with pytest.raises(PermanentError, match="manifest revision is stale"):
|
||||
await run_job(
|
||||
client,
|
||||
_job(revision="r1", extra={"_manifest": {"version": 1}}),
|
||||
sources=sources,
|
||||
)
|
||||
|
||||
client.create_document_from_source.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manifest_upsert_ingests_when_revision_matches():
|
||||
client = _mock_client()
|
||||
client.create_document_from_source.return_value = Document(
|
||||
id="doc-42",
|
||||
content="x",
|
||||
uri="https://example.com/a.pdf",
|
||||
metadata={"md5": "abcd", "source_revision": "r1"},
|
||||
)
|
||||
sources: list[Source] = [_StubSource("src", "r1")]
|
||||
|
||||
result = await run_job(
|
||||
client,
|
||||
_job(revision="r1", extra={"_manifest": {"version": 1}}),
|
||||
sources=sources,
|
||||
)
|
||||
|
||||
assert result.document_id == "doc-42"
|
||||
client.create_document_from_source.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_proceeds_when_source_unresolvable():
|
||||
"""No configured source for the job: the probe can't run, so the delete
|
||||
|
|
|
|||
|
|
@ -309,6 +309,78 @@ async def test_per_source_retry_policy_overrides_default(jobs, sync, tmp_path):
|
|||
assert queued[0].max_attempts == 9
|
||||
|
||||
|
||||
# --- dry-run collection ---
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_reports_changes_without_mutating_queue_or_sync(
|
||||
fs_config, jobs, sync
|
||||
):
|
||||
source = _StubSource(
|
||||
"src",
|
||||
[
|
||||
[
|
||||
_event("file:///a.md", revision="r1"),
|
||||
_event(
|
||||
"file:///b.md",
|
||||
kind=SourceEventKind.UNCHANGED,
|
||||
revision="r2",
|
||||
),
|
||||
_event("file:///gone.md", kind=SourceEventKind.DELETE),
|
||||
]
|
||||
],
|
||||
)
|
||||
poller = _periodic(source, fs_config, jobs, sync)
|
||||
|
||||
ok, summary, changes = await poller._dry_run_once()
|
||||
|
||||
assert ok is True
|
||||
assert summary.upsert_count == 1
|
||||
assert summary.delete_count == 1
|
||||
assert summary.unchanged_count == 1
|
||||
assert summary.ignored_delete_count == 0
|
||||
assert [(c.op, c.uri, c.revision) for c in changes] == [
|
||||
(JobOp.UPSERT, "file:///a.md", "r1"),
|
||||
(JobOp.DELETE, "file:///gone.md", None),
|
||||
]
|
||||
assert await jobs.list_jobs(source_id="src") == []
|
||||
assert await sync.list_known_uris("src") == set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_counts_ignored_deletes_when_orphan_delete_disabled(
|
||||
fs_config, jobs, sync
|
||||
):
|
||||
cfg = fs_config.model_copy(update={"delete_orphans": False})
|
||||
source = _StubSource(
|
||||
"src", [[_event("file:///gone.md", kind=SourceEventKind.DELETE)]]
|
||||
)
|
||||
poller = _periodic(source, cfg, jobs, sync)
|
||||
|
||||
ok, summary, changes = await poller._dry_run_once()
|
||||
|
||||
assert ok is True
|
||||
assert summary.delete_count == 0
|
||||
assert summary.ignored_delete_count == 1
|
||||
assert changes == []
|
||||
assert await jobs.list_jobs(source_id="src") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_skips_when_queue_has_pending_work(fs_config, jobs, sync):
|
||||
await jobs.enqueue("src", "file:///already.md", op=JobOp.UPSERT)
|
||||
source = _StubSource("src", [[_event("file:///a.md")]])
|
||||
poller = _periodic(source, fs_config, jobs, sync)
|
||||
|
||||
ok, summary, changes = await poller._dry_run_once()
|
||||
|
||||
assert ok is False
|
||||
assert summary.source_id == "src"
|
||||
assert changes == []
|
||||
assert source.discover_calls == 0
|
||||
assert poller.last_skip_reason == "pending_work"
|
||||
|
||||
|
||||
# --- PollerManager lifecycle ---
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ pruning, not embedding."""
|
|||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
|
@ -23,7 +24,11 @@ from haiku.rag.config import (
|
|||
WorkerConfig,
|
||||
)
|
||||
from haiku.rag.ingester.app import IngesterApp
|
||||
from haiku.rag.ingester.batch import BatchChange, BatchManifest
|
||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||
from haiku.rag.ingester.queue.migrations import open_queue
|
||||
from haiku.rag.ingester.queue.models import JobOp
|
||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||
from haiku.rag.ingester.workers.pool import WorkerPool
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
|
@ -83,6 +88,10 @@ def _mock_client() -> AsyncMock:
|
|||
return client
|
||||
|
||||
|
||||
def _manifest(*changes: BatchChange) -> BatchManifest:
|
||||
return BatchManifest(generated_at=datetime.now(UTC), changes=list(changes))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def use_client(monkeypatch):
|
||||
"""Make IngesterApp's internally-created HaikuRAG resolve to the given
|
||||
|
|
@ -123,6 +132,30 @@ async def test_run_batch_drains_upserts(tmp_path, use_client):
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_reports_progress(tmp_path, use_client):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
(tmp_path / "b.md").write_text("world")
|
||||
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
progress = []
|
||||
|
||||
report = await IngesterApp(
|
||||
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||
).run_batch(progress_callback=progress.append)
|
||||
|
||||
assert report.succeeded == 2
|
||||
assert report.dead == 0
|
||||
assert progress
|
||||
assert progress[-1].total == 2
|
||||
assert progress[-1].completed == 2
|
||||
assert progress[-1].succeeded == 2
|
||||
assert progress[-1].dead == 0
|
||||
assert progress[-1].queued == 0
|
||||
assert progress[-1].claimed == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_prunes_orphans(tmp_path, use_client):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
|
|
@ -239,6 +272,290 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
|
|||
client.create_document_from_source.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_dry_run_reports_manifest_without_mutating_queue(tmp_path):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
config = _config(tmp_path)
|
||||
db_path = tmp_path / "db.lancedb"
|
||||
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
sync = SyncStateRepo(engine)
|
||||
await sync.upsert("local", (tmp_path / "gone.md").as_uri(), revision="old")
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
report = await IngesterApp(config=config, db_path=db_path).run_batch_dry_run()
|
||||
|
||||
assert report.failed_sweeps == []
|
||||
assert report.manifest.version == 1
|
||||
assert [(change.op, change.uri) for change in report.manifest.changes] == [
|
||||
(JobOp.UPSERT, (tmp_path / "a.md").as_uri()),
|
||||
(JobOp.DELETE, (tmp_path / "gone.md").as_uri()),
|
||||
]
|
||||
source_summary = report.manifest.sources[0]
|
||||
assert source_summary.source_id == "local"
|
||||
assert source_summary.upsert_count == 1
|
||||
assert source_summary.delete_count == 1
|
||||
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
jobs = JobRepo(engine)
|
||||
sync = SyncStateRepo(engine)
|
||||
assert await jobs.list_jobs(source_id="local") == []
|
||||
assert await sync.list_known_uris("local") == {(tmp_path / "gone.md").as_uri()}
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_drains_changes_without_sweeping(
|
||||
tmp_path, use_client, monkeypatch
|
||||
):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
revision = str((tmp_path / "a.md").stat().st_mtime_ns)
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
sweep_all = AsyncMock(side_effect=AssertionError("manifest replay must not sweep"))
|
||||
monkeypatch.setattr(PollerManager, "sweep_all", sweep_all)
|
||||
|
||||
report = await IngesterApp(
|
||||
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(
|
||||
_manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
revision=revision,
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
assert report.succeeded == 1
|
||||
assert report.dead == 0
|
||||
client.create_document_from_source.assert_awaited_once()
|
||||
sweep_all.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_rejects_stale_upsert_revision(
|
||||
tmp_path, use_client
|
||||
):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
|
||||
report = await IngesterApp(
|
||||
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(
|
||||
_manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
revision="stale",
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
assert report.succeeded == 0
|
||||
assert report.dead == 1
|
||||
client.create_document_from_source.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_delete_uses_manifest_even_if_file_reappears(
|
||||
tmp_path, use_client
|
||||
):
|
||||
path = tmp_path / "gone.md"
|
||||
path.write_text("back")
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
|
||||
report = await IngesterApp(
|
||||
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(
|
||||
_manifest(
|
||||
BatchChange(
|
||||
op=JobOp.DELETE,
|
||||
source_id="local",
|
||||
uri=path.as_uri(),
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
assert report.succeeded == 1
|
||||
assert report.dead == 0
|
||||
client.delete_document.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_resumes_same_manifest_work(tmp_path, use_client):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
revision = str((tmp_path / "a.md").stat().st_mtime_ns)
|
||||
config = _config(tmp_path)
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
manifest = _manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
revision=revision,
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
jobs = JobRepo(engine)
|
||||
await jobs.enqueue(
|
||||
"local",
|
||||
(tmp_path / "a.md").as_uri(),
|
||||
op=JobOp.UPSERT,
|
||||
revision=revision,
|
||||
extra={
|
||||
"_manifest": {
|
||||
"version": manifest.version,
|
||||
"generated_at": manifest.generated_at.isoformat(),
|
||||
"discovered_at": manifest.changes[0].discovered_at.isoformat(),
|
||||
}
|
||||
},
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
report = await IngesterApp(
|
||||
config=config, db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(manifest)
|
||||
|
||||
assert report.succeeded == 1
|
||||
assert report.dead == 0
|
||||
client.create_document_from_source.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_rejects_non_manifest_pending_work(
|
||||
tmp_path, use_client
|
||||
):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
config = _config(tmp_path)
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
jobs = JobRepo(engine)
|
||||
await jobs.enqueue("local", (tmp_path / "a.md").as_uri(), op=JobOp.UPSERT)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||
await IngesterApp(
|
||||
config=config, db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(
|
||||
_manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_rejects_different_manifest_pending_work(
|
||||
tmp_path, use_client
|
||||
):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
config = _config(tmp_path)
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
manifest = _manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
jobs = JobRepo(engine)
|
||||
await jobs.enqueue(
|
||||
"local",
|
||||
(tmp_path / "a.md").as_uri(),
|
||||
op=JobOp.UPSERT,
|
||||
extra={
|
||||
"_manifest": {
|
||||
"version": 1,
|
||||
"generated_at": "2026-01-01T00:00:00+00:00",
|
||||
"discovered_at": manifest.changes[0].discovered_at.isoformat(),
|
||||
}
|
||||
},
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||
await IngesterApp(
|
||||
config=config, db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(manifest)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_rejects_unrelated_pending_work(
|
||||
tmp_path, use_client
|
||||
):
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
config = _config(tmp_path)
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
engine = await open_queue(config.ingester.queue)
|
||||
try:
|
||||
jobs = JobRepo(engine)
|
||||
await jobs.enqueue("other", "file:///outside.md", op=JobOp.UPSERT)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||
await IngesterApp(
|
||||
config=config, db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(
|
||||
_manifest(
|
||||
BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=(tmp_path / "a.md").as_uri(),
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_from_manifest_rejects_duplicate_changes(tmp_path, use_client):
|
||||
path = tmp_path / "a.md"
|
||||
path.write_text("hello")
|
||||
client = _mock_client()
|
||||
use_client(client)
|
||||
change = BatchChange(
|
||||
op=JobOp.UPSERT,
|
||||
source_id="local",
|
||||
uri=path.as_uri(),
|
||||
discovered_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate change"):
|
||||
await IngesterApp(
|
||||
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||
).run_batch_from_manifest(_manifest(change, change))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_batch_aborts_when_all_workers_die(
|
||||
tmp_path, use_client, monkeypatch, caplog
|
||||
|
|
|
|||
Loading…
Reference in a new issue