swap aioboto3 for obstore in the S3 client path
This commit is contained in:
Yiorgis Gozadinos
parent
f87d894b01
commit
ef82e6eb08
8 changed files with 247 additions and 385 deletions
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in a new issue