diff --git a/CHANGELOG.md b/CHANGELOG.md index 36bb8854..1efc42be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 ` 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. diff --git a/docs/ingester.md b/docs/ingester.md index bd60684b..e89b6125 100644 --- a/docs/ingester.md +++ b/docs/ingester.md @@ -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-.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 diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index c211f581..dd08372c 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/ingester/batch.py b/haiku_rag_slim/haiku/rag/ingester/batch.py new file mode 100644 index 00000000..d982d4e2 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/batch.py @@ -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] = [] diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index 82f758f1..21d67355 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -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-.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( diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index 093a8cd0..c3d6a859 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index c05dda5f..a947ddf0 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 9ad32d84..680a63c8 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index f6c4bc90..977abbac 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -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, diff --git a/tests/ingester/test_cli.py b/tests/ingester/test_cli.py index 1e69cdad..385456ab 100644 --- a/tests/ingester/test_cli.py +++ b/tests/ingester/test_cli.py @@ -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 --- diff --git a/tests/ingester/test_pipeline.py b/tests/ingester/test_pipeline.py index 34468670..882c9ee2 100644 --- a/tests/ingester/test_pipeline.py +++ b/tests/ingester/test_pipeline.py @@ -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 diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index eca4d418..c9daed5e 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -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 --- diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index d9643c12..bf87f851 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -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