From b0d0ac588df7c19f3c91c6033e1a0e26c106d50f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 27 May 2026 14:37:25 +0300 Subject: [PATCH] Split snapshot APIs and resolve_fetcher by intent --- haiku_rag_slim/haiku/rag/client/documents.py | 25 +++-- .../haiku/rag/ingester/pollers/base.py | 7 +- .../haiku/rag/ingester/queue/repository.py | 30 ++++-- .../haiku/rag/ingester/sources/__init__.py | 8 +- .../haiku/rag/ingester/sources/base.py | 7 +- .../haiku/rag/ingester/sources/fs.py | 16 ++-- .../haiku/rag/ingester/sources/http.py | 27 +++--- .../haiku/rag/ingester/sources/registry.py | 51 +++++----- .../haiku/rag/ingester/sources/s3.py | 17 ++-- .../haiku/rag/ingester/sources/webdav.py | 12 ++- .../haiku/rag/ingester/workers/pipeline.py | 10 +- .../haiku/rag/ingester/workers/pool.py | 10 +- tests/ingester/test_api.py | 2 +- tests/ingester/test_fs_source.py | 4 +- tests/ingester/test_http_source.py | 18 ++-- tests/ingester/test_pollers.py | 15 +-- tests/ingester/test_queue.py | 30 +++--- tests/ingester/test_resolve_fetcher.py | 96 ++++++++++--------- tests/ingester/test_s3_source.py | 2 +- tests/ingester/test_webdav_source.py | 9 +- tests/ingester/test_workers.py | 8 +- 21 files changed, 233 insertions(+), 171 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index d8611c46..574a309e 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -10,7 +10,11 @@ from haiku.rag.client.processing import ( ) from haiku.rag.client.titles import resolve_title from haiku.rag.converters import get_converter -from haiku.rag.ingester.sources import FetchResult, resolve_fetcher +from haiku.rag.ingester.sources import ( + FetchResult, + resolve_adhoc_fetcher, + resolve_configured_source, +) from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items @@ -372,15 +376,16 @@ async def create_document_from_source( f"Unsupported file extension: {local_path.suffix}" ) - # Single resource — resolve the right Source adapter for this URI. - # `sources` (configured, in-order) wins over scheme-based adhoc adapters - # so worker fetches reuse the authenticated source the poller used. - fetcher = resolve_fetcher( - source_str, - sources=sources, - source_id=source_id, - storage_options=storage_options, - ) + # Worker jobs carry source_id from the poller; strict lookup so a + # renamed/removed source surfaces as a DLQ instead of silently dropping + # credentials. Ad-hoc CLI calls (no source_id) fall back to scheme-based + # adapters when no configured source matches. + if source_id is not None: + fetcher = resolve_configured_source(source_str, source_id, sources) + else: + fetcher = resolve_adhoc_fetcher( + source_str, sources=sources, storage_options=storage_options + ) # The stored URI is what we look up + persist by. For an explicit uri # override, use it as-is. For a file:// input the source string is diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index 1ee107e7..616a4dfd 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -108,13 +108,16 @@ class BasePoller: ) return False try: - snapshot = await self._sync.get_snapshot(self.source_id) + revisions = await self._sync.get_revision_snapshot(self.source_id) + known = await self._sync.list_known_uris(self.source_id) counts = { SourceEventKind.UPSERT: 0, SourceEventKind.DELETE: 0, SourceEventKind.UNCHANGED: 0, } - async for event in self.source.discover(since=snapshot): + async for event in self.source.discover( + since=revisions, known_uris=known + ): counts[event.kind] += 1 await self._handle_event(event) self._breaker.record_success() diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 6bff9901..c8934cd5 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -230,7 +230,8 @@ class JobRepo: return _row_to_job(row) async def cancel(self, job_id: str) -> bool: - """Delete a queued or claimed job. Returns True if a row was removed.""" + """True iff a queued/claimed row was removed; terminal jobs aren't + cancellable (succeeded/dead rows are kept for history).""" async with self._lock: async with self._conn.execute( "DELETE FROM jobs WHERE id=? AND status IN ('queued', 'claimed') RETURNING id", @@ -399,20 +400,33 @@ class SyncStateRepo: # connection. See JobRepo for the SQLite cursor + commit constraint. self._lock = lock or asyncio.Lock() - async def get_snapshot(self, source_id: str) -> dict[str, str | None]: - """uri -> revision map for the source. Revision is None for URIs - known to the source but never successfully ingested with a revision - (HTTP responses without ETag / Last-Modified, or a worker that DLQ'd - before completion). Sources compare with `previous == current`; None - compares as not-equal so the URI is treated as changed.""" + async def get_revision_snapshot(self, source_id: str) -> dict[str, str]: + """uri -> revision map for URIs that have a stored revision. Sources + compare current revision against this map to decide UPSERT vs + UNCHANGED. Rows without a revision (HTTP without ETag, or a worker + that didn't complete) are excluded — they have no revision to + compare against; the closing-loop DELETE diff uses list_known_uris + instead.""" async with self._lock: async with self._conn.execute( - "SELECT uri, revision FROM sync_state WHERE source_id=?", + "SELECT uri, revision FROM sync_state WHERE source_id=? AND revision IS NOT NULL", (source_id,), ) as cursor: rows = await cursor.fetchall() return {row["uri"]: row["revision"] for row in rows} + async def list_known_uris(self, source_id: str) -> set[str]: + """Every URI the source has ever produced. Used by the closing-loop + diff in discover() so a URI previously seen but no longer visible + (FS file deleted, HTTP URL removed from config) emits DELETE.""" + async with self._lock: + async with self._conn.execute( + "SELECT uri FROM sync_state WHERE source_id=?", + (source_id,), + ) as cursor: + rows = await cursor.fetchall() + return {row["uri"] for row in rows} + async def get_row(self, source_id: str, uri: str) -> SyncStateRow | None: async with self._lock: async with self._conn.execute( diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/__init__.py b/haiku_rag_slim/haiku/rag/ingester/sources/__init__.py index 4ece8894..15e1072f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/__init__.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/__init__.py @@ -8,7 +8,10 @@ from haiku.rag.ingester.sources.base import ( from haiku.rag.ingester.sources.filter import FileFilter from haiku.rag.ingester.sources.fs import FSSource from haiku.rag.ingester.sources.http import HTTPSource -from haiku.rag.ingester.sources.registry import resolve_fetcher +from haiku.rag.ingester.sources.registry import ( + resolve_adhoc_fetcher, + resolve_configured_source, +) from haiku.rag.ingester.sources.s3 import S3Source from haiku.rag.ingester.sources.webdav import WebDAVSource @@ -23,5 +26,6 @@ __all__ = [ "SourceEvent", "SourceEventKind", "WebDAVSource", - "resolve_fetcher", + "resolve_adhoc_fetcher", + "resolve_configured_source", ] diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/base.py b/haiku_rag_slim/haiku/rag/ingester/sources/base.py index 7beea83f..32099b35 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/base.py @@ -9,7 +9,7 @@ from pydantic import BaseModel, Field # uri -> revision. Captures what revisions of which URIs we had last seen # for a given source. Passed to discover() so the source can yield only # UPSERT / UNCHANGED / DELETE deltas instead of a full re-scan. -RevisionSnapshot = Mapping[str, str | None] +RevisionSnapshot = Mapping[str, str] class SourceEventKind(StrEnum): @@ -64,5 +64,8 @@ class Source(Protocol): async def fetch(self, uri: str) -> FetchResult: ... def discover( - self, since: RevisionSnapshot | None = None + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, ) -> AsyncIterator[SourceEvent]: ... diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index b6cbd391..8b944e04 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -102,9 +102,13 @@ class FSSource: ) async def discover( - self, since: RevisionSnapshot | None = None + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, ) -> AsyncIterator[SourceEvent]: - snapshot: dict[str, str | None] = dict(since) if since else {} + snapshot: dict[str, str] = dict(since) if since else {} + known = known_uris or set() now = datetime.now(UTC) seen: set[str] = set() @@ -151,11 +155,9 @@ class FSSource: discovered_at=now, ) - # Anything in the snapshot we didn't encounter during the walk is - # gone from the source. Emit DELETE so the poller can clean up. - for uri in snapshot: - if uri in seen: - continue + # Anything previously known to this source that we didn't encounter + # during the walk is gone. Emit DELETE so the poller can clean up. + for uri in known - seen: yield SourceEvent( source_id=self.source_id, uri=uri, diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/http.py b/haiku_rag_slim/haiku/rag/ingester/sources/http.py index 82dab875..83c0bada 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/http.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/http.py @@ -83,19 +83,24 @@ class HTTPSource: ) async def discover( - self, since: RevisionSnapshot | None = None + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, ) -> AsyncIterator[SourceEvent]: # HTTP has no listing concept — discover() only reports on what is - # currently configured in self.urls. URLs that were previously in - # config but aren't now emit DELETE so the poller can clean up - # alongside the in-source 410 signal. + # currently configured in self.urls. URLs that were previously + # known to this source (sync_state) but aren't in the current + # config emit DELETE so the poller can clean up alongside the + # in-source 410 signal. # # 410 Gone is the one real source-side deletion signal: the origin # explicitly says "permanently gone". 404 and other failures are # ambiguous (transient outage, misconfigured URL, auth blip), so we # fall back to UPSERT with no revision and let the worker decide # via GET. - snapshot: dict[str, str | None] = dict(since) if since else {} + snapshot: dict[str, str] = dict(since) if since else {} + known = known_uris or set() now = datetime.now(UTC) configured = set(self.urls) @@ -147,13 +152,11 @@ class HTTPSource: discovered_at=now, ) - # Anything previously ingested by this source that's no longer in - # config (URL removed from `urls`) emits DELETE so delete_orphans - # can clean up. Without this, removing a URL from config leaves the - # document and sync_state indefinitely. - for url in snapshot: - if url in configured: - continue + # Anything previously known to this source that's no longer in + # config emits DELETE so delete_orphans can clean up. Without this, + # removing a URL from config leaves the document and sync_state + # indefinitely. + for url in known - configured: yield SourceEvent( source_id=self.source_id, uri=url, diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/registry.py b/haiku_rag_slim/haiku/rag/ingester/sources/registry.py index c102509d..793c74c0 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/registry.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/registry.py @@ -9,34 +9,41 @@ from haiku.rag.ingester.sources.http import HTTPSource from haiku.rag.ingester.sources.s3 import S3Source -def resolve_fetcher( +def resolve_configured_source( + uri: str, + source_id: str, + sources: Iterable[Source] | None, +) -> Source: + """Strict lookup: return the configured source with this id, or raise. + + Worker jobs carry source_id from when they were enqueued. Falling back + to an ad-hoc fetcher would silently drop credentials when a source has + been renamed or removed from config — better to raise and let the job + DLQ so the misconfiguration surfaces. + """ + for src in sources or (): + if src.source_id == source_id: + if not src.supports(uri): + raise UnsupportedSourceError( + f"Source {source_id!r} doesn't support URI {uri!r}" + ) + return src + raise UnsupportedSourceError( + f"No configured source with id {source_id!r} for URI {uri!r}" + ) + + +def resolve_adhoc_fetcher( uri: str, - sources: Iterable[Source] | None = None, *, - source_id: str | None = None, + sources: Iterable[Source] | None = None, storage_options: dict[str, str] | None = None, ) -> Source: - """Pick a Source adapter for ``uri``. + """Best-effort lookup for one-shot fetches (e.g. ``add-src ``). - When ``source_id`` is given (worker path), the source with that id is - used so credentials/headers of the configured source are reused rather - than picking whichever source happens to match ``supports(uri)`` first. - Without ``source_id`` (ad-hoc ``add-src ``), the first configured - source whose ``supports(uri)`` returns True wins, then a scheme-based - ad-hoc adapter. + Configured ``sources`` win when one matches; otherwise a scheme-based + adapter is built so any URI can be fetched without configuration. """ - if source_id is not None: - for src in sources or (): - if src.source_id == source_id: - if not src.supports(uri): - raise UnsupportedSourceError( - f"Source {source_id!r} doesn't support URI {uri!r}" - ) - return src - raise UnsupportedSourceError( - f"No configured source with id {source_id!r} for URI {uri!r}" - ) - if sources: for src in sources: if src.supports(uri): diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py index 1b5fe5d8..0b0b229f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py @@ -103,13 +103,17 @@ class S3Source: ) async def discover( - self, since: RevisionSnapshot | None = None + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, ) -> AsyncIterator[SourceEvent]: import obstore # type: ignore[import-not-found] from haiku.rag.s3 import make_s3_store - snapshot: dict[str, str | None] = dict(since) if since else {} + snapshot: dict[str, str] = dict(since) if since else {} + known = known_uris or set() now = datetime.now(UTC) seen: set[str] = set() store = make_s3_store(self.bucket, self.storage_options) @@ -136,11 +140,10 @@ class S3Source: discovered_at=now, ) - # URIs in the snapshot that no longer appear under the prefix have - # been deleted upstream — emit DELETE so the poller cleans up. - for uri in snapshot: - if uri in seen: - continue + # URIs previously known to this source that no longer appear under + # the prefix have been deleted upstream — emit DELETE so the poller + # cleans up. + for uri in known - seen: yield SourceEvent( source_id=self.source_id, uri=uri, diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py b/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py index f50cce10..7dd6a15e 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py @@ -251,9 +251,13 @@ class WebDAVSource: ) async def discover( - self, since: RevisionSnapshot | None = None + self, + since: RevisionSnapshot | None = None, + *, + known_uris: set[str] | None = None, ) -> AsyncIterator[SourceEvent]: - snapshot: dict[str, str | None] = dict(since) if since else {} + snapshot: dict[str, str] = dict(since) if since else {} + known = known_uris or set() now = datetime.now(UTC) seen: set[str] = set() @@ -291,9 +295,7 @@ class WebDAVSource: discovered_at=now, ) - for uri in snapshot: - if uri in seen: - continue + for uri in known - seen: yield SourceEvent( source_id=self.source_id, uri=uri, diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index 41229915..3cf97c47 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -73,11 +73,11 @@ async def run_job( sources: list["Source"] | None = None, ) -> JobResult: """Execute the work described by `job`. `sources` is the list of - configured Source adapters — `resolve_fetcher` prefers them over - URI-scheme adhoc adapters so workers reuse the same authenticated / - pre-configured fetch context the pollers used at discovery. Raises - PermanentError or TransientError; the worker uses that to decide - dead vs retry.""" + configured Source adapters; the client looks up `job.source_id` + against it via `resolve_configured_source` so workers reuse the + authenticated/pre-configured fetch context the pollers used at + discovery. Raises PermanentError or TransientError; the worker + uses that to decide dead vs retry.""" extra = job.extra or {} parent_ctx = extra.get("_otel") attach = attach_context(parent_ctx) if parent_ctx else nullcontext() diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index b013f376..4f204a27 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -22,12 +22,10 @@ _WORKER_BREAKER_COOLDOWN_S = 60.0 class WorkerPool: - """Asyncio-based pool. N worker tasks share a bounded Semaphore, each - pulling jobs from the queue and running them through `run_job`. Reaper - task resets claims older than `claim_timeout_s` so a crashed worker - doesn't strand its job. - - Lifecycle: build it, await start(), let it run, await stop(). + """`worker_count` async tasks each pull jobs from the queue and run them + through `run_job`. A reaper task resets claims older than + `claim_timeout_s` so a crashed worker doesn't strand its job. Lifecycle: + build it, await start(), let it run, await stop(). """ def __init__( diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py index db2395a8..4e5143d9 100644 --- a/tests/ingester/test_api.py +++ b/tests/ingester/test_api.py @@ -376,7 +376,7 @@ class _StubSource: async def fetch(self, uri) -> FetchResult: # pragma: no cover raise NotImplementedError - async def discover(self, since=None): + async def discover(self, since=None, *, known_uris=None): events = self._sweeps.pop(0) if self._sweeps else [] for event in events: yield event diff --git a/tests/ingester/test_fs_source.py b/tests/ingester/test_fs_source.py index e8420cea..df32304f 100644 --- a/tests/ingester/test_fs_source.py +++ b/tests/ingester/test_fs_source.py @@ -115,8 +115,8 @@ async def test_fs_source_discover_changed_yields_upsert(fs_root: Path): @pytest.mark.asyncio async def test_fs_source_discover_emits_delete_for_missing(fs_root: Path): src = FSSource(root=fs_root, supported_extensions=[".md", ".txt"]) - snapshot = {(fs_root / "ghost.md").as_uri(): "999"} - events = [e async for e in src.discover(since=snapshot)] + known = {(fs_root / "ghost.md").as_uri()} + events = [e async for e in src.discover(known_uris=known)] deletes = [e for e in events if e.kind is SourceEventKind.DELETE] assert len(deletes) == 1 assert deletes[0].uri == (fs_root / "ghost.md").as_uri() diff --git a/tests/ingester/test_http_source.py b/tests/ingester/test_http_source.py index 9e30a683..8abe9cd7 100644 --- a/tests/ingester/test_http_source.py +++ b/tests/ingester/test_http_source.py @@ -290,7 +290,10 @@ async def test_discover_emits_delete_for_removed_url(): source_id="x", urls=["https://example.com/a.md"], transport=transport ) events = [ - e async for e in src.discover(since={"https://example.com/gone.md": "old"}) + e + async for e in src.discover( + known_uris={"https://example.com/a.md", "https://example.com/gone.md"} + ) ] by_uri = {e.uri: e for e in events} assert by_uri["https://example.com/a.md"].kind is not SourceEventKind.DELETE @@ -299,10 +302,10 @@ async def test_discover_emits_delete_for_removed_url(): @pytest.mark.asyncio -async def test_discover_emits_delete_for_removed_url_with_none_revision(): - """Same as above but the snapshot entry's revision is None — an HTTP - response without ETag or Last-Modified produces that state, and these - rows must still trigger config-removal DELETE.""" +async def test_discover_emits_delete_for_removed_url_with_no_revision_tracked(): + """known_uris alone determines config-removal DELETE — a URL the + source has seen before but never had a revision for (HTTP without + ETag/Last-Modified) still triggers DELETE when dropped from config.""" transport = _transport( { ("HEAD", "https://example.com/a.md"): httpx.Response( @@ -314,7 +317,10 @@ async def test_discover_emits_delete_for_removed_url_with_none_revision(): source_id="x", urls=["https://example.com/a.md"], transport=transport ) events = [ - e async for e in src.discover(since={"https://example.com/no-etag.md": None}) + e + async for e in src.discover( + known_uris={"https://example.com/a.md", "https://example.com/no-etag.md"} + ) ] by_uri = {e.uri: e for e in events} assert by_uri["https://example.com/no-etag.md"].kind is SourceEventKind.DELETE diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index c1b1d48a..9c5c4ad4 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -69,7 +69,7 @@ class _StubSource: async def fetch(self, uri: str) -> FetchResult: # pragma: no cover raise NotImplementedError - async def discover(self, since=None): + async def discover(self, since=None, *, known_uris=None): self.discover_calls += 1 if self.fail_with is not None: raise self.fail_with @@ -130,10 +130,11 @@ async def test_upsert_event_enqueues_job_and_touches_sync_state(fs_config, jobs, assert queued[0].revision == "r1" # Pollers DO NOT write revision to sync_state — the worker does that - # after a successful ingest. The URI shows up in the snapshot with - # revision=None until then. - snapshot = await sync.get_snapshot("src") - assert snapshot == {"file:///a.md": None} + # after a successful ingest. The URI shows up in list_known_uris from + # the moment the poller emits the event, but the revision_snapshot + # stays empty until ingestion completes. + assert await sync.get_revision_snapshot("src") == {} + assert await sync.list_known_uris("src") == {"file:///a.md"} @pytest.mark.asyncio @@ -145,7 +146,7 @@ async def test_unchanged_event_touches_sync_state_no_job(fs_config, jobs, sync): await poller._sweep_once() assert await jobs.list_jobs(source_id="src") == [] - assert await sync.get_snapshot("src") == {"file:///a.md": "r1"} + assert await sync.get_revision_snapshot("src") == {"file:///a.md": "r1"} @pytest.mark.asyncio @@ -232,7 +233,7 @@ async def test_dead_job_does_not_clear_sync_state_revision(fs_config, jobs, sync assert row.revision == "r1" await poller._sweep_once() - assert await sync.get_snapshot("src") == {"file:///a.md": "r1"} + assert await sync.get_revision_snapshot("src") == {"file:///a.md": "r1"} @pytest.mark.asyncio diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index 152bc72e..7404fa1f 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -720,48 +720,48 @@ def test_repos_default_to_independent_locks_when_used_alone(conn): @pytest.mark.asyncio -async def test_sync_state_get_snapshot_empty(sync): - assert await sync.get_snapshot("unknown") == {} +async def test_sync_state_get_revision_snapshot_empty(sync): + assert await sync.get_revision_snapshot("unknown") == {} @pytest.mark.asyncio async def test_sync_state_upsert_and_get(sync): await sync.upsert("s", "u1", revision="abc", content_hash="m1") await sync.upsert("s", "u2", revision="def", content_hash="m2") - assert await sync.get_snapshot("s") == {"u1": "abc", "u2": "def"} + assert await sync.get_revision_snapshot("s") == {"u1": "abc", "u2": "def"} @pytest.mark.asyncio async def test_sync_state_upsert_overwrites(sync): await sync.upsert("s", "u1", revision="abc", content_hash="m1") await sync.upsert("s", "u1", revision="def", content_hash="m2", ingested=True) - assert await sync.get_snapshot("s") == {"u1": "def"} + assert await sync.get_revision_snapshot("s") == {"u1": "def"} @pytest.mark.asyncio async def test_sync_state_delete_removes_entry(sync): await sync.upsert("s", "u1", revision="abc", content_hash="m1") await sync.delete("s", "u1") - assert await sync.get_snapshot("s") == {} + assert await sync.get_revision_snapshot("s") == {} @pytest.mark.asyncio async def test_sync_state_snapshot_scoped_per_source(sync): await sync.upsert("s1", "u", revision="abc", content_hash="m") await sync.upsert("s2", "u", revision="def", content_hash="m") - assert await sync.get_snapshot("s1") == {"u": "abc"} - assert await sync.get_snapshot("s2") == {"u": "def"} + assert await sync.get_revision_snapshot("s1") == {"u": "abc"} + assert await sync.get_revision_snapshot("s2") == {"u": "def"} @pytest.mark.asyncio -async def test_sync_state_snapshot_includes_null_revision_rows(sync): - """A URI known to a source but never successfully ingested with a - revision (HTTP without ETag/Last-Modified, or a worker DLQ'd before - completion) is still in the snapshot — with revision=None. Needed so - config-removal DELETE detection sees these rows.""" +async def test_sync_state_revision_snapshot_excludes_null_revision_rows(sync): + """get_revision_snapshot returns only rows with a stored revision — + sources compare against this for UPSERT/UNCHANGED. Rows without a + revision (HTTP without ETag, worker DLQ'd before completion) are + visible via list_known_uris instead.""" await sync.upsert("s", "u", revision=None, content_hash=None) - snapshot = await sync.get_snapshot("s") - assert snapshot == {"u": None} + assert await sync.get_revision_snapshot("s") == {} + assert await sync.list_known_uris("s") == {"u"} @pytest.mark.asyncio @@ -773,7 +773,7 @@ async def test_sync_state_upsert_preserves_revision_when_none(sync): assert row is not None assert row.revision == "v1" assert row.content_hash == "hash-v1" - assert await sync.get_snapshot("s") == {"u": "v1"} + assert await sync.get_revision_snapshot("s") == {"u": "v1"} @pytest.mark.asyncio diff --git a/tests/ingester/test_resolve_fetcher.py b/tests/ingester/test_resolve_fetcher.py index 51f2c30e..7da7debb 100644 --- a/tests/ingester/test_resolve_fetcher.py +++ b/tests/ingester/test_resolve_fetcher.py @@ -2,36 +2,41 @@ from pathlib import Path import pytest -from haiku.rag.ingester.sources import resolve_fetcher +from haiku.rag.ingester.sources import ( + resolve_adhoc_fetcher, + resolve_configured_source, +) from haiku.rag.ingester.sources.fs import FSSource from haiku.rag.ingester.sources.http import HTTPSource from haiku.rag.ingester.sources.s3 import S3Source +# --- resolve_adhoc_fetcher --- -def test_resolves_fs_for_file_uri(): - src = resolve_fetcher("file:///tmp/sample.md") + +def test_adhoc_resolves_fs_for_file_uri(): + src = resolve_adhoc_fetcher("file:///tmp/sample.md") assert isinstance(src, FSSource) -def test_resolves_fs_for_bare_path(): - src = resolve_fetcher("/tmp/sample.md") +def test_adhoc_resolves_fs_for_bare_path(): + src = resolve_adhoc_fetcher("/tmp/sample.md") assert isinstance(src, FSSource) -def test_resolves_http(): - src = resolve_fetcher("https://example.com/x.pdf") +def test_adhoc_resolves_http(): + src = resolve_adhoc_fetcher("https://example.com/x.pdf") assert isinstance(src, HTTPSource) -def test_resolves_s3_scopes_to_bucket(): - src = resolve_fetcher("s3://my-bucket/key.pdf") +def test_adhoc_resolves_s3_scopes_to_bucket(): + src = resolve_adhoc_fetcher("s3://my-bucket/key.pdf") assert isinstance(src, S3Source) assert src.bucket == "my-bucket" assert src.prefix == "" -def test_resolves_s3_forwards_storage_options(): - src = resolve_fetcher( +def test_adhoc_resolves_s3_forwards_storage_options(): + src = resolve_adhoc_fetcher( "s3://my-bucket/key.pdf", storage_options={"endpoint": "http://seaweed:8333", "allow_http": "true"}, ) @@ -39,64 +44,69 @@ def test_resolves_s3_forwards_storage_options(): assert src.storage_options["endpoint"] == "http://seaweed:8333" -def test_unknown_scheme_raises(): +def test_adhoc_unknown_scheme_raises(): with pytest.raises(ValueError, match="No source adapter"): - resolve_fetcher("ftp://example.com/x") + resolve_adhoc_fetcher("ftp://example.com/x") -def test_configured_source_matches_first(tmp_path: Path): +def test_adhoc_configured_source_matches_first(tmp_path: Path): (tmp_path / "a.md").write_text("hi") fs = FSSource(root=tmp_path) - chosen = resolve_fetcher((tmp_path / "a.md").as_uri(), sources=[fs]) + chosen = resolve_adhoc_fetcher((tmp_path / "a.md").as_uri(), sources=[fs]) assert chosen is fs -def test_configured_source_falls_through_when_no_match(): +def test_adhoc_configured_source_falls_through_when_no_match(): fs = FSSource(root=Path("/tmp")) # https doesn't match an FS source — fall back to ad-hoc HTTPSource - chosen = resolve_fetcher("https://example.com/x", sources=[fs]) + chosen = resolve_adhoc_fetcher("https://example.com/x", sources=[fs]) assert isinstance(chosen, HTTPSource) assert chosen is not fs -def test_source_id_prefers_matching_source_over_first_supports(): - """Two HTTPSource configs with different auth headers: passing - source_id picks the right one. Without source_id, the first one to - return True from supports(uri) wins regardless of which config - actually owns the job.""" +def test_adhoc_returns_first_supporting_source_with_multiple_http(): + """Without source_id (ad-hoc mode), the first configured HTTP source + that supports the URI wins. Used by `add-src` from the CLI.""" arxiv = HTTPSource(source_id="arxiv", headers={"Authorization": "Bearer A"}) intranet = HTTPSource(source_id="intranet", headers={"Authorization": "Bearer B"}) + chosen = resolve_adhoc_fetcher("https://example.com/x", sources=[arxiv, intranet]) + assert chosen is arxiv - by_id = resolve_fetcher( - "https://example.com/x", sources=[arxiv, intranet], source_id="intranet" + +# --- resolve_configured_source --- + + +def test_configured_picks_by_source_id(): + """Worker path: pick the source by id so credentials/headers of the + configured source are reused instead of falling to whichever matched + supports(uri) first.""" + arxiv = HTTPSource(source_id="arxiv", headers={"Authorization": "Bearer A"}) + intranet = HTTPSource(source_id="intranet", headers={"Authorization": "Bearer B"}) + assert ( + resolve_configured_source( + "https://example.com/x", "intranet", [arxiv, intranet] + ) + is intranet ) - assert by_id is intranet - - without_id = resolve_fetcher("https://example.com/x", sources=[arxiv, intranet]) - assert without_id is arxiv -def test_source_id_raises_when_matched_source_rejects_uri(): - """source_id selects the source by identity, but that source still - has to support the URI. Mismatch means a stale enqueue from a - reconfigured source — surface it instead of silently picking another.""" +def test_configured_raises_when_matched_source_rejects_uri(): + """source_id selects by identity, but the source still has to support + the URI. Mismatch = stale enqueue from a reconfigured source.""" fs = FSSource(root=Path("/tmp"), source_id="local") with pytest.raises(ValueError, match="doesn't support URI"): - resolve_fetcher("https://example.com/x", sources=[fs], source_id="local") + resolve_configured_source("https://example.com/x", "local", [fs]) -def test_source_id_raises_when_no_source_matches(): - """A worker job carries source_id from when it was enqueued. If the - operator renamed/removed that source, falling through to an ad-hoc - unauthenticated adapter would silently drop credentials. Raise instead - so the job DLQs and the misconfiguration is visible.""" +def test_configured_raises_when_no_source_matches(): + """Worker job carries a source_id from when it was enqueued. If the + operator renamed/removed that source, falling back to an ad-hoc adapter + would silently drop credentials. Raise instead so the job DLQs.""" arxiv = HTTPSource(source_id="arxiv", headers={"Authorization": "Bearer A"}) with pytest.raises(ValueError, match="No configured source with id 'gone'"): - resolve_fetcher("https://example.com/x", sources=[arxiv], source_id="gone") + resolve_configured_source("https://example.com/x", "gone", [arxiv]) -def test_source_id_with_no_sources_list_still_raises(): - """Same strict mode when sources is None: source_id is a strong claim - that must match a configured source.""" +def test_configured_raises_when_sources_list_is_none(): with pytest.raises(ValueError, match="No configured source with id"): - resolve_fetcher("https://example.com/x", source_id="arxiv") + resolve_configured_source("https://example.com/x", "arxiv", None) diff --git a/tests/ingester/test_s3_source.py b/tests/ingester/test_s3_source.py index 44674fb0..469a0600 100644 --- a/tests/ingester/test_s3_source.py +++ b/tests/ingester/test_s3_source.py @@ -166,7 +166,7 @@ async def test_discover_unchanged_against_matching_snapshot(fake_s3_listing): async def test_discover_emits_delete_for_missing_keys(fake_s3_listing): fake_s3_listing([[_meta("file1.md", "abc")]]) src = S3Source(uri="s3://bucket/", supported_extensions=[".md"]) - events = [e async for e in src.discover(since={"s3://bucket/gone.md": "old"})] + events = [e async for e in src.discover(known_uris={"s3://bucket/gone.md"})] deletes = [e for e in events if e.kind is SourceEventKind.DELETE] assert len(deletes) == 1 assert deletes[0].uri == "s3://bucket/gone.md" diff --git a/tests/ingester/test_webdav_source.py b/tests/ingester/test_webdav_source.py index 443d05cd..682d6e03 100644 --- a/tests/ingester/test_webdav_source.py +++ b/tests/ingester/test_webdav_source.py @@ -229,11 +229,12 @@ async def test_discover_emits_delete_for_files_no_longer_listed(): base_url="https://nc.example.com/dav/", transport=_transport(handler), ) - snapshot = { - "https://nc.example.com/dav/a.md": "rev-a", - "https://nc.example.com/dav/gone.md": "rev-old", + revisions = {"https://nc.example.com/dav/a.md": "rev-a"} + known = { + "https://nc.example.com/dav/a.md", + "https://nc.example.com/dav/gone.md", } - events = [event async for event in src.discover(since=snapshot)] + events = [event async for event in src.discover(since=revisions, known_uris=known)] kinds = {e.uri: e.kind for e in events} assert kinds == { "https://nc.example.com/dav/a.md": SourceEventKind.UNCHANGED, diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 1e7fd26f..1d32b35b 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -80,7 +80,7 @@ async def test_drain_marks_job_succeeded_and_writes_sync_state(client, jobs, syn assert refreshed is not None assert refreshed.status is JobStatus.SUCCEEDED - snapshot = await sync.get_snapshot("src") + snapshot = await sync.get_revision_snapshot("src") assert snapshot == {"s3://b/k.md": "e1"} @@ -101,7 +101,7 @@ async def test_drain_delete_op_removes_sync_state(client, jobs, sync): assert refreshed.status is JobStatus.SUCCEEDED client.delete_document.assert_awaited_once_with("doc-1") - snapshot = await sync.get_snapshot("src") + snapshot = await sync.get_revision_snapshot("src") assert snapshot == {} @@ -119,7 +119,7 @@ async def test_permanent_error_marks_dead_no_reschedule(client, jobs, sync): assert refreshed.status is JobStatus.DEAD assert refreshed.last_error == "unsupported" # sync_state is NOT written on failure - assert await sync.get_snapshot("src") == {} + assert await sync.get_revision_snapshot("src") == {} @pytest.mark.asyncio @@ -343,7 +343,7 @@ async def test_worker_loses_claim_to_reaper_does_not_write_sync_state( assert refreshed.status is JobStatus.CLAIMED assert refreshed.claimed_by == "worker-B" # And sync_state must be untouched. - assert await sync.get_snapshot("src") == {} + assert await sync.get_revision_snapshot("src") == {} @pytest.mark.asyncio