add HTTP and S3 source adapters + resolve_fetcher
This commit is contained in:
parent
d0a730ef60
commit
b9637cd625
7 changed files with 827 additions and 0 deletions
|
|
@ -7,13 +7,19 @@ from haiku.rag.ingester.sources.base import (
|
||||||
)
|
)
|
||||||
from haiku.rag.ingester.sources.filter import FileFilter
|
from haiku.rag.ingester.sources.filter import FileFilter
|
||||||
from haiku.rag.ingester.sources.fs import FSSource
|
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.s3 import S3Source
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"FetchResult",
|
"FetchResult",
|
||||||
"FileFilter",
|
"FileFilter",
|
||||||
"FSSource",
|
"FSSource",
|
||||||
|
"HTTPSource",
|
||||||
"RevisionSnapshot",
|
"RevisionSnapshot",
|
||||||
|
"S3Source",
|
||||||
"Source",
|
"Source",
|
||||||
"SourceEvent",
|
"SourceEvent",
|
||||||
"SourceEventKind",
|
"SourceEventKind",
|
||||||
|
"resolve_fetcher",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
134
haiku_rag_slim/haiku/rag/ingester/sources/http.py
Normal file
134
haiku_rag_slim/haiku/rag/ingester/sources/http.py
Normal file
|
|
@ -0,0 +1,134 @@
|
||||||
|
import hashlib
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources.base import (
|
||||||
|
FetchResult,
|
||||||
|
RevisionSnapshot,
|
||||||
|
SourceEvent,
|
||||||
|
SourceEventKind,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_revision(headers: httpx.Headers) -> tuple[str | None, dict[str, str]]:
|
||||||
|
extra: dict[str, str] = {}
|
||||||
|
etag = (headers.get("etag") or "").strip('"').strip()
|
||||||
|
last_modified = (headers.get("last-modified") or "").strip()
|
||||||
|
if etag:
|
||||||
|
extra["etag"] = etag
|
||||||
|
if last_modified:
|
||||||
|
extra["last_modified"] = last_modified
|
||||||
|
# Prefer ETag — it's a stronger validator. Fall back to Last-Modified.
|
||||||
|
revision = etag or last_modified or None
|
||||||
|
return revision, extra
|
||||||
|
|
||||||
|
|
||||||
|
class HTTPSource:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
source_id: str,
|
||||||
|
urls: list[str] | None = None,
|
||||||
|
headers: dict[str, str] | None = None,
|
||||||
|
transport: httpx.AsyncBaseTransport | None = None,
|
||||||
|
) -> None:
|
||||||
|
self.source_id = source_id
|
||||||
|
self.urls = list(urls or [])
|
||||||
|
self.headers = dict(headers or {})
|
||||||
|
# transport is for testing — production callers leave it None and httpx
|
||||||
|
# uses its real transport.
|
||||||
|
self._transport = transport
|
||||||
|
|
||||||
|
def supports(self, uri: str) -> bool:
|
||||||
|
return urlparse(uri).scheme in ("http", "https")
|
||||||
|
|
||||||
|
def _client(self) -> httpx.AsyncClient:
|
||||||
|
return httpx.AsyncClient(headers=self.headers, transport=self._transport)
|
||||||
|
|
||||||
|
async def fetch(self, uri: str) -> FetchResult:
|
||||||
|
async with self._client() as http:
|
||||||
|
response = await http.get(uri)
|
||||||
|
response.raise_for_status()
|
||||||
|
body = response.content
|
||||||
|
content_type = (
|
||||||
|
response.headers.get("content-type", "application/octet-stream")
|
||||||
|
.split(";")[0]
|
||||||
|
.strip()
|
||||||
|
.lower()
|
||||||
|
)
|
||||||
|
revision, extra = _extract_revision(response.headers)
|
||||||
|
return FetchResult(
|
||||||
|
uri=uri,
|
||||||
|
body=body,
|
||||||
|
content_type=content_type,
|
||||||
|
content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(),
|
||||||
|
revision=revision,
|
||||||
|
extra_metadata=extra,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def discover(
|
||||||
|
self, since: RevisionSnapshot | None = None
|
||||||
|
) -> AsyncIterator[SourceEvent]:
|
||||||
|
# HTTP has no listing concept — discover() only reports on what is
|
||||||
|
# currently configured in self.urls. Config drift (URIs that were
|
||||||
|
# configured before but aren't now) is not visible here; the poller
|
||||||
|
# layer detects that by diffing self.urls across sweeps.
|
||||||
|
#
|
||||||
|
# 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] = dict(since) if since else {}
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
|
||||||
|
async with self._client() as http:
|
||||||
|
for url in self.urls:
|
||||||
|
try:
|
||||||
|
head = await http.head(url)
|
||||||
|
except Exception:
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=url,
|
||||||
|
kind=SourceEventKind.UPSERT,
|
||||||
|
revision=None,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
if head.status_code == 410:
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=url,
|
||||||
|
kind=SourceEventKind.DELETE,
|
||||||
|
revision=None,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
if head.is_error:
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=url,
|
||||||
|
kind=SourceEventKind.UPSERT,
|
||||||
|
revision=None,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
revision, _ = _extract_revision(head.headers)
|
||||||
|
if revision is not None and snapshot.get(url) == revision:
|
||||||
|
kind = SourceEventKind.UNCHANGED
|
||||||
|
else:
|
||||||
|
kind = SourceEventKind.UPSERT
|
||||||
|
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=url,
|
||||||
|
kind=kind,
|
||||||
|
revision=revision,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
42
haiku_rag_slim/haiku/rag/ingester/sources/registry.py
Normal file
42
haiku_rag_slim/haiku/rag/ingester/sources/registry.py
Normal file
|
|
@ -0,0 +1,42 @@
|
||||||
|
from collections.abc import Iterable
|
||||||
|
from pathlib import Path
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources.base import 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
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_fetcher(
|
||||||
|
uri: str,
|
||||||
|
sources: Iterable[Source] | 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 <uri>``) work without any
|
||||||
|
configuration.
|
||||||
|
"""
|
||||||
|
if sources:
|
||||||
|
for src in sources:
|
||||||
|
if src.supports(uri):
|
||||||
|
return src
|
||||||
|
|
||||||
|
scheme = urlparse(uri).scheme
|
||||||
|
if scheme in ("", "file"):
|
||||||
|
# Root only matters for discover(); fetch() needs an absolute path
|
||||||
|
# that already encodes the location, so any root is correct.
|
||||||
|
return FSSource(root=Path("/"))
|
||||||
|
if scheme in ("http", "https"):
|
||||||
|
return HTTPSource(source_id="http:adhoc")
|
||||||
|
if scheme == "s3":
|
||||||
|
bucket = urlparse(uri).netloc
|
||||||
|
if not bucket:
|
||||||
|
raise ValueError(f"Invalid S3 URI: {uri}")
|
||||||
|
return S3Source(uri=f"s3://{bucket}/", storage_options=storage_options)
|
||||||
|
|
||||||
|
raise ValueError(f"No source adapter for URI scheme {scheme!r}: {uri}")
|
||||||
136
haiku_rag_slim/haiku/rag/ingester/sources/s3.py
Normal file
136
haiku_rag_slim/haiku/rag/ingester/sources/s3.py
Normal file
|
|
@ -0,0 +1,136 @@
|
||||||
|
import hashlib
|
||||||
|
import mimetypes
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources.base import (
|
||||||
|
FetchResult,
|
||||||
|
RevisionSnapshot,
|
||||||
|
SourceEvent,
|
||||||
|
SourceEventKind,
|
||||||
|
)
|
||||||
|
from haiku.rag.ingester.sources.filter import (
|
||||||
|
FileFilter,
|
||||||
|
_default_supported_extensions,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_s3_uri(uri: str) -> tuple[str, str]:
|
||||||
|
parsed = urlparse(uri)
|
||||||
|
if parsed.scheme != "s3" or not parsed.netloc:
|
||||||
|
raise ValueError(f"Invalid S3 URI: {uri}")
|
||||||
|
return parsed.netloc, parsed.path.lstrip("/")
|
||||||
|
|
||||||
|
|
||||||
|
class S3Source:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
uri: str,
|
||||||
|
storage_options: dict[str, str] | None = None,
|
||||||
|
ignore_patterns: list[str] | None = None,
|
||||||
|
include_patterns: list[str] | None = None,
|
||||||
|
supported_extensions: list[str] | None = None,
|
||||||
|
source_id: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
self.bucket, self.prefix = _parse_s3_uri(uri)
|
||||||
|
# uri_prefix is the canonical "everything I own" — used by supports()
|
||||||
|
# to scope dispatch, and by discover() to build per-key URIs.
|
||||||
|
self.uri_prefix = f"s3://{self.bucket}/{self.prefix}"
|
||||||
|
self.source_id = source_id or f"s3:{self.bucket}/{self.prefix}"
|
||||||
|
self.storage_options = dict(storage_options or {})
|
||||||
|
self.supported_extensions = (
|
||||||
|
list(supported_extensions)
|
||||||
|
if supported_extensions is not None
|
||||||
|
else _default_supported_extensions()
|
||||||
|
)
|
||||||
|
self.filter = FileFilter(
|
||||||
|
ignore_patterns=ignore_patterns,
|
||||||
|
include_patterns=include_patterns,
|
||||||
|
supported_extensions=self.supported_extensions,
|
||||||
|
)
|
||||||
|
|
||||||
|
def supports(self, uri: str) -> bool:
|
||||||
|
return uri.startswith(self.uri_prefix)
|
||||||
|
|
||||||
|
async def fetch(self, uri: str) -> FetchResult:
|
||||||
|
import obstore # type: ignore[import-not-found]
|
||||||
|
|
||||||
|
from haiku.rag.s3 import make_s3_store
|
||||||
|
|
||||||
|
bucket, key = _parse_s3_uri(uri)
|
||||||
|
store = make_s3_store(bucket, self.storage_options)
|
||||||
|
|
||||||
|
head = await obstore.head_async(store, key)
|
||||||
|
etag = (head.get("e_tag") or "").strip('"').strip() or None
|
||||||
|
|
||||||
|
resp = await obstore.get_async(store, key)
|
||||||
|
body = await resp.bytes_async()
|
||||||
|
# obstore returns a Bytes view; convert to plain bytes so the rest of
|
||||||
|
# the pipeline (hashing, tempfile write) doesn't have to know.
|
||||||
|
body = bytes(body)
|
||||||
|
|
||||||
|
content_type, _ = mimetypes.guess_type(key)
|
||||||
|
if not content_type:
|
||||||
|
content_type = "application/octet-stream"
|
||||||
|
|
||||||
|
extra: dict[str, str] = {}
|
||||||
|
if etag is not None:
|
||||||
|
extra["etag"] = etag
|
||||||
|
|
||||||
|
return FetchResult(
|
||||||
|
uri=uri,
|
||||||
|
body=body,
|
||||||
|
content_type=content_type,
|
||||||
|
content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(),
|
||||||
|
revision=etag,
|
||||||
|
extra_metadata=extra,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def discover(
|
||||||
|
self, since: RevisionSnapshot | None = None
|
||||||
|
) -> AsyncIterator[SourceEvent]:
|
||||||
|
import obstore # type: ignore[import-not-found]
|
||||||
|
|
||||||
|
from haiku.rag.s3 import make_s3_store
|
||||||
|
|
||||||
|
snapshot: dict[str, str] = dict(since) if since else {}
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
seen: set[str] = set()
|
||||||
|
store = make_s3_store(self.bucket, self.storage_options)
|
||||||
|
|
||||||
|
async for batch in obstore.list(store, prefix=self.prefix or None):
|
||||||
|
for obj in batch:
|
||||||
|
key = obj["path"]
|
||||||
|
if not self.filter.include_file(key):
|
||||||
|
continue
|
||||||
|
uri = f"s3://{self.bucket}/{key}"
|
||||||
|
seen.add(uri)
|
||||||
|
revision = (obj.get("e_tag") or "").strip('"').strip() or None
|
||||||
|
|
||||||
|
if revision is not None and snapshot.get(uri) == revision:
|
||||||
|
kind = SourceEventKind.UNCHANGED
|
||||||
|
else:
|
||||||
|
kind = SourceEventKind.UPSERT
|
||||||
|
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=uri,
|
||||||
|
kind=kind,
|
||||||
|
revision=revision,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
|
|
||||||
|
# URIs we previously synced but that no longer appear under the
|
||||||
|
# prefix have been deleted upstream.
|
||||||
|
for uri in snapshot:
|
||||||
|
if uri in seen:
|
||||||
|
continue
|
||||||
|
yield SourceEvent(
|
||||||
|
source_id=self.source_id,
|
||||||
|
uri=uri,
|
||||||
|
kind=SourceEventKind.DELETE,
|
||||||
|
revision=None,
|
||||||
|
discovered_at=now,
|
||||||
|
)
|
||||||
252
tests/ingester/test_http_source.py
Normal file
252
tests/ingester/test_http_source.py
Normal file
|
|
@ -0,0 +1,252 @@
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources.base import SourceEventKind
|
||||||
|
from haiku.rag.ingester.sources.http import HTTPSource
|
||||||
|
|
||||||
|
|
||||||
|
def _transport(routes: dict[tuple[str, str], httpx.Response]) -> httpx.MockTransport:
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
key = (request.method, str(request.url))
|
||||||
|
if key not in routes:
|
||||||
|
return httpx.Response(404)
|
||||||
|
return routes[key]
|
||||||
|
|
||||||
|
return httpx.MockTransport(handler)
|
||||||
|
|
||||||
|
|
||||||
|
def test_supports_http_and_https():
|
||||||
|
src = HTTPSource(source_id="default")
|
||||||
|
assert src.supports("http://example.com/a.pdf")
|
||||||
|
assert src.supports("https://example.com/a.pdf")
|
||||||
|
assert not src.supports("file:///tmp/a.pdf")
|
||||||
|
assert not src.supports("s3://bucket/a.pdf")
|
||||||
|
|
||||||
|
|
||||||
|
def test_source_id_is_user_provided():
|
||||||
|
assert HTTPSource(source_id="arxiv").source_id == "arxiv"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_returns_bytes_and_md5_and_etag():
|
||||||
|
body = b"hello world"
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("GET", "https://example.com/a.md"): httpx.Response(
|
||||||
|
200,
|
||||||
|
content=body,
|
||||||
|
headers={
|
||||||
|
"content-type": "text/markdown",
|
||||||
|
"etag": '"abc123"',
|
||||||
|
"last-modified": "Wed, 21 Oct 2025 07:28:00 GMT",
|
||||||
|
},
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(source_id="default", transport=transport)
|
||||||
|
result = await src.fetch("https://example.com/a.md")
|
||||||
|
assert result.uri == "https://example.com/a.md"
|
||||||
|
assert result.body == body
|
||||||
|
assert result.content_hash == hashlib.md5(body, usedforsecurity=False).hexdigest()
|
||||||
|
assert result.content_type == "text/markdown"
|
||||||
|
# etag preferred over last-modified, surrounding quotes stripped
|
||||||
|
assert result.revision == "abc123"
|
||||||
|
assert result.extra_metadata["etag"] == "abc123"
|
||||||
|
assert result.extra_metadata["last_modified"] == "Wed, 21 Oct 2025 07:28:00 GMT"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_falls_back_to_last_modified_when_no_etag():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("GET", "https://example.com/a"): httpx.Response(
|
||||||
|
200,
|
||||||
|
content=b"x",
|
||||||
|
headers={
|
||||||
|
"content-type": "application/pdf",
|
||||||
|
"last-modified": "Wed, 21 Oct 2025 07:28:00 GMT",
|
||||||
|
},
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(source_id="default", transport=transport)
|
||||||
|
result = await src.fetch("https://example.com/a")
|
||||||
|
assert result.revision == "Wed, 21 Oct 2025 07:28:00 GMT"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_no_revision_when_neither_header_present():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("GET", "https://example.com/a"): httpx.Response(
|
||||||
|
200, content=b"x", headers={"content-type": "text/plain"}
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(source_id="default", transport=transport)
|
||||||
|
result = await src.fetch("https://example.com/a")
|
||||||
|
assert result.revision is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_strips_content_type_parameters():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("GET", "https://example.com/a"): httpx.Response(
|
||||||
|
200,
|
||||||
|
content=b"x",
|
||||||
|
headers={"content-type": "text/html; charset=utf-8"},
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(source_id="default", transport=transport)
|
||||||
|
result = await src.fetch("https://example.com/a")
|
||||||
|
assert result.content_type == "text/html"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_raises_on_error_status():
|
||||||
|
src = HTTPSource(source_id="default", transport=_transport({}))
|
||||||
|
with pytest.raises(httpx.HTTPStatusError):
|
||||||
|
await src.fetch("https://example.com/missing")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_sends_configured_headers():
|
||||||
|
seen: dict[str, str] = {}
|
||||||
|
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
seen.update(request.headers)
|
||||||
|
return httpx.Response(200, content=b"x", headers={"content-type": "text/plain"})
|
||||||
|
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="default",
|
||||||
|
headers={"Authorization": "Bearer abc"},
|
||||||
|
transport=httpx.MockTransport(handler),
|
||||||
|
)
|
||||||
|
await src.fetch("https://example.com/a")
|
||||||
|
assert seen.get("authorization") == "Bearer abc"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_empty_when_no_urls_configured():
|
||||||
|
src = HTTPSource(source_id="default")
|
||||||
|
assert [e async for e in src.discover()] == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_yields_upsert_for_each_configured_url():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("HEAD", "https://example.com/a.md"): httpx.Response(
|
||||||
|
200, headers={"etag": '"abc"'}
|
||||||
|
),
|
||||||
|
("HEAD", "https://example.com/b.md"): httpx.Response(
|
||||||
|
200, headers={"etag": '"def"'}
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="x",
|
||||||
|
urls=["https://example.com/a.md", "https://example.com/b.md"],
|
||||||
|
transport=transport,
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert {e.uri for e in events} == {
|
||||||
|
"https://example.com/a.md",
|
||||||
|
"https://example.com/b.md",
|
||||||
|
}
|
||||||
|
assert all(e.kind is SourceEventKind.UPSERT for e in events)
|
||||||
|
assert {e.revision for e in events} == {"abc", "def"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_unchanged_against_matching_snapshot():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("HEAD", "https://example.com/a.md"): httpx.Response(
|
||||||
|
200, headers={"etag": '"abc"'}
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="x", urls=["https://example.com/a.md"], transport=transport
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover(since={"https://example.com/a.md": "abc"})]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].kind is SourceEventKind.UNCHANGED
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_emits_delete_on_410_gone():
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("HEAD", "https://example.com/retired.md"): httpx.Response(410),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="x",
|
||||||
|
urls=["https://example.com/retired.md"],
|
||||||
|
transport=transport,
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].kind is SourceEventKind.DELETE
|
||||||
|
assert events[0].uri == "https://example.com/retired.md"
|
||||||
|
assert events[0].revision is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("status", [404, 401, 403, 405, 500, 502, 503])
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_treats_non_410_errors_as_upsert(status):
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("HEAD", "https://example.com/a.md"): httpx.Response(status),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="x", urls=["https://example.com/a.md"], transport=transport
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].kind is SourceEventKind.UPSERT
|
||||||
|
assert events[0].revision is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_treats_network_errors_as_upsert():
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
raise httpx.ConnectError("boom")
|
||||||
|
|
||||||
|
src = HTTPSource(
|
||||||
|
source_id="x",
|
||||||
|
urls=["https://example.com/a.md"],
|
||||||
|
transport=httpx.MockTransport(handler),
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].kind is SourceEventKind.UPSERT
|
||||||
|
assert events[0].revision is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_does_not_emit_delete_for_unconfigured_uri():
|
||||||
|
"""Config drift is the poller's job, not HTTPSource's. A URI that was in
|
||||||
|
the snapshot but is no longer configured must not appear as a DELETE."""
|
||||||
|
transport = _transport(
|
||||||
|
{
|
||||||
|
("HEAD", "https://example.com/a.md"): httpx.Response(
|
||||||
|
200, headers={"etag": '"abc"'}
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
src = HTTPSource(
|
||||||
|
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"})
|
||||||
|
]
|
||||||
|
assert {e.uri for e in events} == {"https://example.com/a.md"}
|
||||||
|
assert all(e.kind is not SourceEventKind.DELETE for e in events)
|
||||||
59
tests/ingester/test_resolve_fetcher.py
Normal file
59
tests/ingester/test_resolve_fetcher.py
Normal file
|
|
@ -0,0 +1,59 @@
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources import resolve_fetcher
|
||||||
|
from haiku.rag.ingester.sources.fs import FSSource
|
||||||
|
from haiku.rag.ingester.sources.http import HTTPSource
|
||||||
|
from haiku.rag.ingester.sources.s3 import S3Source
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolves_fs_for_file_uri():
|
||||||
|
src = resolve_fetcher("file:///tmp/sample.md")
|
||||||
|
assert isinstance(src, FSSource)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolves_fs_for_bare_path():
|
||||||
|
src = resolve_fetcher("/tmp/sample.md")
|
||||||
|
assert isinstance(src, FSSource)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolves_http():
|
||||||
|
src = resolve_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")
|
||||||
|
assert isinstance(src, S3Source)
|
||||||
|
assert src.bucket == "my-bucket"
|
||||||
|
assert src.prefix == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolves_s3_forwards_storage_options():
|
||||||
|
src = resolve_fetcher(
|
||||||
|
"s3://my-bucket/key.pdf",
|
||||||
|
storage_options={"endpoint": "http://seaweed:8333", "allow_http": "true"},
|
||||||
|
)
|
||||||
|
assert isinstance(src, S3Source)
|
||||||
|
assert src.storage_options["endpoint"] == "http://seaweed:8333"
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_scheme_raises():
|
||||||
|
with pytest.raises(ValueError, match="No source adapter"):
|
||||||
|
resolve_fetcher("ftp://example.com/x")
|
||||||
|
|
||||||
|
|
||||||
|
def test_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])
|
||||||
|
assert chosen is fs
|
||||||
|
|
||||||
|
|
||||||
|
def test_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])
|
||||||
|
assert isinstance(chosen, HTTPSource)
|
||||||
|
assert chosen is not fs
|
||||||
198
tests/ingester/test_s3_source.py
Normal file
198
tests/ingester/test_s3_source.py
Normal file
|
|
@ -0,0 +1,198 @@
|
||||||
|
import hashlib
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.ingester.sources.base import SourceEventKind
|
||||||
|
from haiku.rag.ingester.sources.s3 import S3Source
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def fake_obstore_io(monkeypatch):
|
||||||
|
import obstore
|
||||||
|
|
||||||
|
head_async = AsyncMock()
|
||||||
|
get_async = AsyncMock()
|
||||||
|
monkeypatch.setattr(obstore, "head_async", head_async)
|
||||||
|
monkeypatch.setattr(obstore, "get_async", get_async)
|
||||||
|
return head_async, get_async
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def fake_s3_listing(monkeypatch):
|
||||||
|
import obstore
|
||||||
|
|
||||||
|
batches: list[list[dict]] = []
|
||||||
|
|
||||||
|
def list_obs(_store, *_, **__):
|
||||||
|
async def _iter():
|
||||||
|
for batch in batches:
|
||||||
|
yield batch
|
||||||
|
|
||||||
|
return _iter()
|
||||||
|
|
||||||
|
monkeypatch.setattr(obstore, "list", MagicMock(side_effect=list_obs))
|
||||||
|
|
||||||
|
def set_batches(new_batches):
|
||||||
|
batches.clear()
|
||||||
|
batches.extend(new_batches)
|
||||||
|
|
||||||
|
return set_batches
|
||||||
|
|
||||||
|
|
||||||
|
def _meta(path: str, etag: str = "abc") -> dict:
|
||||||
|
return {"path": path, "e_tag": f'"{etag}"', "size": 0, "last_modified": None}
|
||||||
|
|
||||||
|
|
||||||
|
def _get_result(data: bytes) -> MagicMock:
|
||||||
|
result = MagicMock()
|
||||||
|
result.bytes_async = AsyncMock(return_value=data)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def test_supports_s3_only():
|
||||||
|
src = S3Source(uri="s3://bucket/")
|
||||||
|
assert src.supports("s3://bucket/a.md")
|
||||||
|
assert not src.supports("https://example.com/a.md")
|
||||||
|
assert not src.supports("file:///tmp/a.md")
|
||||||
|
|
||||||
|
|
||||||
|
def test_supports_scopes_to_prefix():
|
||||||
|
src = S3Source(uri="s3://bucket/incoming/")
|
||||||
|
assert src.supports("s3://bucket/incoming/a.md")
|
||||||
|
assert not src.supports("s3://bucket/other/a.md")
|
||||||
|
|
||||||
|
|
||||||
|
def test_source_id_uses_bucket_and_prefix():
|
||||||
|
assert S3Source(uri="s3://bucket/prefix/").source_id == "s3:bucket/prefix/"
|
||||||
|
assert S3Source(uri="s3://bucket/").source_id == "s3:bucket/"
|
||||||
|
|
||||||
|
|
||||||
|
def test_invalid_uri_raises():
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
S3Source(uri="https://example.com/")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
S3Source(uri="s3:///key")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_returns_bytes_md5_etag(fake_obstore_io):
|
||||||
|
head_async, get_async = fake_obstore_io
|
||||||
|
body = b"S3 hosted content"
|
||||||
|
head_async.return_value = {"e_tag": '"abc123"'}
|
||||||
|
get_async.return_value = _get_result(body)
|
||||||
|
|
||||||
|
src = S3Source(uri="s3://bucket/")
|
||||||
|
result = await src.fetch("s3://bucket/folder/file.txt")
|
||||||
|
|
||||||
|
assert result.uri == "s3://bucket/folder/file.txt"
|
||||||
|
assert result.body == body
|
||||||
|
assert result.content_hash == hashlib.md5(body, usedforsecurity=False).hexdigest()
|
||||||
|
assert result.content_type == "text/plain"
|
||||||
|
assert result.revision == "abc123"
|
||||||
|
assert result.extra_metadata["etag"] == "abc123"
|
||||||
|
head_async.assert_awaited_once()
|
||||||
|
get_async.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_handles_missing_etag(fake_obstore_io):
|
||||||
|
head_async, get_async = fake_obstore_io
|
||||||
|
head_async.return_value = {}
|
||||||
|
get_async.return_value = _get_result(b"x")
|
||||||
|
|
||||||
|
src = S3Source(uri="s3://bucket/")
|
||||||
|
result = await src.fetch("s3://bucket/file.txt")
|
||||||
|
assert result.revision is None
|
||||||
|
assert "etag" not in result.extra_metadata
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_invalid_uri_raises(fake_obstore_io):
|
||||||
|
src = S3Source(uri="s3://bucket/")
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
await src.fetch("s3:///no-bucket")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_yields_upsert_for_new_keys(fake_s3_listing):
|
||||||
|
fake_s3_listing(
|
||||||
|
[
|
||||||
|
[
|
||||||
|
_meta("file1.md", "abc"),
|
||||||
|
_meta("subfolder/file2.md", "def"),
|
||||||
|
]
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
src = S3Source(uri="s3://bucket/", supported_extensions=[".md"])
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
uris = {e.uri for e in events}
|
||||||
|
assert uris == {
|
||||||
|
"s3://bucket/file1.md",
|
||||||
|
"s3://bucket/subfolder/file2.md",
|
||||||
|
}
|
||||||
|
assert all(e.kind is SourceEventKind.UPSERT for e in events)
|
||||||
|
assert {e.revision for e in events} == {"abc", "def"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_unchanged_against_matching_snapshot(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/file1.md": "abc"})]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].kind is SourceEventKind.UNCHANGED
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
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"})]
|
||||||
|
deletes = [e for e in events if e.kind is SourceEventKind.DELETE]
|
||||||
|
assert len(deletes) == 1
|
||||||
|
assert deletes[0].uri == "s3://bucket/gone.md"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_respects_prefix(fake_s3_listing):
|
||||||
|
fake_s3_listing([[_meta("incoming/file1.md", "abc")]])
|
||||||
|
src = S3Source(uri="s3://bucket/incoming/", supported_extensions=[".md"])
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert len(events) == 1
|
||||||
|
assert events[0].uri == "s3://bucket/incoming/file1.md"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_respects_extension_filter(fake_s3_listing):
|
||||||
|
fake_s3_listing(
|
||||||
|
[
|
||||||
|
[
|
||||||
|
_meta("a.md", "x"),
|
||||||
|
_meta("b.log", "y"),
|
||||||
|
]
|
||||||
|
]
|
||||||
|
)
|
||||||
|
src = S3Source(uri="s3://bucket/", supported_extensions=[".md"])
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert {e.uri for e in events} == {"s3://bucket/a.md"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_discover_respects_ignore_patterns(fake_s3_listing):
|
||||||
|
fake_s3_listing(
|
||||||
|
[
|
||||||
|
[
|
||||||
|
_meta("a.md", "x"),
|
||||||
|
_meta("draft-b.md", "y"),
|
||||||
|
]
|
||||||
|
]
|
||||||
|
)
|
||||||
|
src = S3Source(
|
||||||
|
uri="s3://bucket/",
|
||||||
|
supported_extensions=[".md"],
|
||||||
|
ignore_patterns=["draft-*"],
|
||||||
|
)
|
||||||
|
events = [e async for e in src.discover()]
|
||||||
|
assert {e.uri for e in events} == {"s3://bucket/a.md"}
|
||||||
Loading…
Reference in a new issue