swap aioboto3 for obstore in the S3 client path

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 obstore.head_async(store, key)
-        head = await s3.head_object(Bucket=bucket, Key=key)
+    etag = (head.get("e_tag") or "").strip('"')
-        etag = head["ETag"].strip('"')
-        content_type = (head.get("ContentType") or "").lower()
-        if not content_type:
-            content_type = "application/octet-stream"
 
-        existing_doc = await client.get_document_by_uri(stored_uri)
+    existing_doc = await client.get_document_by_uri(stored_uri)
-        if existing_doc and existing_doc.metadata.get("etag") == etag:
+    if existing_doc and existing_doc.metadata.get("etag") == etag:
-            updated = False
+        updated = False
-            if title is not None and title != existing_doc.title:
+        if title is not None and title != existing_doc.title:
-                existing_doc.title = title
+            existing_doc.title = title
-                updated = True
+            updated = True
 
-            merged_metadata = {**(existing_doc.metadata or {}), **metadata}
+        merged_metadata = {**(existing_doc.metadata or {}), **metadata}
-            if merged_metadata != existing_doc.metadata:
+        if merged_metadata != existing_doc.metadata:
-                existing_doc.metadata = merged_metadata
+            existing_doc.metadata = merged_metadata
-                updated = True
+            updated = True
 
-            if updated:
+        if updated:
-                return await client.document_repository.update(existing_doc)
+            return await client.document_repository.update(existing_doc)
-            return existing_doc
+        return existing_doc
 
-        file_extension = get_extension_from_content_type_or_url(url, content_type)
+    content_type, _ = mimetypes.guess_type(key)
-        if file_extension not in supported_extensions:
+    if not content_type:
-            raise ValueError(
+        content_type = "application/octet-stream"
-                f"Unsupported content type/extension: {content_type}/{file_extension}"
-            )
 
-        get_resp = await s3.get_object(Bucket=bucket, Key=key)
+    file_extension = get_extension_from_content_type_or_url(url, content_type)
-        body = await get_resp["Body"].read()
+    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()
 

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] = {}
+        import obstore  # type: ignore[import-not-found]
-        session, client_kwargs = self._make_s3_session(self.entry.storage_options)
 
-        async with session.client("s3", **client_kwargs) as s3:
+        uris_seen: dict[str, str] = {}
-            paginator = s3.get_paginator("list_objects_v2")
+        store = self._make_s3_store(self.bucket, self.entry.storage_options)
-            async for page in paginator.paginate(
+
-                Bucket=self.bucket, Prefix=self.prefix
+        async for batch in obstore.list(store, prefix=self.prefix or None):
-            ):
+            for obj in batch:
-                for obj in page.get("Contents", []):
+                key = obj["path"]
-                    key = obj["Key"]
+                if not self.filter.include_file(key):
-                    if not self.filter.include_file(key):
+                    continue
-                        continue
+                uri = f"s3://{self.bucket}/{key}"
-                    uri = f"s3://{self.bucket}/{key}"
+                uris_seen[uri] = (obj.get("e_tag") or "").strip('"')
-                    uris_seen[uri] = obj["ETag"].strip('"')
 
         existing_etags = await self._existing_etags_under_prefix()
 

haiku_rag_slim/haiku/rag/s3.py

@@ -1,42 +1,39 @@
 from typing import Any
 
 
