Add side-effect-free batch dry-run discovery

This commit is contained in:
Yiorgis Gozadinos 2026-06-22 11:28:53 +03:00
parent 6f2ab0963c
commit 4b573bfebd
No known key found for this signature in database
6 changed files with 282 additions and 0 deletions

View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING
from pydantic import BaseModel from pydantic import BaseModel
from haiku.rag.config import AppConfig 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.metadata import build_providers, load_metadata_providers
from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.pollers.manager import PollerManager
from haiku.rag.ingester.queue.migrations import open_queue from haiku.rag.ingester.queue.migrations import open_queue
@ -129,6 +130,35 @@ class IngesterApp:
await self._engine.dispose() await self._engine.dispose()
self._engine = None 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: async def _stop_pool(self) -> None:
"""Stop the worker pool, honouring the shutdown grace, then drain any """Stop the worker pool, honouring the shutdown grace, then drain any
cancel-cleanup release tasks before the queue connection closes.""" cancel-cleanup release tasks before the queue connection closes."""
@ -244,6 +274,14 @@ class IngesterApp:
await self._stop_pool() await self._stop_pool()
await self._pollers.close_sources() 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): async def _maybe_start_api(self, api: bool):
"""Spin up the FastAPI control plane on an asyncio task. Returns """Spin up the FastAPI control plane on an asyncio task. Returns
(task, server) or (None, None) when the API is disabled.""" (task, server) or (None, None) when the API is disabled."""

View 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] = []

View file

@ -4,6 +4,7 @@ import random
from datetime import UTC, datetime from datetime import UTC, datetime
from haiku.rag.config import SourceConfig 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.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.queue.models import JobOp, SyncRow from haiku.rag.ingester.queue.models import JobOp, SyncRow
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
@ -165,6 +166,82 @@ class BasePoller:
) )
return False 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( async def _handle_event(
self, self,
event: SourceEvent, event: SourceEvent,

View file

@ -1,9 +1,11 @@
import asyncio import asyncio
import logging import logging
from collections.abc import Sequence from collections.abc import Sequence
from datetime import UTC, datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from haiku.rag.config import FSSourceConfig, SourceConfig 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.base import BasePoller
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.pollers.factory import build_source from haiku.rag.ingester.pollers.factory import build_source
@ -90,6 +92,27 @@ class PollerManager:
failed.append(poller.source_id) failed.append(poller.source_id)
return failed 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: async def stop(self) -> None:
for poller in self._pollers: for poller in self._pollers:
await poller.stop() await poller.stop()

View file

@ -309,6 +309,78 @@ async def test_per_source_retry_policy_overrides_default(jobs, sync, tmp_path):
assert queued[0].max_attempts == 9 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 --- # --- PollerManager lifecycle ---

View file

@ -24,6 +24,9 @@ from haiku.rag.config import (
) )
from haiku.rag.ingester.app import IngesterApp from haiku.rag.ingester.app import IngesterApp
from haiku.rag.ingester.pollers.manager import PollerManager 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.ingester.workers.pool import WorkerPool
from haiku.rag.store.models.document import Document 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() 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 @pytest.mark.asyncio
async def test_run_batch_aborts_when_all_workers_die( async def test_run_batch_aborts_when_all_workers_die(
tmp_path, use_client, monkeypatch, caplog tmp_path, use_client, monkeypatch, caplog