diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index a0b0003a..100bc51a 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -463,15 +463,17 @@ async def _create_or_update_document_from_s3( """Create or update a document from an s3:// URL. Two-stage change detection: - - HeadObject ETag matches stored metadata["etag"] → skip GetObject and re-chunk. + - HEAD ETag matches stored metadata["etag"] → skip GET and re-chunk. - ETag differs but content MD5 matches → refresh etag only, no re-chunk. - Otherwise download, convert, chunk, embed. ``uri`` overrides the s3:// URL as the stored document identifier. """ + import obstore # type: ignore[import-not-found] + from haiku.rag.client.processing import get_extension_from_content_type_or_url from haiku.rag.embeddings import embed_chunks - from haiku.rag.s3 import make_s3_session + from haiku.rag.s3 import make_s3_store metadata = metadata or {} stored_uri = uri if uri is not None else url @@ -485,39 +487,39 @@ async def _create_or_update_document_from_s3( converter = get_converter(client._config) supported_extensions = converter.supported_extensions - session, client_kwargs = make_s3_session(storage_options) + store = make_s3_store(bucket, storage_options) - async with session.client("s3", **client_kwargs) as s3: - head = await s3.head_object(Bucket=bucket, Key=key) - etag = head["ETag"].strip('"') - content_type = (head.get("ContentType") or "").lower() - if not content_type: - content_type = "application/octet-stream" + head = await obstore.head_async(store, key) + etag = (head.get("e_tag") or "").strip('"') - existing_doc = await client.get_document_by_uri(stored_uri) - if existing_doc and existing_doc.metadata.get("etag") == etag: - updated = False - if title is not None and title != existing_doc.title: - existing_doc.title = title - updated = True + existing_doc = await client.get_document_by_uri(stored_uri) + if existing_doc and existing_doc.metadata.get("etag") == etag: + updated = False + if title is not None and title != existing_doc.title: + existing_doc.title = title + updated = True - merged_metadata = {**(existing_doc.metadata or {}), **metadata} - if merged_metadata != existing_doc.metadata: - existing_doc.metadata = merged_metadata - updated = True + merged_metadata = {**(existing_doc.metadata or {}), **metadata} + if merged_metadata != existing_doc.metadata: + existing_doc.metadata = merged_metadata + updated = True - if updated: - return await client.document_repository.update(existing_doc) - return existing_doc + if updated: + return await client.document_repository.update(existing_doc) + return existing_doc - file_extension = get_extension_from_content_type_or_url(url, content_type) - if file_extension not in supported_extensions: - raise ValueError( - f"Unsupported content type/extension: {content_type}/{file_extension}" - ) + content_type, _ = mimetypes.guess_type(key) + if not content_type: + content_type = "application/octet-stream" - get_resp = await s3.get_object(Bucket=bucket, Key=key) - body = await get_resp["Body"].read() + file_extension = get_extension_from_content_type_or_url(url, content_type) + if file_extension not in supported_extensions: + raise ValueError( + f"Unsupported content type/extension: {content_type}/{file_extension}" + ) + + get_resp = await obstore.get_async(store, key) + body = await get_resp.bytes_async() md5_hash = hashlib.md5(body, usedforsecurity=False).hexdigest() diff --git a/haiku_rag_slim/haiku/rag/monitor.py b/haiku_rag_slim/haiku/rag/monitor.py index 7cabdfd4..40e18a5a 100644 --- a/haiku_rag_slim/haiku/rag/monitor.py +++ b/haiku_rag_slim/haiku/rag/monitor.py @@ -234,7 +234,7 @@ class S3Watcher: entry: S3MonitorEntry, supported_extensions: list[str], ) -> None: - from haiku.rag.s3 import make_s3_session + from haiku.rag.s3 import make_s3_store parsed = urlparse(entry.uri) if not parsed.netloc: @@ -245,7 +245,7 @@ class S3Watcher: self.bucket = parsed.netloc self.prefix = parsed.path.lstrip("/") self.uri_prefix = f"s3://{self.bucket}/{self.prefix}" - self._make_s3_session = make_s3_session + self._make_s3_store = make_s3_store self.filter = FileFilter( ignore_patterns=entry.ignore_patterns or None, include_patterns=entry.include_patterns or None, @@ -265,20 +265,18 @@ class S3Watcher: logger.error(f"S3 watcher refresh failed for {self.entry.uri}: {e}") async def refresh(self) -> None: - uris_seen: dict[str, str] = {} - session, client_kwargs = self._make_s3_session(self.entry.storage_options) + import obstore # type: ignore[import-not-found] - async with session.client("s3", **client_kwargs) as s3: - paginator = s3.get_paginator("list_objects_v2") - async for page in paginator.paginate( - Bucket=self.bucket, Prefix=self.prefix - ): - for obj in page.get("Contents", []): - key = obj["Key"] - if not self.filter.include_file(key): - continue - uri = f"s3://{self.bucket}/{key}" - uris_seen[uri] = obj["ETag"].strip('"') + uris_seen: dict[str, str] = {} + store = self._make_s3_store(self.bucket, self.entry.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}" + uris_seen[uri] = (obj.get("e_tag") or "").strip('"') existing_etags = await self._existing_etags_under_prefix() diff --git a/haiku_rag_slim/haiku/rag/s3.py b/haiku_rag_slim/haiku/rag/s3.py index 9c9d489f..6629a7e6 100644 --- a/haiku_rag_slim/haiku/rag/s3.py +++ b/haiku_rag_slim/haiku/rag/s3.py @@ -1,42 +1,39 @@ from typing import Any -def make_s3_session( - storage_options: dict[str, str] | None, -) -> tuple[Any, dict[str, Any]]: - """Build an aioboto3 Session and S3 client kwargs from LanceDB-style options. +def make_s3_store(bucket: str, storage_options: dict[str, str] | None) -> Any: + """Build an obstore `S3Store` from LanceDB-style storage_options. - Accepts the same dict shape as `LanceDBConfig.storage_options` so users can - copy-paste their LanceDB credentials. Recognized keys: aws_access_key_id, - aws_secret_access_key, aws_session_token, region (or region_name), endpoint - (or endpoint_url), allow_http. Empty/missing keys fall back to the boto3 - default credential chain (environment variables, IAM role, AWS profile). + Accepts the same dict shape as `LanceDBConfig.storage_options` (the same + Rust `object_store` crate is used by both LanceDB and obstore, so the keys + line up). Recognized keys include aws_access_key_id, aws_secret_access_key, + aws_session_token, region (or aws_region), endpoint (or aws_endpoint), + allow_http. Empty/missing keys fall back to the AWS default credential + chain (environment variables, IAM role, AWS profile). """ try: - import aioboto3 # type: ignore[import-not-found] # ty: ignore[unresolved-import] + from obstore.store import ( + S3Store, # type: ignore[import-not-found] + ) except ImportError as e: raise ImportError( - "aioboto3 is required for s3:// sources. " + "obstore is required for s3:// sources. " "Install with: pip install haiku.rag-slim[s3]" ) from e - storage_options = storage_options or {} - session_kwargs: dict[str, Any] = {} - client_kwargs: dict[str, Any] = {} + options = dict(storage_options or {}) + allow_http = str(options.pop("allow_http", "")).lower() == "true" - for key in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token"): - if key in storage_options: - session_kwargs[key] = storage_options[key] + # Custom endpoints (SeaweedFS, MinIO, etc.) need path-style requests; AWS + # accepts virtual-hosted-style by default. + has_custom_endpoint = bool(options.get("endpoint") or options.get("aws_endpoint")) - region = storage_options.get("region") or storage_options.get("region_name") - if region: - session_kwargs["region_name"] = region + kwargs: dict[str, Any] = {} + if options: + kwargs["config"] = options + if allow_http: + kwargs["client_options"] = {"allow_http": True} + if has_custom_endpoint: + kwargs["virtual_hosted_style_request"] = False - endpoint = storage_options.get("endpoint") or storage_options.get("endpoint_url") - if endpoint: - client_kwargs["endpoint_url"] = endpoint - - if str(storage_options.get("allow_http", "")).lower() == "true": - client_kwargs["use_ssl"] = False - - return aioboto3.Session(**session_kwargs), client_kwargs + return S3Store(bucket, **kwargs) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index a8d3789d..83b83af9 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -44,7 +44,7 @@ dependencies = [ # Document processing docling = ["docling>=2.91.0", "opencv-python-headless>=4.13.0.92"] # S3 / object-storage monitoring -s3 = ["aioboto3>=13"] +s3 = ["obstore>=0.9,<0.10"] # Embedding providers voyageai = ["pydantic-ai-slim[voyageai]"] # Rerankers diff --git a/pyproject.toml b/pyproject.toml index 2539c311..26755586 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,14 +58,6 @@ exclude = [ [tool.hatch.build.targets.wheel] only-include = ["/README.md"] -[tool.uv] -conflicts = [ - [ - { package = "haiku.rag-slim", extra = "s3" }, - { package = "pydantic-ai-slim", extra = "bedrock" }, - ], -] - [tool.uv.workspace] members = ["haiku_rag_slim", "evaluations"] @@ -77,6 +69,7 @@ members = ["haiku_rag_slim", "evaluations"] dev = [ "haiku.rag-evals", "datasets>=4.8.4", + "obstore>=0.9,<0.10", "mkdocs>=1.6.1", "mkdocs-material>=9.7.6", "pre-commit>=4.5.1", diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index 9dfe90be..156fe872 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -14,7 +14,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig, LanceDBConfig, S3MonitorEntry from haiku.rag.store.engine import Store -HAS_AIOBOTO3 = importlib.util.find_spec("aioboto3") is not None +HAS_OBSTORE = importlib.util.find_spec("obstore") is not None S3_ENDPOINT = "http://localhost:8333" S3_BUCKET = "test-bucket" @@ -161,42 +161,28 @@ async def test_app_info_empty_db(tmp_path, capsys): # local — these tests verify the watcher path, not LanceDB-on-S3. -_aioboto3_required = pytest.mark.skipif( - not HAS_AIOBOTO3, - reason="aioboto3 not installed (uv sync --extra s3)", +_obstore_required = pytest.mark.skipif( + not HAS_OBSTORE, + reason="obstore not installed (uv sync --extra s3)", ) -async def _put_object(prefix: str, key: str, body: bytes) -> None: - import aioboto3 # type: ignore[import-not-found] # ty: ignore[unresolved-import] +def _watcher_store(): + from haiku.rag.s3 import make_s3_store - session = aioboto3.Session( - aws_access_key_id=S3_STORAGE_OPTIONS["aws_access_key_id"], - aws_secret_access_key=S3_STORAGE_OPTIONS["aws_secret_access_key"], - region_name=S3_STORAGE_OPTIONS["region"], - ) - async with session.client( - "s3", endpoint_url=S3_STORAGE_OPTIONS["endpoint"], use_ssl=False - ) as s3: - try: - await s3.create_bucket(Bucket=S3_BUCKET) - except Exception: - pass # bucket already exists - await s3.put_object(Bucket=S3_BUCKET, Key=f"{prefix}/{key}", Body=body) + return make_s3_store(S3_BUCKET, S3_STORAGE_OPTIONS) + + +async def _put_object(prefix: str, key: str, body: bytes) -> None: + import obstore + + await obstore.put_async(_watcher_store(), f"{prefix}/{key}", body) async def _delete_object(prefix: str, key: str) -> None: - import aioboto3 # type: ignore[import-not-found] # ty: ignore[unresolved-import] + import obstore - session = aioboto3.Session( - aws_access_key_id=S3_STORAGE_OPTIONS["aws_access_key_id"], - aws_secret_access_key=S3_STORAGE_OPTIONS["aws_secret_access_key"], - region_name=S3_STORAGE_OPTIONS["region"], - ) - async with session.client( - "s3", endpoint_url=S3_STORAGE_OPTIONS["endpoint"], use_ssl=False - ) as s3: - await s3.delete_object(Bucket=S3_BUCKET, Key=f"{prefix}/{key}") + await obstore.delete_async(_watcher_store(), f"{prefix}/{key}") def _watcher_entry(prefix: str, **overrides) -> S3MonitorEntry: @@ -210,7 +196,7 @@ def _watcher_entry(prefix: str, **overrides) -> S3MonitorEntry: ) -@_aioboto3_required +@_obstore_required @pytest.mark.asyncio async def test_s3_watcher_initial_sweep(tmp_path): from haiku.rag.monitor import S3Watcher @@ -236,7 +222,7 @@ async def test_s3_watcher_initial_sweep(tmp_path): ] -@_aioboto3_required +@_obstore_required @pytest.mark.asyncio async def test_s3_watcher_detects_new_object(tmp_path): from haiku.rag.monitor import S3Watcher @@ -258,7 +244,7 @@ async def test_s3_watcher_detects_new_object(tmp_path): assert await rag.count_documents() == 2 -@_aioboto3_required +@_obstore_required @pytest.mark.asyncio async def test_s3_watcher_detects_modified_object(tmp_path): from haiku.rag.monitor import S3Watcher @@ -290,7 +276,7 @@ async def test_s3_watcher_detects_modified_object(tmp_path): assert "new content body" in second.content -@_aioboto3_required +@_obstore_required @pytest.mark.asyncio async def test_s3_watcher_orphan_deletion(tmp_path): from haiku.rag.monitor import S3Watcher diff --git a/tests/test_s3_monitor.py b/tests/test_s3_monitor.py index 0f946b1f..b125c592 100644 --- a/tests/test_s3_monitor.py +++ b/tests/test_s3_monitor.py @@ -1,5 +1,4 @@ import asyncio -import sys from unittest.mock import AsyncMock, MagicMock import pytest @@ -10,59 +9,53 @@ from haiku.rag.store.models.document import Document @pytest.fixture -def fake_aioboto3(monkeypatch): - fake = MagicMock() - monkeypatch.setitem(sys.modules, "aioboto3", fake) - return fake +def s3_listing(monkeypatch): + """Patch `obstore.list_obs` with an async-iterator returning controllable batches. + + Returns `(set_batches, list_mock)`. `set_batches([[meta, ...], ...])` + seeds the next call's pages. + """ + import obstore + + batches: list[list[MagicMock]] = [] + + def list_obs(_store, *_, **__): + async def _iter(): + for batch in batches: + yield batch + + return _iter() + + list_mock = MagicMock(side_effect=list_obs) + monkeypatch.setattr(obstore, "list", list_mock) + + def set_batches(new_batches): + batches.clear() + batches.extend(new_batches) + + return set_batches, list_mock -@pytest.fixture -def s3_paginator(): - """Return (paginator_mock, set_pages) — `set_pages([page, ...])` rewires the iterator.""" - pages: list[dict] = [] - - async def _paginate(**_kwargs): - for page in pages: - yield page - - paginator = MagicMock() - paginator.paginate.side_effect = lambda **kw: _paginate(**kw) - - def set_pages(new_pages): - pages.clear() - pages.extend(new_pages) - - return paginator, set_pages - - -@pytest.fixture -def fake_s3_client(fake_aioboto3, s3_paginator): - paginator, set_pages = s3_paginator - s3_client = MagicMock() - s3_client.get_paginator.return_value = paginator - - client_ctx = AsyncMock() - client_ctx.__aenter__.return_value = s3_client - client_ctx.__aexit__.return_value = None - - session = MagicMock() - session.client.return_value = client_ctx - fake_aioboto3.Session.return_value = session - - return s3_client, set_pages - - -def _entry(**kwargs): - base = { - "uri": "s3://my-bucket/incoming/", - "poll_interval": 60, - "delete_orphans": False, - "ignore_patterns": [], - "include_patterns": [], - "storage_options": {}, +def _meta(path: str, etag: str) -> dict: + # Real obstore ObjectMeta is a TypedDict; raw S3 ETags include quotes. + return { + "path": path, + "e_tag": f'"{etag}"', + "size": 0, + "last_modified": None, } - base.update(kwargs) - return S3MonitorEntry(**base) + + +def _entry(**kwargs) -> S3MonitorEntry: + return S3MonitorEntry( + uri=kwargs.pop("uri", "s3://my-bucket/incoming/"), + poll_interval=kwargs.pop("poll_interval", 60), + delete_orphans=kwargs.pop("delete_orphans", False), + ignore_patterns=kwargs.pop("ignore_patterns", []), + include_patterns=kwargs.pop("include_patterns", []), + storage_options=kwargs.pop("storage_options", {}), + **kwargs, + ) def _doc(uri: str, etag: str, doc_id: str | None = None) -> Document: @@ -75,18 +68,9 @@ def _doc(uri: str, etag: str, doc_id: str | None = None) -> Document: @pytest.mark.asyncio -async def test_s3_watcher_refresh_upserts_new_objects(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages( - [ - { - "Contents": [ - {"Key": "incoming/a.txt", "ETag": '"abc"'}, - {"Key": "incoming/b.txt", "ETag": '"def"'}, - ] - } - ] - ) +async def test_s3_watcher_refresh_upserts_new_objects(s3_listing): + set_batches, _ = s3_listing + set_batches([[_meta("incoming/a.txt", "abc"), _meta("incoming/b.txt", "def")]]) from haiku.rag.monitor import S3Watcher @@ -110,9 +94,9 @@ async def test_s3_watcher_refresh_upserts_new_objects(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_skips_unchanged_etag(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}]) +async def test_s3_watcher_skips_unchanged_etag(s3_listing): + set_batches, _ = s3_listing + set_batches([[_meta("incoming/a.txt", "abc")]]) from haiku.rag.monitor import S3Watcher @@ -126,9 +110,9 @@ async def test_s3_watcher_skips_unchanged_etag(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_upserts_when_etag_differs(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"new"'}]}]) +async def test_s3_watcher_upserts_when_etag_differs(s3_listing): + set_batches, _ = s3_listing + set_batches([[_meta("incoming/a.txt", "new")]]) from haiku.rag.monitor import S3Watcher @@ -147,9 +131,9 @@ async def test_s3_watcher_upserts_when_etag_differs(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_strips_etag_quotes(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}]) +async def test_s3_watcher_strips_etag_quotes(s3_listing): + set_batches, _ = s3_listing + set_batches([[_meta("incoming/a.txt", "abc")]]) from haiku.rag.monitor import S3Watcher @@ -165,9 +149,9 @@ async def test_s3_watcher_strips_etag_quotes(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_deletes_orphans_when_enabled(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}]) +async def test_s3_watcher_deletes_orphans_when_enabled(s3_listing): + set_batches, _ = s3_listing + set_batches([[_meta("incoming/a.txt", "abc")]]) from haiku.rag.monitor import S3Watcher @@ -189,9 +173,9 @@ async def test_s3_watcher_deletes_orphans_when_enabled(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_does_not_delete_orphans_when_disabled(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages([{"Contents": []}]) +async def test_s3_watcher_does_not_delete_orphans_when_disabled(s3_listing): + set_batches, _ = s3_listing + set_batches([[]]) from haiku.rag.monitor import S3Watcher @@ -211,10 +195,10 @@ async def test_s3_watcher_does_not_delete_orphans_when_disabled(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_orphan_scope_is_per_entry(fake_s3_client): +async def test_s3_watcher_orphan_scope_is_per_entry(s3_listing): """A doc under a different bucket prefix must not be touched.""" - s3, set_pages = fake_s3_client - set_pages([{"Contents": []}]) + set_batches, _ = s3_listing + set_batches([[]]) from haiku.rag.monitor import S3Watcher @@ -234,17 +218,15 @@ async def test_s3_watcher_orphan_scope_is_per_entry(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_applies_include_and_ignore_patterns(fake_s3_client): - s3, set_pages = fake_s3_client - set_pages( +async def test_s3_watcher_applies_include_and_ignore_patterns(s3_listing): + set_batches, _ = s3_listing + set_batches( [ - { - "Contents": [ - {"Key": "incoming/keep.md", "ETag": '"1"'}, - {"Key": "incoming/draft.md", "ETag": '"2"'}, - {"Key": "incoming/skip.txt", "ETag": '"3"'}, - ] - } + [ + _meta("incoming/keep.md", "1"), + _meta("incoming/draft.md", "2"), + _meta("incoming/skip.txt", "3"), + ] ] ) @@ -271,24 +253,27 @@ async def test_s3_watcher_applies_include_and_ignore_patterns(fake_s3_client): @pytest.mark.asyncio -async def test_s3_watcher_observe_survives_transient_list_failure(fake_s3_client): +async def test_s3_watcher_observe_survives_transient_list_failure(s3_listing): """First refresh succeeds; second refresh raises; loop survives and recovers.""" - s3, set_pages = fake_s3_client + set_batches, list_mock = s3_listing - pages_initial = [{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}] - pages_after = [{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}] + pages_initial = [[_meta("incoming/a.txt", "abc")]] + pages_after = [[_meta("incoming/a.txt", "abc")]] - set_pages(pages_initial) paginate_calls = {"n": 0} - async def paginate_side_effect(**_kwargs): + def list_obs_side_effect(_store, *_, **__): paginate_calls["n"] += 1 if paginate_calls["n"] == 2: raise RuntimeError("transient list failure") - for page in pages_after if paginate_calls["n"] > 1 else pages_initial: - yield page - s3.get_paginator.return_value.paginate.side_effect = paginate_side_effect + async def _iter(): + for batch in pages_after if paginate_calls["n"] > 1 else pages_initial: + yield batch + + return _iter() + + list_mock.side_effect = list_obs_side_effect from haiku.rag.monitor import S3Watcher @@ -305,7 +290,6 @@ async def test_s3_watcher_observe_survives_transient_list_failure(fake_s3_client ) task = asyncio.create_task(watcher.observe()) - # Let three iterations run: initial refresh, transient failure, recovery. for _ in range(20): await asyncio.sleep(0) if paginate_calls["n"] >= 3: @@ -334,15 +318,13 @@ async def test_s3_watcher_invalid_uri_rejected(): @pytest.mark.asyncio -async def test_serve_starts_one_s3_task_per_entry(monkeypatch, fake_s3_client): +async def test_serve_starts_one_s3_task_per_entry(monkeypatch, s3_listing): """`serve` wires one S3Watcher task per MonitorConfig.s3 entry.""" from haiku.rag import app as app_module - captured_tasks: list = [] original_create_task = asyncio.create_task def tracking_create_task(coro, *args, **kwargs): - captured_tasks.append(coro) return original_create_task(coro, *args, **kwargs) monkeypatch.setattr(app_module.asyncio, "create_task", tracking_create_task) @@ -370,7 +352,6 @@ async def test_serve_starts_one_s3_task_per_entry(monkeypatch, fake_s3_client): monkeypatch.setattr(app_module.S3Watcher, "observe", fake_s3_observe) - # Provide a dummy supported_extensions to skip docling import. class _Conv: supported_extensions = [".txt"] @@ -387,6 +368,5 @@ async def test_serve_starts_one_s3_task_per_entry(monkeypatch, fake_s3_client): await app.serve(enable_monitor=True, enable_mcp=False) - # Both S3 entries should have triggered observe(), plus the FileWatcher. assert fw_observe_calls["n"] == 1 assert s3_observe_calls["n"] == 2 diff --git a/tests/test_s3_source.py b/tests/test_s3_source.py index c5f0ba7b..061accac 100644 --- a/tests/test_s3_source.py +++ b/tests/test_s3_source.py @@ -7,116 +7,70 @@ from haiku.rag.client import HaikuRAG @pytest.fixture -def fake_aioboto3(monkeypatch): - """Install a fake aioboto3 module in sys.modules. +def fake_obstore_io(monkeypatch): + """Patch `obstore.head_async` and `obstore.get_async` with controllable mocks. - Returns the module itself; callers configure `Session` to control behavior. + Returns a tuple `(head_async, get_async)`. Tests assign `.return_value` to + each to seed responses. Real `obstore.store.S3Store` is left intact — + constructing the store still exercises the storage_options path. """ - fake = MagicMock() - monkeypatch.setitem(sys.modules, "aioboto3", fake) - return fake + 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_client(fake_aioboto3): - """Configure fake_aioboto3.Session to return a controllable S3 client. - - Returns a MagicMock representing the S3 client (head_object, get_object). - Tests override its return values per-case. - """ - s3_client = MagicMock() - s3_client.head_object = AsyncMock() - s3_client.get_object = AsyncMock() - - client_ctx = AsyncMock() - client_ctx.__aenter__.return_value = s3_client - client_ctx.__aexit__.return_value = None - - session = MagicMock() - session.client.return_value = client_ctx - - fake_aioboto3.Session.return_value = session - return s3_client +def _meta(etag: str) -> dict: + return {"e_tag": etag, "path": "ignored", "size": 0, "last_modified": None} -def _streaming_body(data: bytes) -> AsyncMock: - body = AsyncMock() - body.read.return_value = data - return body +def _get_result(data: bytes) -> MagicMock: + result = MagicMock() + result.bytes_async = AsyncMock(return_value=data) + return result -def test_make_s3_session_translates_lancedb_keys(fake_aioboto3): - from haiku.rag.s3 import make_s3_session +def test_make_s3_store_accepts_lancedb_keys(): + from haiku.rag.s3 import make_s3_store - storage_options = { - "endpoint": "http://seaweed:8333", - "region": "us-east-1", - "aws_access_key_id": "AKIA", - "aws_secret_access_key": "secret", - "allow_http": "true", - } - - session, client_kwargs = make_s3_session(storage_options) - - fake_aioboto3.Session.assert_called_once_with( - aws_access_key_id="AKIA", - aws_secret_access_key="secret", - region_name="us-east-1", + store = make_s3_store( + "my-bucket", + { + "endpoint": "http://seaweed:8333", + "region": "us-east-1", + "aws_access_key_id": "AKIA", + "aws_secret_access_key": "secret", + "allow_http": "true", + }, ) - assert client_kwargs == { - "endpoint_url": "http://seaweed:8333", - "use_ssl": False, - } - assert session is fake_aioboto3.Session.return_value + assert store is not None # construction must not raise -def test_make_s3_session_accepts_native_aliases(fake_aioboto3): - from haiku.rag.s3 import make_s3_session +def test_make_s3_store_no_options_uses_default_chain(): + from haiku.rag.s3 import make_s3_store - _, client_kwargs = make_s3_session( - {"region_name": "eu-west-1", "endpoint_url": "https://s3"} - ) - - fake_aioboto3.Session.assert_called_once_with(region_name="eu-west-1") - assert client_kwargs == {"endpoint_url": "https://s3"} + store = make_s3_store("my-bucket", None) + assert store is not None -def test_make_s3_session_no_options_uses_default_chain(fake_aioboto3): - from haiku.rag.s3 import make_s3_session +def test_make_s3_store_missing_obstore_raises_actionable_error(monkeypatch): + monkeypatch.setitem(sys.modules, "obstore.store", None) - session, client_kwargs = make_s3_session(None) - - fake_aioboto3.Session.assert_called_once_with() - assert client_kwargs == {} - - -def test_make_s3_session_allow_http_only_when_truthy(fake_aioboto3): - from haiku.rag.s3 import make_s3_session - - _, ck = make_s3_session({"allow_http": "false"}) - assert "use_ssl" not in ck - - _, ck = make_s3_session({"allow_http": "True"}) - assert ck["use_ssl"] is False - - -def test_make_s3_session_missing_aioboto3_raises_actionable_error(monkeypatch): - monkeypatch.setitem(sys.modules, "aioboto3", None) - - from haiku.rag.s3 import make_s3_session + from haiku.rag.s3 import make_s3_store with pytest.raises(ImportError, match=r"haiku\.rag-slim\[s3\]"): - make_s3_session({}) + make_s3_store("my-bucket", {}) @pytest.mark.asyncio -async def test_create_document_from_s3_new(fake_s3_client, temp_db_path): +async def test_create_document_from_s3_new(fake_obstore_io, temp_db_path): + head_async, get_async = fake_obstore_io text = b"S3 hosted content" - fake_s3_client.head_object.return_value = { - "ETag": '"abc123"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)} + head_async.return_value = _meta('"abc123"') + get_async.return_value = _get_result(text) async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document_from_source("s3://my-bucket/folder/file.txt") @@ -125,92 +79,76 @@ async def test_create_document_from_s3_new(fake_s3_client, temp_db_path): assert doc.metadata["etag"] == "abc123" # quotes stripped assert doc.metadata["md5"] # real content MD5 assert doc.metadata["md5"] != "abc123" - fake_s3_client.head_object.assert_awaited_once() - fake_s3_client.get_object.assert_awaited_once() + head_async.assert_awaited_once() + get_async.assert_awaited_once() @pytest.mark.asyncio async def test_create_document_from_s3_skips_when_etag_unchanged( - fake_s3_client, temp_db_path + fake_obstore_io, temp_db_path ): + head_async, get_async = fake_obstore_io text = b"S3 hosted content" - fake_s3_client.head_object.return_value = { - "ETag": '"abc123"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)} + head_async.return_value = _meta('"abc123"') + get_async.return_value = _get_result(text) async with HaikuRAG(temp_db_path, create=True) as client: first = await client.create_document_from_source("s3://my-bucket/file.txt") - # Second call with the same ETag must not GetObject. - fake_s3_client.get_object.reset_mock() - # Re-arm body so a stray call would still produce something readable. - fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)} + # Second call with the same ETag must not GET. + get_async.reset_mock() + get_async.return_value = _get_result(text) # re-arm just in case second = await client.create_document_from_source("s3://my-bucket/file.txt") assert second.id == first.id - assert fake_s3_client.head_object.await_count == 2 - fake_s3_client.get_object.assert_not_awaited() + assert head_async.await_count == 2 + get_async.assert_not_awaited() @pytest.mark.asyncio async def test_create_document_from_s3_etag_changed_md5_same_skips_rechunk( - fake_s3_client, temp_db_path + fake_obstore_io, temp_db_path ): """Multipart re-upload of same content: etag changes, MD5 doesn't. - Expected: GetObject runs to verify, but no re-chunk; only metadata.etag updates. + Expected: GET runs to verify, but no re-chunk; only metadata.etag updates. """ + head_async, get_async = fake_obstore_io text = b"S3 hosted content" - fake_s3_client.head_object.return_value = { - "ETag": '"abc123"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)} + head_async.return_value = _meta('"abc123"') + get_async.return_value = _get_result(text) async with HaikuRAG(temp_db_path, create=True) as client: first = await client.create_document_from_source("s3://my-bucket/file.txt") original_md5 = first.metadata["md5"] original_updated_at = first.updated_at - # Same bytes, different ETag (simulating a multipart re-upload). - fake_s3_client.head_object.return_value = { - "ETag": '"def456-2"', # multipart-style ETag - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)} + # Same bytes, different ETag (multipart re-upload). + head_async.return_value = _meta('"def456-2"') + get_async.return_value = _get_result(text) second = await client.create_document_from_source("s3://my-bucket/file.txt") assert second.id == first.id - assert second.metadata["md5"] == original_md5 # MD5 unchanged - assert second.metadata["etag"] == "def456-2" # ETag refreshed + assert second.metadata["md5"] == original_md5 + assert second.metadata["etag"] == "def456-2" assert second.updated_at >= original_updated_at - # GetObject ran once (initial create) plus once more for the etag-changed compare. - assert fake_s3_client.get_object.await_count == 2 + assert get_async.await_count == 2 # initial create + etag-changed compare @pytest.mark.asyncio async def test_create_document_from_s3_etag_changed_md5_changed_rechunks( - fake_s3_client, temp_db_path + fake_obstore_io, temp_db_path ): - fake_s3_client.head_object.return_value = { - "ETag": '"abc123"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(b"original text")} + head_async, get_async = fake_obstore_io + head_async.return_value = _meta('"abc123"') + get_async.return_value = _get_result(b"original text") async with HaikuRAG(temp_db_path, create=True) as client: first = await client.create_document_from_source("s3://my-bucket/file.txt") - fake_s3_client.head_object.return_value = { - "ETag": '"new999"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = { - "Body": _streaming_body(b"different text now") - } + head_async.return_value = _meta('"new999"') + get_async.return_value = _get_result(b"different text now") second = await client.create_document_from_source("s3://my-bucket/file.txt") @@ -222,40 +160,8 @@ async def test_create_document_from_s3_etag_changed_md5_changed_rechunks( @pytest.mark.asyncio async def test_create_document_from_s3_rejects_invalid_uri( - fake_s3_client, temp_db_path + fake_obstore_io, temp_db_path ): async with HaikuRAG(temp_db_path, create=True) as client: with pytest.raises(ValueError, match="Invalid S3 URI"): await client.create_document_from_source("s3://only-bucket-no-key") - - -@pytest.mark.asyncio -async def test_create_document_from_s3_passes_storage_options( - fake_aioboto3, fake_s3_client, temp_db_path -): - fake_s3_client.head_object.return_value = { - "ETag": '"abc"', - "ContentType": "text/plain", - } - fake_s3_client.get_object.return_value = {"Body": _streaming_body(b"hello")} - - async with HaikuRAG(temp_db_path, create=True) as client: - await client.create_document_from_source( - "s3://bucket/key.txt", - storage_options={ - "endpoint": "http://seaweed:8333", - "region": "us-east-1", - "allow_http": "true", - "aws_access_key_id": "AKIA", - "aws_secret_access_key": "secret", - }, - ) - - fake_aioboto3.Session.assert_called_with( - aws_access_key_id="AKIA", - aws_secret_access_key="secret", - region_name="us-east-1", - ) - fake_aioboto3.Session.return_value.client.assert_called_with( - "s3", endpoint_url="http://seaweed:8333", use_ssl=False - )