-def make_s3_session(
+def make_s3_store(bucket: str, storage_options: dict[str, str] | None) -> Any:
-    storage_options: dict[str, str] | None,
+    """Build an obstore `S3Store` from LanceDB-style storage_options.
-) -> tuple[Any, dict[str, Any]]:
-    """Build an aioboto3 Session and S3 client kwargs from LanceDB-style options.
 
-    Accepts the same dict shape as `LanceDBConfig.storage_options` so users can
+    Accepts the same dict shape as `LanceDBConfig.storage_options` (the same
-    copy-paste their LanceDB credentials. Recognized keys: aws_access_key_id,
+    Rust `object_store` crate is used by both LanceDB and obstore, so the keys
-    aws_secret_access_key, aws_session_token, region (or region_name), endpoint
+    line up). Recognized keys include aws_access_key_id, aws_secret_access_key,
-    (or endpoint_url), allow_http. Empty/missing keys fall back to the boto3
+    aws_session_token, region (or aws_region), endpoint (or aws_endpoint),
-    default credential chain (environment variables, IAM role, AWS profile).
+    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 {}
+    options = dict(storage_options or {})
-    session_kwargs: dict[str, Any] = {}
+    allow_http = str(options.pop("allow_http", "")).lower() == "true"
-    client_kwargs: dict[str, Any] = {}
 
-    for key in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token"):
+    # Custom endpoints (SeaweedFS, MinIO, etc.) need path-style requests; AWS
-        if key in storage_options:
+    # accepts virtual-hosted-style by default.
-            session_kwargs[key] = storage_options[key]
+    has_custom_endpoint = bool(options.get("endpoint") or options.get("aws_endpoint"))
 
-    region = storage_options.get("region") or storage_options.get("region_name")
+    kwargs: dict[str, Any] = {}
-    if region:
+    if options:
-        session_kwargs["region_name"] = region
+        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")
+    return S3Store(bucket, **kwargs)
-    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

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

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",

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(
+_obstore_required = pytest.mark.skipif(
-    not HAS_AIOBOTO3,
+    not HAS_OBSTORE,
-    reason="aioboto3 not installed (uv sync --extra s3)",
+    reason="obstore not installed (uv sync --extra s3)",
 )
 
 
-async def _put_object(prefix: str, key: str, body: bytes) -> None:
+def _watcher_store():
-    import aioboto3  # type: ignore[import-not-found]  # ty: ignore[unresolved-import]
+    from haiku.rag.s3 import make_s3_store
 
-    session = aioboto3.Session(
+    return make_s3_store(S3_BUCKET, S3_STORAGE_OPTIONS)
-        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 def _put_object(prefix: str, key: str, body: bytes) -> None:
-    )
+    import obstore
-    async with session.client(
+
-        "s3", endpoint_url=S3_STORAGE_OPTIONS["endpoint"], use_ssl=False
+    await obstore.put_async(_watcher_store(), f"{prefix}/{key}", body)
-    ) 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)
 
 
 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(
+    await obstore.delete_async(_watcher_store(), f"{prefix}/{key}")
-        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}")
 
 
 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

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):
+def s3_listing(monkeypatch):
-    fake = MagicMock()
+    """Patch `obstore.list_obs` with an async-iterator returning controllable batches.
-    monkeypatch.setitem(sys.modules, "aioboto3", fake)
+
-    return fake
+    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 _meta(path: str, etag: str) -> dict:
-def s3_paginator():
+    # Real obstore ObjectMeta is a TypedDict; raw S3 ETags include quotes.
-    """Return (paginator_mock, set_pages) — `set_pages([page, ...])` rewires the iterator."""
+    return {
-    pages: list[dict] = []
+        "path": path,
-
+        "e_tag": f'"{etag}"',
-    async def _paginate(**_kwargs):
+        "size": 0,
-        for page in pages:
+        "last_modified": None,
-            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": {},
     }
-    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):
+async def test_s3_watcher_refresh_upserts_new_objects(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages(
+    set_batches([[_meta("incoming/a.txt", "abc"), _meta("incoming/b.txt", "def")]])
-        [
-            {
-                "Contents": [
-                    {"Key": "incoming/a.txt", "ETag": '"abc"'},
-                    {"Key": "incoming/b.txt", "ETag": '"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):
+async def test_s3_watcher_skips_unchanged_etag(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
+    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):
+async def test_s3_watcher_upserts_when_etag_differs(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"new"'}]}])
+    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):
+async def test_s3_watcher_strips_etag_quotes(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
+    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):
+async def test_s3_watcher_deletes_orphans_when_enabled(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages([{"Contents": [{"Key": "incoming/a.txt", "ETag": '"abc"'}]}])
+    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):
+async def test_s3_watcher_does_not_delete_orphans_when_disabled(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages([{"Contents": []}])
+    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_batches, _ = s3_listing
-    set_pages([{"Contents": []}])
+    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):
+async def test_s3_watcher_applies_include_and_ignore_patterns(s3_listing):
-    s3, set_pages = fake_s3_client
+    set_batches, _ = s3_listing
-    set_pages(
+    set_batches(
         [
-            {
+            [
-                "Contents": [
+                _meta("incoming/keep.md", "1"),
-                    {"Key": "incoming/keep.md", "ETag": '"1"'},
+                _meta("incoming/draft.md", "2"),
-                    {"Key": "incoming/draft.md", "ETag": '"2"'},
+                _meta("incoming/skip.txt", "3"),
-                    {"Key": "incoming/skip.txt", "ETag": '"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_initial = [[_meta("incoming/a.txt", "abc")]]
-    pages_after = [{"Contents": [{"Key": "incoming/a.txt", "ETag": '"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

tests/test_s3_source.py

@@ -7,116 +7,70 @@ from haiku.rag.client import HaikuRAG
 
 
 @pytest.fixture
-def fake_aioboto3(monkeypatch):
+def fake_obstore_io(monkeypatch):
-    """Install a fake aioboto3 module in sys.modules.
+    """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()
+    import obstore
-    monkeypatch.setitem(sys.modules, "aioboto3", fake)
+
-    return fake
+    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 _meta(etag: str) -> dict:
-def fake_s3_client(fake_aioboto3):
+    return {"e_tag": etag, "path": "ignored", "size": 0, "last_modified": None}
-    """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 _streaming_body(data: bytes) -> AsyncMock:
+def _get_result(data: bytes) -> MagicMock:
-    body = AsyncMock()
+    result = MagicMock()
-    body.read.return_value = data
+    result.bytes_async = AsyncMock(return_value=data)
-    return body
+    return result
 
 
-def test_make_s3_session_translates_lancedb_keys(fake_aioboto3):
+def test_make_s3_store_accepts_lancedb_keys():
-    from haiku.rag.s3 import make_s3_session
+    from haiku.rag.s3 import make_s3_store
 
