From d7fdd61feda21c1f2d0b221ce59030d754d1bda2 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 27 May 2026 13:36:03 +0300 Subject: [PATCH] Resolve worker source by source_id, not just supports(uri) --- haiku_rag_slim/haiku/rag/client/__init__.py | 2 ++ haiku_rag_slim/haiku/rag/client/documents.py | 6 ++++- .../haiku/rag/ingester/sources/registry.py | 26 ++++++++++++++----- .../haiku/rag/ingester/workers/pipeline.py | 1 + tests/ingester/test_pipeline.py | 6 ++--- tests/ingester/test_resolve_fetcher.py | 26 +++++++++++++++++++ 6 files changed, 56 insertions(+), 11 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 4cadb1ee..3946d4b3 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -216,6 +216,7 @@ class HaikuRAG: uri: str | None = None, storage_options: dict[str, str] | None = None, sources: "list[Source] | None" = None, + source_id: str | None = None, ) -> Document | list[Document]: from haiku.rag.client.documents import create_document_from_source @@ -227,6 +228,7 @@ class HaikuRAG: uri=uri, storage_options=storage_options, sources=sources, + source_id=source_id, ) async def update_document( diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index c51face9..d8611c46 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -308,6 +308,7 @@ async def create_document_from_source( uri: str | None = None, storage_options: dict[str, str] | None = None, sources: "list[Source] | None" = None, + source_id: str | None = None, ) -> Document | list[Document]: """Create or update document(s) from a file path, directory, or URL. @@ -375,7 +376,10 @@ async def create_document_from_source( # `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, storage_options=storage_options + source_str, + sources=sources, + source_id=source_id, + storage_options=storage_options, ) # The stored URI is what we look up + persist by. For an explicit uri diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/registry.py b/haiku_rag_slim/haiku/rag/ingester/sources/registry.py index 6ed14b97..707830b2 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/registry.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/registry.py @@ -13,19 +13,31 @@ def resolve_fetcher( uri: str, sources: Iterable[Source] | None = None, *, + source_id: str | None = None, storage_options: dict[str, str] | None = None, ) -> Source: """Pick a Source adapter for ``uri``. - Configured ``sources`` win — the first whose ``supports(uri)`` returns True - is returned. Without a configured match, an ad-hoc adapter is built from - the URI scheme so one-shot calls (``add-src ``) work without any - configuration. + 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. """ if sources: - for src in sources: - if src.supports(uri): - return src + if source_id is not None: + for src in sources: + 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 + else: + for src in sources: + if src.supports(uri): + return src scheme = urlparse(uri).scheme if scheme in ("", "file"): diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index 45b9babf..41229915 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -103,6 +103,7 @@ async def run_job( result = await client.create_document_from_source( job.uri, sources=sources, + source_id=job.source_id, ) # Directory ingestion returns list[Document] — workers ingest single # resources, so a list here is a programming error in the caller. diff --git a/tests/ingester/test_pipeline.py b/tests/ingester/test_pipeline.py index 0ac1f0d1..76a6b66a 100644 --- a/tests/ingester/test_pipeline.py +++ b/tests/ingester/test_pipeline.py @@ -58,14 +58,14 @@ async def test_upsert_calls_create_document_from_source_and_returns_metadata(): assert result.content_hash == "abcd" assert result.deleted is False client.create_document_from_source.assert_awaited_once_with( - "https://example.com/a.pdf", sources=None + "https://example.com/a.pdf", sources=None, source_id="src" ) @pytest.mark.asyncio async def test_upsert_threads_configured_sources_to_client(): """The list of configured Source adapters reaches the client so - resolve_fetcher can pick the authenticated one over an adhoc adapter.""" + resolve_fetcher can pick the right one by source_id.""" from haiku.rag.ingester.sources.http import HTTPSource client = _mock_client() @@ -76,7 +76,7 @@ async def test_upsert_threads_configured_sources_to_client(): await run_job(client, _job(), sources=[configured]) client.create_document_from_source.assert_awaited_once_with( - "https://example.com/a.pdf", sources=[configured] + "https://example.com/a.pdf", sources=[configured], source_id="src" ) diff --git a/tests/ingester/test_resolve_fetcher.py b/tests/ingester/test_resolve_fetcher.py index 060c9e9f..d162c7a4 100644 --- a/tests/ingester/test_resolve_fetcher.py +++ b/tests/ingester/test_resolve_fetcher.py @@ -57,3 +57,29 @@ def test_configured_source_falls_through_when_no_match(): chosen = resolve_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.""" + arxiv = HTTPSource(source_id="arxiv", headers={"Authorization": "Bearer A"}) + intranet = HTTPSource(source_id="intranet", headers={"Authorization": "Bearer B"}) + + by_id = resolve_fetcher( + "https://example.com/x", sources=[arxiv, intranet], source_id="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.""" + 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")