diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index c211f581..c88902c2 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING from pydantic import BaseModel from haiku.rag.config import AppConfig +from haiku.rag.ingester.batch import BatchDryRunReport 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 @@ -129,6 +130,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.""" @@ -244,6 +274,14 @@ class IngesterApp: 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 _maybe_start_api(self, api: bool): """Spin up the FastAPI control plane on an asyncio task. Returns (task, server) or (None, None) when the API is disabled.""" 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/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/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..c9030965 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -24,6 +24,9 @@ from haiku.rag.config import ( ) from haiku.rag.ingester.app import IngesterApp 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 @@ -239,6 +242,42 @@ 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_aborts_when_all_workers_die( tmp_path, use_client, monkeypatch, caplog