-    storage_options = {
+    store = make_s3_store(
-        "endpoint": "http://seaweed:8333",
+        "my-bucket",
-        "region": "us-east-1",
+        {
-        "aws_access_key_id": "AKIA",
+            "endpoint": "http://seaweed:8333",
-        "aws_secret_access_key": "secret",
+            "region": "us-east-1",
-        "allow_http": "true",
+            "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",
     )
-    assert client_kwargs == {
+    assert store is not None  # construction must not raise
-        "endpoint_url": "http://seaweed:8333",
-        "use_ssl": False,
-    }
-    assert session is fake_aioboto3.Session.return_value
 
 
-def test_make_s3_session_accepts_native_aliases(fake_aioboto3):
+def test_make_s3_store_no_options_uses_default_chain():
-    from haiku.rag.s3 import make_s3_session
+    from haiku.rag.s3 import make_s3_store
 
-    _, client_kwargs = make_s3_session(
+    store = make_s3_store("my-bucket", None)
-        {"region_name": "eu-west-1", "endpoint_url": "https://s3"}
+    assert store is not None
-    )
-
-    fake_aioboto3.Session.assert_called_once_with(region_name="eu-west-1")
-    assert client_kwargs == {"endpoint_url": "https://s3"}
 
 
-def test_make_s3_session_no_options_uses_default_chain(fake_aioboto3):
+def test_make_s3_store_missing_obstore_raises_actionable_error(monkeypatch):
-    from haiku.rag.s3 import make_s3_session
+    monkeypatch.setitem(sys.modules, "obstore.store", None)
 
-    session, client_kwargs = make_s3_session(None)
+    from haiku.rag.s3 import make_s3_store
-
-    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
 
     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 = {
+    head_async.return_value = _meta('"abc123"')
-        "ETag": '"abc123"',
+    get_async.return_value = _get_result(text)
-        "ContentType": "text/plain",
-    }
-    fake_s3_client.get_object.return_value = {"Body": _streaming_body(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()
+    head_async.assert_awaited_once()
-    fake_s3_client.get_object.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 = {
+    head_async.return_value = _meta('"abc123"')
-        "ETag": '"abc123"',
+    get_async.return_value = _get_result(text)
-        "ContentType": "text/plain",
-    }
-    fake_s3_client.get_object.return_value = {"Body": _streaming_body(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.
+        # Second call with the same ETag must not GET.
-        fake_s3_client.get_object.reset_mock()
+        get_async.reset_mock()
-        # Re-arm body so a stray call would still produce something readable.
+        get_async.return_value = _get_result(text)  # re-arm just in case
-        fake_s3_client.get_object.return_value = {"Body": _streaming_body(text)}
         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
+    assert head_async.await_count == 2
-    fake_s3_client.get_object.assert_not_awaited()
+    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 = {
+    head_async.return_value = _meta('"abc123"')
-        "ETag": '"abc123"',
+    get_async.return_value = _get_result(text)
-        "ContentType": "text/plain",
-    }
-    fake_s3_client.get_object.return_value = {"Body": _streaming_body(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).
+        # Same bytes, different ETag (multipart re-upload).
-        fake_s3_client.head_object.return_value = {
+        head_async.return_value = _meta('"def456-2"')
-            "ETag": '"def456-2"',  # multipart-style ETag
+        get_async.return_value = _get_result(text)
-            "ContentType": "text/plain",
-        }
-        fake_s3_client.get_object.return_value = {"Body": _streaming_body(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["md5"] == original_md5
-    assert second.metadata["etag"] == "def456-2"  # ETag refreshed
+    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 get_async.await_count == 2  # initial create + etag-changed compare
-    assert fake_s3_client.get_object.await_count == 2
 
 
 @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 = {
+    head_async, get_async = fake_obstore_io
-        "ETag": '"abc123"',
+    head_async.return_value = _meta('"abc123"')
-        "ContentType": "text/plain",
+    get_async.return_value = _get_result(b"original text")
-    }
-    fake_s3_client.get_object.return_value = {"Body": _streaming_body(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 = {
+        head_async.return_value = _meta('"new999"')
-            "ETag": '"new999"',
+        get_async.return_value = _get_result(b"different text now")
-            "ContentType": "text/plain",
-        }
-        fake_s3_client.get_object.return_value = {
-            "Body": _streaming_body(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
-    )