diff --git a/CHANGELOG.md b/CHANGELOG.md index 68b8b5c2..a77c7aaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,11 @@ - Searches in one model response deduplicate their results: evidence a sibling search already showed collapses to a reference line, and a picture attaches once per response. +- `Store(location, config)`, `connect_lancedb(location, config)`, + `gather_database_info(location, config)` and `run_doctor(config, location, ...)` + take the database location, a path or a URI. `ConnectionMode.of(location)` + replaces `ConnectionMode.from_config`. `DatabaseRef.connection()` and + `default_db_path` removed. ## [0.81.0] - 2026-09-01 diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 76c82637..a8bb4962 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -1,5 +1,4 @@ import logging -from functools import cached_property from pathlib import Path from typing import TYPE_CHECKING @@ -64,21 +63,10 @@ class HaikuRAGApp: [ref] = self.scope.databases return ref - @cached_property - def _connection(self) -> "tuple[AppConfig, Path]": - """How to open the one database this command works on, directly. - - Derived per database from its configured location. - """ - from haiku.rag.client.session import default_db_path - - config, db_path = self._one.connection(self.config) - return config, db_path or default_db_path(config) - @property - def _store_config(self) -> AppConfig: - """The configuration for opening the one database directly.""" - return self._connection[0] + def _location(self) -> "Path | str": + """Where the one database this command works on is.""" + return self._one.location @property def _is_local(self) -> bool: @@ -91,12 +79,9 @@ class HaikuRAGApp: @property def _path(self) -> Path: - """The path of the one database this command works on. - - A database behind a URI has none of its own, and the default stands in: - the URI in `_store_config` is what decides where it connects. - """ - return self._connection[1] + """The path of the one local database this command works on.""" + assert self._one.db_path is not None + return self._one.db_path @property def display_path(self) -> "Path | str": @@ -140,7 +125,7 @@ class HaikuRAGApp: self.console.print("[red]Database path does not exist.[/red]") return - info = await gather_database_info(self._store_config, self._path) + info = await gather_database_info(self._location, self.config) if not info.exists: self.console.print( @@ -282,8 +267,8 @@ class HaikuRAGApp: cm = status if status is not None else nullcontext() with cm: report = await run_doctor( - self._store_config, - self._path, + self.config, + self._location, dict(os.environ), duplicates_out=duplicates_out, on_progress=on_progress, @@ -340,8 +325,8 @@ class HaikuRAGApp: return async with Store( - self._path, - config=self._store_config, + self._location, + config=self.config, skip_validation=True, read_only=True, skip_migration_check=True, @@ -415,15 +400,15 @@ class HaikuRAGApp: """ from haiku.rag.store.engine import Store - return Store(self._path, config=self._store_config, read_only=self.read_only) + return Store(self._location, config=self.config, read_only=self.read_only) def _tag_read_store(self) -> "Store": """Read-only store for tag inspection; works on old or drifted DBs.""" from haiku.rag.store.engine import Store return Store( - self._path, - config=self._store_config, + self._location, + config=self.config, skip_validation=True, skip_migration_check=True, read_only=True, @@ -760,8 +745,8 @@ class HaikuRAGApp: from haiku.rag.store.engine import Store async with Store( - self._path, - config=self._store_config, + self._location, + config=self.config, skip_validation=True, skip_migration_check=True, read_only=self.read_only, diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index d256543b..aaa3c8ed 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -42,14 +42,8 @@ def run_chat( config.qa.model = model_config config.analysis.model = model_config - # The capabilities read the databases the scope covers, not what the - # configuration names: a `--db PATH` selection is outside the - # configuration, and a `--db-name NAME` selection is narrower than it. - if scope.covers_multiple: - capability_config, capability_db_path = config, None - else: - capability_config, capability_db_path = scope.databases[0].connection(config) - + # The app opens the scope and lends that client to the capabilities, which + # read what `--db PATH` or `--db-name NAME` selected. enabled = capabilities or ["rag"] capability_list = [] defer_loading = len(enabled) > 1 @@ -68,8 +62,7 @@ def run_chat( capability_list.append( create_capability( - db_path=capability_db_path, - config=capability_config, + config=config, defer_loading=defer_loading, vision=driving_model.vision, ) @@ -80,8 +73,7 @@ def run_chat( capability_list.append( create_capability( - db_path=capability_db_path, - config=capability_config, + config=config, defer_loading=defer_loading, vision=driving_model.vision, ) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index ce1b1a52..b5dccbdb 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -148,10 +148,12 @@ class ChatApp(App): # a client whose __aenter__ failed. await client.__aenter__() self.client = client - # Lent to the capabilities: already the databases they were built for, - # and one connection per database however many capabilities read it. + # Lent to the capabilities, with the scope it covers: one connection + # per database however many capabilities read it, and the analysis + # sandbox is built over the same selection. for capability in self._capabilities: capability.borrowed_rag = client + capability.scope = self.scope self._agent = Agent( self._model, diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 3634343e..965d2ff2 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -21,7 +21,6 @@ from haiku.rag.client.session import ( FederatedSession, SingleDatabaseSession, aclose_quietly, - default_db_path, ) from haiku.rag.config import AppConfig, get_config from haiku.rag.converters import get_converter @@ -145,9 +144,6 @@ class HaikuRAG: of nothing to search. """ self._configured = config if config is not None else get_config() - # What the caller configured, kept intact: entering derives a - # single-database configuration from it, and every re-entry derives - # from the configured set. self._config = self._configured self._requested_db_path = Path(db_path) if db_path is not None else None if self._requested_db_path is not None and sources is not None: @@ -339,15 +335,12 @@ class HaikuRAG: return self [ref] = scope.databases - self._config, db_path = ref.connection(self._configured) - self._session = await SingleDatabaseSession( - db_path if db_path is not None else default_db_path(self._config), + ref, self._config, skip_validation=self._skip_validation, create=self._create, read_only=self._read_only, - source=ref.name, ).open() return self diff --git a/haiku_rag_slim/haiku/rag/client/scope.py b/haiku_rag_slim/haiku/rag/client/scope.py index dd739bb5..6c53af75 100644 --- a/haiku_rag_slim/haiku/rag/client/scope.py +++ b/haiku_rag_slim/haiku/rag/client/scope.py @@ -43,15 +43,10 @@ class DatabaseRef: uri, db_path = locate_database(location) return cls(name=name, uri=uri, db_path=db_path) - def connection(self, config: AppConfig) -> tuple[AppConfig, Path | None]: - """The configuration and path to open this one database with. - - A copy: the caller's configuration still names whatever set it named. - """ - one = config.model_copy(deep=True) - one.lancedb.databases = {} - one.lancedb.uri = self.uri - return one, self.db_path + @property + def location(self) -> Path | str: + """Where the database is: its path, or its URI.""" + return self.db_path if self.db_path is not None else self.uri @dataclass(frozen=True) diff --git a/haiku_rag_slim/haiku/rag/client/session.py b/haiku_rag_slim/haiku/rag/client/session.py index b2507146..01dec0b3 100644 --- a/haiku_rag_slim/haiku/rag/client/session.py +++ b/haiku_rag_slim/haiku/rag/client/session.py @@ -43,36 +43,30 @@ async def aclose_quietly(closeable: Any, what: str) -> None: logger.debug("Closing the %s failed on teardown", what, exc_info=True) -def default_db_path(config: AppConfig) -> Path: - """Where a database lives when its location names no path.""" - return config.storage.data_dir / "haiku.rag.lancedb" - - class SingleDatabaseSession: """One database: its store, its repositories, and their lifecycle. Everything that needs a store lives here, so nothing above has to ask whether - it has one. ``source`` is the configured name this database answers to, or - None where nothing names it. + it has one. Built from the resolved reference: ``source`` is the configured + name it answers to, or None where nothing names it, and the store receives + its location. - ``db_path``, ``config``, ``read_only`` and ``source`` are readable: a client + ``ref``, ``config``, ``read_only`` and ``source`` are readable: a client borrowing this session reports them as its own. """ def __init__( self, - db_path: Path | str, + ref: DatabaseRef, config: AppConfig, *, skip_validation: bool = False, create: bool = False, read_only: bool = False, - source: str | None = None, ) -> None: - self.db_path = db_path + self.ref = ref self.config = config self.read_only = read_only - self.source = source self._skip_validation = skip_validation self._create = create self._vacuum_tasks: set[asyncio.Task] = set() @@ -80,19 +74,25 @@ class SingleDatabaseSession: self._vacuum_dirty = False @property - def location(self) -> Path | str: - """Configured URI or local path for this database. + def source(self) -> str | None: + return self.ref.name - Not `db_path`, which is a placeholder where a URI holds the database. - """ - return self.config.lancedb.uri or self.db_path + @property + def location(self) -> Path | str: + """Where this database is: its path, or its URI.""" + return self.ref.location + + @property + def db_path(self) -> Path | None: + """The local path, or None for a database behind a URI.""" + return self.ref.db_path async def open(self) -> "SingleDatabaseSession": """Connect, validate, and build the repositories.""" failure: str | None = None try: self.store = Store( - self.db_path, + self.location, config=self.config, skip_validation=self._skip_validation, create=self._create, @@ -309,14 +309,11 @@ class FederatedSession: Registered here because a cancelled `gather` discards its results. """ - ref = self._refs[name] - one, db_path = ref.connection(self._config) self._sessions[name] = await SingleDatabaseSession( - db_path if db_path is not None else default_db_path(one), - one, + self._refs[name], + self._config, skip_validation=self._skip_validation, read_only=self._read_only, - source=ref.name, ).open() async def aclose(self) -> None: diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 570eceba..40e20222 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -1080,7 +1080,7 @@ async def run_provider_checks( async def run_doctor( config: AppConfig, - db_path: Path, + location: Path | str, environ: dict[str, str], duplicates_out: Path | None = None, on_progress: Callable[[str], None] | None = None, @@ -1092,7 +1092,7 @@ async def run_doctor( """ notify = on_progress or (lambda _label: None) notify("Inspecting tables") - db = await connect_lancedb(config, db_path) + db = await connect_lancedb(location, config) stats = await get_database_stats(db) results: list[CheckResult] = [] @@ -1110,7 +1110,7 @@ async def run_doctor( missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]] if not missing: async with Store( - db_path, + location, config=config, skip_validation=True, read_only=True, diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py index 2f0dc6c2..631d937e 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py @@ -1,6 +1,5 @@ from fastapi import APIRouter, Depends, HTTPException, status -from haiku.rag.client.session import default_db_path from haiku.rag.ingester.api.server import APIState, get_state from haiku.rag.store.info import DatabaseInfo, gather_database_info @@ -21,5 +20,4 @@ async def database(state: APIState = Depends(get_state)) -> DatabaseInfo: detail="database not configured", ) [ref] = state.scope.databases - one, db_path = ref.connection(state.config) - return await gather_database_info(one, db_path or default_db_path(one)) + return await gather_database_info(ref.location, state.config) diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index 136a1850..95244252 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -20,13 +20,12 @@ async def database_lines(client: "HaikuRAG") -> list[str]: Reported through the connection the client already holds. A failure becomes a line of the report, and the other databases still report. """ - from haiku.rag.store.engine import ConnectionMode from haiku.rag.store.info import get_database_stats lines: list[str] = [] db_path = client.store.db_path - if client.store._connection_mode == ConnectionMode.LOCAL and not db_path.exists(): + if db_path is not None and not db_path.exists(): return ["[red]Database path does not exist.[/red]"] try: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 262455ce..85cb5efe 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -38,11 +38,12 @@ class ConnectionMode(Enum): OBJECT_STORAGE = "object_storage" @staticmethod - def from_config(config: AppConfig) -> "ConnectionMode": - uri = config.lancedb.uri - if not uri: + def of(location: Path | str) -> "ConnectionMode": + """How a location is connected to: a path is local, `db://` is LanceDB + Cloud, any other scheme is object storage.""" + if isinstance(location, Path) or "://" not in location: return ConnectionMode.LOCAL - if uri.startswith("db://"): + if location.startswith("db://"): return ConnectionMode.CLOUD return ConnectionMode.OBJECT_STORAGE @@ -72,8 +73,10 @@ def _session(config: AppConfig) -> lancedb.Session: async def connect_lancedb( - config: AppConfig, db_path: Path | None = None + location: Path | str, config: AppConfig ) -> lancedb.AsyncConnection: + """Connect to the database at `location`, with the connection settings + (credentials, storage options, caches, consistency) from `config`.""" interval = config.lancedb.read_consistency_interval_seconds kwargs: dict[str, Any] = { "session": _session(config), @@ -81,22 +84,19 @@ async def connect_lancedb( timedelta(seconds=interval) if interval is not None else None ), } - mode = ConnectionMode.from_config(config) + mode = ConnectionMode.of(location) if mode == ConnectionMode.CLOUD: return await lancedb.connect_async( - uri=config.lancedb.uri, + uri=str(location), api_key=config.lancedb.api_key, region=config.lancedb.region, **kwargs, ) - elif mode == ConnectionMode.OBJECT_STORAGE: + if mode == ConnectionMode.OBJECT_STORAGE: if config.lancedb.storage_options: kwargs["storage_options"] = config.lancedb.storage_options - return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs) - else: - if db_path is None: - raise ValueError("No lancedb.uri configured and no db_path provided") - return await lancedb.connect_async(db_path.absolute(), **kwargs) + return await lancedb.connect_async(uri=str(location), **kwargs) + return await lancedb.connect_async(Path(location).absolute(), **kwargs) def _stored_vector_dim(settings: dict) -> int | None: @@ -180,14 +180,24 @@ class TagInfo: class Store: def __init__( self, - db_path: Path | str, + location: Path | str, config: AppConfig | None = None, skip_validation: bool = False, create: bool = False, read_only: bool = False, skip_migration_check: bool = False, ): - self.db_path: Path = Path(db_path) + """A store over the database at `location`, a local path or a URI. + + `config` supplies connection settings; where the database is comes + from `location` alone. + """ + self._location: Path | str = location + self.db_path: Path | None = ( + Path(location) + if ConnectionMode.of(location) == ConnectionMode.LOCAL + else None + ) self._config = config if config is not None else get_config() self._read_only = read_only self._create = create @@ -200,7 +210,7 @@ class Store: self._rebuild_lock = asyncio.Lock() self._is_new_db = False - if self._connection_mode == ConnectionMode.LOCAL: + if self.db_path is not None: if not self.db_path.exists(): if not create: raise FileNotFoundError( @@ -231,7 +241,7 @@ class Store: async def _initialize(self): """Perform async initialization: connect to LanceDB, init tables, validate.""" self.db: lancedb.AsyncConnection = await connect_lancedb( - self._config, self.db_path + self.location, self._config ) # Read once and thread onward: on object storage each of these is a @@ -392,9 +402,14 @@ class Store: needed = datetime.now() - oldest + TAG_RETENTION_MARGIN return max(retention, needed) + @property + def location(self) -> Path | str: + """Where this store connected: a local path, or a URI.""" + return self._location + @property def _connection_mode(self) -> ConnectionMode: - return ConnectionMode.from_config(self._config) + return ConnectionMode.of(self._location) async def _ensure_vector_index(self) -> None: """Create or rebuild vector index on chunks table. diff --git a/haiku_rag_slim/haiku/rag/store/info.py b/haiku_rag_slim/haiku/rag/store/info.py index 3c7b1dc4..ade5d624 100644 --- a/haiku_rag_slim/haiku/rag/store/info.py +++ b/haiku_rag_slim/haiku/rag/store/info.py @@ -96,15 +96,15 @@ class DatabaseInfo(BaseModel): packages: dict[str, str] = Field(default_factory=dict) -async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo: +async def gather_database_info(location: Path | str, config: AppConfig) -> DatabaseInfo: """Collect read-only database state without going through Store, so a database missing tables (e.g. pre-migration) still reports what it can.""" from haiku.rag.store.upgrades import get_pending_upgrades from haiku.rag.utils import get_package_versions - display_path = config.lancedb.uri or str(db_path) + display_path = str(location) - db = await connect_lancedb(config, db_path) + db = await connect_lancedb(location, config) stats = await get_database_stats(db) if not any(entry["exists"] for entry in stats.values()): diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py index b44251f3..547bde3f 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py @@ -77,8 +77,11 @@ async def _apply_split_document_meta(store: Store) -> None: Exception ): # pragma: no cover - defensive; stats() failure shouldn't block the split live_bytes = 0 - free_bytes = shutil.disk_usage(store.db_path).free - if live_bytes and free_bytes < live_bytes: + # A database behind a URI has no local disk to run out of. + free_bytes = ( + shutil.disk_usage(store.db_path).free if store.db_path is not None else None + ) + if live_bytes and free_bytes is not None and free_bytes < live_bytes: logger.warning( "Skipping post-migration vacuum: need ~%.2f GB free to compact the " "documents table, have %.2f GB. Run `haiku-rag vacuum` once you have " diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index ea5ea9b0..cc5774cb 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -84,17 +84,18 @@ def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch): from haiku.rag.chat import run_chat run_chat(scope=DatabaseScope.resolve(config, database_name="b")) + named_scope = chat_app.call_args.kwargs["scope"] [named] = chat_app.call_args.kwargs["capabilities"] run_chat(scope=DatabaseScope.resolve(config)) + covering_scope = chat_app.call_args.kwargs["scope"] [covering] = chat_app.call_args.kwargs["capabilities"] - # The chat lends its own client, so this scope is the fallback: it places - # the named database alone. - [placed] = named.scope.databases - assert placed.db_path == tmp_path / "b.lancedb" - assert named.config.lancedb.databases == {} - assert covering.scope.names == ("a", "b") + # The app opens the scope it is handed and lends that client to the + # capabilities, which keep the configuration as the caller named it. + assert named_scope.names == ("b",) + assert covering_scope.names == ("a", "b") + assert set(named.config.lancedb.databases) == {"a", "b"} assert set(covering.config.lancedb.databases) == {"a", "b"} @@ -684,6 +685,37 @@ class TestLendingTheClient: assert borrowed == [client] * len(app._capabilities) assert borrowed + @pytest.mark.asyncio + async def test_mounting_gives_every_capability_the_apps_scope(self, tmp_path): + """A capability built over the configured set covers what the chat + selected once mounted: the analysis sandbox is built over that scope.""" + from haiku.rag.chat.app import ChatApp + from haiku.rag.client.scope import DatabaseScope + from haiku.rag.config.models import AppConfig, LanceDBConfig + + config = AppConfig( + lancedb=LanceDBConfig( + databases={ + "a": str(tmp_path / "a.lancedb"), + "b": str(tmp_path / "b.lancedb"), + } + ) + ) + selected = DatabaseScope.resolve(config, database_name="b") + capability = create_capability(config=config) + assert capability.scope.covers_multiple + + client = _make_mock_client() + app = ChatApp(scope=selected, capabilities=[capability], read_only=True) + with ( + patch("haiku.rag.chat.app.HaikuRAG") as stub_rag, + _covering_returns(stub_rag, client), + ): + async with app.run_test(): + pass + + assert capability.scope == selected + class TestDocumentSelectionIdentity: """Two documents can share a title, within a corpus and across databases, so diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index 192af9b8..995910f1 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -7,7 +7,6 @@ import pytest from haiku.rag.capabilities._tools import search_corpus from haiku.rag.capabilities.rag import RAGState, create_capability from haiku.rag.client import HaikuRAG -from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.session import FederatedSession from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.store.exceptions import UnknownDatabaseError @@ -256,12 +255,9 @@ class TestLendingANamedClient: await _seed(config, "alpha", ["alpha document about cats"]) await _seed(config, "beta", ["beta document about cats"]) - # `run_chat` derives these for a single-database scope. - scope = DatabaseScope.resolve(config, database_name="alpha") - one_config, one_path = scope.databases[0].connection(config) - capability = create_capability( - db_path=one_path, config=one_config, defer_loading=False - ) + # What `run_chat` builds: the capability's own scope is the set, and + # the lent client is what narrows it. + capability = create_capability(config=config, defer_loading=False) async with HaikuRAG(config=config, sources=["alpha"]) as client: # What `ChatApp.on_mount` does. diff --git a/tests/multi_db/test_scope.py b/tests/multi_db/test_scope.py index 5f8f027e..9453ae34 100644 --- a/tests/multi_db/test_scope.py +++ b/tests/multi_db/test_scope.py @@ -125,10 +125,44 @@ class TestOneConfiguredLocation: config = self._config("s3://bucket/one.lancedb") [ref] = DatabaseScope.resolve(config).databases - one, db_path = ref.connection(config) - assert db_path is None - assert ConnectionMode.from_config(one) == ConnectionMode.OBJECT_STORAGE + assert ref.location == "s3://bucket/one.lancedb" + assert ConnectionMode.of(ref.location) == ConnectionMode.OBJECT_STORAGE + + +class TestSessionsOwnTheRef: + """A session is built from the resolved reference and hands storage only + its location; the configuration it keeps is the one the caller named.""" + + @pytest.mark.asyncio + async def test_a_session_opens_the_location_with_the_undivided_config( + self, tmp_path + ): + from haiku.rag.client.session import SingleDatabaseSession + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + [ref] = DatabaseScope.resolve(config, database_name="alpha").databases + + session = await SingleDatabaseSession(ref, config, read_only=True).open() + try: + assert session.source == "alpha" + assert session.location == ref.location + assert session.db_path == ref.location + assert session.store.location == ref.location + assert session.store._config is config + finally: + await session.aclose() + + @pytest.mark.asyncio + async def test_a_client_keeps_the_configuration_it_was_given(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config, sources=["alpha"]) as rag: + assert rag._config is config + assert set(rag._config.lancedb.databases) == {"alpha", "beta"} + assert rag.store.location == tmp_path / "alpha.lancedb" class TestLocate: diff --git a/tests/store/test_database_info.py b/tests/store/test_database_info.py index eecb2896..3930ca08 100644 --- a/tests/store/test_database_info.py +++ b/tests/store/test_database_info.py @@ -69,7 +69,7 @@ async def _seed(temp_db_path, *, version: str, with_items: bool = True): async def test_gather_database_info_reports_tables_and_settings(temp_db_path): await _seed(temp_db_path, version="1.2.3") - info = await gather_database_info(AppConfig(), temp_db_path) + info = await gather_database_info(temp_db_path, AppConfig()) assert info.exists is True assert info.path == str(temp_db_path) @@ -98,7 +98,7 @@ async def test_gather_database_info_flags_missing_table_and_pending_migrations( ): await _seed(temp_db_path, version="0.39.0", with_items=False) - info = await gather_database_info(AppConfig(), temp_db_path) + info = await gather_database_info(temp_db_path, AppConfig()) tables = {t.name: t for t in info.tables} assert tables["document_items"].exists is False @@ -112,7 +112,30 @@ async def test_gather_database_info_empty_database(temp_db_path): await lancedb.connect_async(temp_db_path) # creates the dir, no tables - info = await gather_database_info(AppConfig(), temp_db_path) + info = await gather_database_info(temp_db_path, AppConfig()) assert info.exists is False assert info.path == str(temp_db_path) + + +@pytest.mark.asyncio +async def test_gather_database_info_connects_to_the_location_it_is_given(): + """A remote location is passed to the connection as is and reported back + as the path; the configuration's own `uri` plays no part.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from haiku.rag.config.models import LanceDBConfig + + config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/other.lancedb")) + with patch( + "haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock + ) as mock_connect: + listing = MagicMock() + listing.tables = [] + mock_connect.return_value.list_tables = AsyncMock(return_value=listing) + + info = await gather_database_info("s3://bucket/papers.lancedb", config) + + assert mock_connect.call_args.args[0] == "s3://bucket/papers.lancedb" + assert info.path == "s3://bucket/papers.lancedb" + assert info.exists is False diff --git a/tests/store/test_v0_58_0_migration.py b/tests/store/test_v0_58_0_migration.py index 05e8c8d3..20ca619e 100644 --- a/tests/store/test_v0_58_0_migration.py +++ b/tests/store/test_v0_58_0_migration.py @@ -126,6 +126,43 @@ class TestV0_58_0MigrationEdgeCases: assert set(by_id) == {"a", "b"} # exactly one row each, no duplicates assert len(rows) == 2 + async def test_a_remote_store_has_no_disk_to_check(self, temp_db_path, monkeypatch): + """A store behind a URI has no local path: the reclaim vacuum runs + without a free-disk check.""" + from haiku.rag.store.upgrades import v0_58_0 + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [LegacyDocumentRecord(id="a", content="x", uri="u", metadata="{}")], + ) + await store.set_haiku_version("0.57.0") + + def no_disk(_path): + raise AssertionError("disk_usage consulted for a remote store") + + monkeypatch.setattr(v0_58_0.shutil, "disk_usage", no_disk) + vacuum_calls: list[int] = [] + + async with Store(temp_db_path, skip_migration_check=True) as store: + store.db_path = None + + async def fake_stats(): + return {"total_bytes": 10_000_000} + + monkeypatch.setattr(store.documents_table, "stats", fake_stats) + orig_vacuum = store.vacuum + + async def tracking_vacuum(*args, **kwargs): + vacuum_calls.append(1) + return await orig_vacuum(*args, **kwargs) + + monkeypatch.setattr(store, "vacuum", tracking_vacuum) + + await store.migrate() + + assert vacuum_calls == [1] + async def test_skips_vacuum_when_disk_is_tight(self, temp_db_path, monkeypatch): """When free disk can't cover one compacted copy, the split still completes but the reclaim vacuum is skipped.""" diff --git a/tests/test_database_scope.py b/tests/test_database_scope.py index 50e20e48..b13b075f 100644 --- a/tests/test_database_scope.py +++ b/tests/test_database_scope.py @@ -3,7 +3,6 @@ from pathlib import Path import pytest from haiku.rag.client.scope import DatabaseRef, DatabaseScope -from haiku.rag.client.session import default_db_path from haiku.rag.config.models import AppConfig, LanceDBConfig, StorageConfig from haiku.rag.store.exceptions import ( AmbiguousDatabaseError, @@ -89,16 +88,13 @@ class TestResolution: assert ref.uri == "" def test_a_path_selects_the_database_over_a_configured_uri(self): - """`--db` exists to override what is configured, and the configuration - derived from the ref is what makes the connection follow it.""" + """`--db` exists to override what is configured.""" config = _config(uri="s3://bucket/one.lancedb") scope = DatabaseScope.resolve(config, database_path=Path("/data/local")) [ref] = scope.databases - assert ref.db_path == Path("/data/local") - one, _ = ref.connection(config) - assert one.lancedb.uri == "" + assert ref.location == Path("/data/local") def test_nothing_configured_falls_back_to_the_data_directory(self, tmp_path): config = AppConfig(storage=StorageConfig(data_dir=tmp_path)) @@ -158,61 +154,23 @@ class TestResolution: DatabaseScope(()) -class TestConnectionDerivation: - """Opening one of a set must not disturb the configuration it came from.""" +class TestLocation: + """One value says where a database is: a path for a local one, a URI string + for a remote one. Storage connects to it as given.""" - def test_a_local_location_becomes_a_path(self): + def test_a_local_location_is_a_path(self): config = _config(databases={"alpha": "/data/alpha.lancedb"}) [ref] = DatabaseScope.resolve(config).databases - one, db_path = ref.connection(config) + assert ref.location == Path("/data/alpha.lancedb") - assert db_path == Path("/data/alpha.lancedb") - assert one.lancedb.uri == "" - assert one.lancedb.databases == {} - - def test_a_uri_location_stays_a_uri(self): + def test_a_uri_location_is_the_uri(self): config = _config(databases={"alpha": "s3://bucket/alpha.lancedb"}) [ref] = DatabaseScope.resolve(config).databases - one, db_path = ref.connection(config) + assert ref.location == "s3://bucket/alpha.lancedb" - assert db_path is None - assert one.lancedb.uri == "s3://bucket/alpha.lancedb" - - def test_the_original_configuration_is_untouched(self): - """Rewriting it in place is what left downstream code unable to tell a set - had been named.""" - config = _config(databases={"alpha": "/a.lancedb", "beta": "/b.lancedb"}) - - for ref in DatabaseScope.resolve(config).databases: - ref.connection(config) - - assert config.lancedb.databases == {"alpha": "/a.lancedb", "beta": "/b.lancedb"} - assert config.lancedb.uri == "" - - def test_each_derived_configuration_is_its_own_copy(self): - config = _config(databases={"alpha": "/a.lancedb", "beta": "s3://b/b.lancedb"}) - alpha, beta = DatabaseScope.resolve(config).databases - - one, _ = alpha.connection(config) - other, _ = beta.connection(config) - - assert one is not other - assert one.lancedb.uri == "" - assert other.lancedb.uri == "s3://b/b.lancedb" - - -def test_a_database_behind_a_uri_has_no_path_of_its_own(tmp_path): - """`connection` hands back no path for a URI, and the store still needs one: - the default stands in, and the URI is what decides where it connects.""" - config = AppConfig( - storage=StorageConfig(data_dir=tmp_path), - lancedb=LanceDBConfig(databases={"alpha": "s3://bucket/alpha.lancedb"}), - ) - [ref] = DatabaseScope.resolve(config).databases - - one, db_path = ref.connection(config) - - assert db_path is None - assert default_db_path(one) == tmp_path / "haiku.rag.lancedb" + def test_a_path_the_caller_gave_is_its_location(self): + assert DatabaseRef.at("/data/other.lancedb").location == Path( + "/data/other.lancedb" + ) diff --git a/tests/test_info.py b/tests/test_info.py index 8cf9eb9d..12273529 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -177,8 +177,7 @@ async def test_app_info_opens_a_named_remote_database(tmp_path): app = HaikuRAGApp(scope=scope, config=config) assert app._is_local is False - assert app._store_config.lancedb.uri == "s3://bucket/papers.lancedb" - assert app._store_config.lancedb.databases == {} + assert app._location == "s3://bucket/papers.lancedb" with patch( "haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock @@ -189,13 +188,12 @@ async def test_app_info_opens_a_named_remote_database(tmp_path): mock_db.list_tables = AsyncMock(return_value=mock_list_result) await app.info() - opened = mock_connect.call_args.args[0] - assert opened.lancedb.uri == "s3://bucket/papers.lancedb" + assert mock_connect.call_args.args[0] == "s3://bucket/papers.lancedb" async def test_app_doctor_opens_a_named_remote_database(): - """`run_doctor` connects with the configuration it is handed: the one - derived for the database, not the one naming the set.""" + """`run_doctor` is handed the database's location, not the configuration + naming the set.""" from haiku.rag.client.scope import DatabaseScope config = AppConfig( @@ -207,7 +205,7 @@ async def test_app_doctor_opens_a_named_remote_database(): run.return_value = MagicMock(checks=[], ok=True, duplicates=None) await app.doctor() - assert run.call_args.args[0].lancedb.uri == "s3://bucket/papers.lancedb" + assert run.call_args.args[1] == "s3://bucket/papers.lancedb" async def test_app_info_uses_connect_lancedb_for_remote(tmp_path): @@ -230,9 +228,8 @@ async def test_app_info_uses_connect_lancedb_for_remote(tmp_path): mock_db.list_tables = AsyncMock(return_value=mock_list_result) await app.info() - # The uri decides where it connects; the path argument is not read. mock_connect.assert_called_once() - assert mock_connect.call_args.args[0].lancedb.uri == "s3://bucket/path" + assert mock_connect.call_args.args[0] == "s3://bucket/path" @pytest.mark.asyncio diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 194c2f05..1e9637ae 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -340,23 +340,19 @@ class TestReportedLocation: @staticmethod def _session(location: str): from haiku.rag.client.scope import DatabaseScope - from haiku.rag.client.session import SingleDatabaseSession, default_db_path + from haiku.rag.client.session import SingleDatabaseSession from haiku.rag.config.models import AppConfig, LanceDBConfig config = AppConfig(lancedb=LanceDBConfig(databases={"alpha": location})) [ref] = DatabaseScope.resolve(config, database_name="alpha").databases - one, db_path = ref.connection(config) - return SingleDatabaseSession( - db_path if db_path is not None else default_db_path(one), - one, - source="alpha", - ) + return SingleDatabaseSession(ref, config) def test_a_named_remote_database_reports_its_uri(self): session = self._session("s3://bucket/alpha.lancedb") - assert isinstance(session.db_path, Path) + assert session.db_path is None assert session.location == "s3://bucket/alpha.lancedb" + assert session.source == "alpha" def test_a_named_local_database_reports_its_path(self): session = self._session("/data/alpha.lancedb") diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 2bdc5040..f3997b93 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -4,53 +4,44 @@ from unittest.mock import AsyncMock, patch import pytest from pydantic import ValidationError -from haiku.rag.config import get_config from haiku.rag.config.models import AppConfig, LanceDBConfig from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb class TestConnectionMode: - def test_local_when_uri_empty(self): - config = AppConfig(lancedb=LanceDBConfig(uri="")) - assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL + """A location is classified by itself: a path is local, `db://` is LanceDB + Cloud, any other scheme is object storage.""" + + def test_a_path_is_local(self, tmp_path): + assert ConnectionMode.of(tmp_path / "db.lancedb") == ConnectionMode.LOCAL + + def test_a_schemeless_string_is_local(self): + assert ConnectionMode.of("/data/db.lancedb") == ConnectionMode.LOCAL def test_cloud_when_db_uri(self): - config = AppConfig( - lancedb=LanceDBConfig( - uri="db://my-database", api_key="key", region="us-east-1" - ) - ) - assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD + assert ConnectionMode.of("db://my-database") == ConnectionMode.CLOUD - def test_object_storage_s3(self): - config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) - assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE - - def test_object_storage_gs(self): - config = AppConfig(lancedb=LanceDBConfig(uri="gs://bucket/path")) - assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE - - def test_object_storage_az(self): - config = AppConfig(lancedb=LanceDBConfig(uri="az://container/path")) - assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE - - def test_object_storage_hdfs(self): - config = AppConfig(lancedb=LanceDBConfig(uri="hdfs://namenode/path")) - assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE - - def test_unknown_uri_treated_as_object_storage(self): - config = AppConfig(lancedb=LanceDBConfig(uri="custom://something")) - assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE + @pytest.mark.parametrize( + "uri", + [ + "s3://bucket/path", + "gs://bucket/path", + "az://container/path", + "hdfs://namenode/path", + "custom://something", + ], + ) + def test_any_other_scheme_is_object_storage(self, uri): + assert ConnectionMode.of(uri) == ConnectionMode.OBJECT_STORAGE class TestConnectLancedb: @pytest.mark.asyncio async def test_local_passes_absolute_db_path(self, temp_db_path): - config = AppConfig(lancedb=LanceDBConfig(uri="")) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config, db_path=temp_db_path) + await connect_lancedb(temp_db_path, AppConfig()) mock_connect.assert_awaited_once() assert mock_connect.call_args.args == (temp_db_path.absolute(),) @@ -60,25 +51,34 @@ class TestConnectLancedb: monkeypatch.chdir(tmp_path) relative = Path("db/rag.lancedb") - config = AppConfig(lancedb=LanceDBConfig(uri="")) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config, db_path=relative) + await connect_lancedb(relative, AppConfig()) mock_connect.assert_awaited_once() assert mock_connect.call_args.args == (relative.absolute(),) + @pytest.mark.asyncio + async def test_the_configured_uri_is_not_consulted(self, temp_db_path): + """Storage connects to the location it is handed; placement is the + caller's, and the configuration's own `uri` never redirects it.""" + config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/db.lancedb")) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(temp_db_path, config) + assert mock_connect.call_args.args == (temp_db_path.absolute(),) + assert "uri" not in mock_connect.call_args.kwargs + @pytest.mark.asyncio async def test_cloud_passes_uri_api_key_region(self): config = AppConfig( - lancedb=LanceDBConfig( - uri="db://my-database", api_key="test-key", region="us-west-2" - ) + lancedb=LanceDBConfig(api_key="test-key", region="us-west-2") ) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) + await connect_lancedb("db://my-database", config) mock_connect.assert_awaited_once() kwargs = mock_connect.call_args.kwargs assert kwargs["uri"] == "db://my-database" @@ -89,7 +89,6 @@ class TestConnectLancedb: async def test_object_storage_passes_uri_and_storage_options(self): config = AppConfig( lancedb=LanceDBConfig( - uri="s3://bucket/path", storage_options={ "endpoint": "http://minio:9000", "region": "us-east-1", @@ -99,7 +98,7 @@ class TestConnectLancedb: with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", config) mock_connect.assert_awaited_once() kwargs = mock_connect.call_args.kwargs assert kwargs["uri"] == "s3://bucket/path" @@ -110,23 +109,25 @@ class TestConnectLancedb: @pytest.mark.asyncio async def test_object_storage_without_storage_options(self): - config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", AppConfig()) mock_connect.assert_awaited_once() kwargs = mock_connect.call_args.kwargs assert kwargs["uri"] == "s3://bucket/path" assert "storage_options" not in kwargs - @pytest.mark.asyncio - async def test_local_without_db_path_raises(self): - config = AppConfig(lancedb=LanceDBConfig(uri="")) - with pytest.raises( - ValueError, match="No lancedb.uri configured and no db_path provided" - ): - await connect_lancedb(config) + +def _remote_store(location: str, config: AppConfig | None = None) -> Store: + """A store over a remote location, opened against a mocked connection.""" + return Store( + location, + config=config, + create=True, + skip_validation=True, + skip_migration_check=True, + ) class TestStoreConnectionMode: @@ -134,132 +135,127 @@ class TestStoreConnectionMode: async def test_store_connection_mode_local(self, temp_db_path): async with Store(temp_db_path, create=True) as store: assert store._connection_mode == ConnectionMode.LOCAL + assert store.location == temp_db_path + assert store.db_path == temp_db_path @pytest.mark.asyncio - async def test_store_connection_mode_cloud(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with ( - patch.object(get_config().lancedb, "uri", "db://test-database"), - patch.object(get_config().lancedb, "api_key", "test-api-key"), - patch.object(get_config().lancedb, "region", "us-east-1"), - ): + async def test_a_local_store_ignores_the_configured_uri(self, temp_db_path): + config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/db.lancedb")) + async with Store(temp_db_path, config=config, create=True) as store: + assert store._connection_mode == ConnectionMode.LOCAL + assert store.db_path == temp_db_path + + @pytest.mark.asyncio + async def test_store_connection_mode_cloud(self): + config = AppConfig(lancedb=LanceDBConfig(api_key="key", region="us-east-1")) + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), + ): + async with _remote_store("db://test-database", config) as store: assert store._connection_mode == ConnectionMode.CLOUD + assert store.location == "db://test-database" + assert store.db_path is None @pytest.mark.asyncio - async def test_store_connection_mode_object_storage(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with patch.object(get_config().lancedb, "uri", "s3://bucket/path"): + async def test_store_connection_mode_object_storage(self): + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), + ): + async with _remote_store("s3://bucket/path") as store: assert store._connection_mode == ConnectionMode.OBJECT_STORAGE + assert store.db_path is None + + +def _remote_store_with_mock_tables(location: str) -> Store: + """A remote store whose tables are mocks: the mode decision is under test, + not the tables.""" + store = _remote_store(location) + store.chunks_table = AsyncMock() + return store class TestVacuumByConnectionMode: @pytest.mark.asyncio - async def test_cloud_skips_vacuum(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with ( - patch.object(get_config().lancedb, "uri", "db://test-database"), - patch.object(get_config().lancedb, "api_key", "test-api-key"), - patch.object(get_config().lancedb, "region", "us-east-1"), - ): - with patch.object( - store.chunks_table, "optimize", new_callable=AsyncMock - ) as mock_optimize: - await store.vacuum() - mock_optimize.assert_not_called() + async def test_cloud_skips_vacuum(self): + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), + ): + async with _remote_store_with_mock_tables("db://test-database") as store: + await store.vacuum() + store.chunks_table.optimize.assert_not_awaited() @pytest.mark.asyncio - async def test_object_storage_runs_vacuum(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with patch.object(get_config().lancedb, "uri", "s3://bucket/path"): + async def test_object_storage_runs_vacuum(self): + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), + ): + async with _remote_store_with_mock_tables("s3://bucket/path") as store: + store.chunks_table.tags.list = AsyncMock(return_value={}) with patch.object( - store.chunks_table, "optimize", new_callable=AsyncMock - ) as mock_optimize: + store, "_tables", return_value={"chunks": store.chunks_table} + ): await store.vacuum() - mock_optimize.assert_called() + store.chunks_table.optimize.assert_awaited_once() @pytest.mark.asyncio async def test_local_runs_vacuum(self, temp_db_path): async with Store(temp_db_path, create=True) as store: - with patch.object(get_config().lancedb, "uri", ""): - with patch.object( - store.chunks_table, "optimize", new_callable=AsyncMock - ) as mock_optimize: - await store.vacuum() - mock_optimize.assert_called() + with patch.object( + store.chunks_table, "optimize", new_callable=AsyncMock + ) as mock_optimize: + await store.vacuum() + mock_optimize.assert_called() class TestVectorIndexByConnectionMode: @pytest.mark.asyncio - async def test_cloud_skips_index_creation(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with ( - patch.object(get_config().lancedb, "uri", "db://test-database"), - patch.object(get_config().lancedb, "api_key", "test-api-key"), - patch.object(get_config().lancedb, "region", "us-east-1"), - ): - with patch.object( - store.chunks_table, "count_rows", new_callable=AsyncMock - ) as mock_count: - await store._ensure_vector_index() - mock_count.assert_not_called() - - @pytest.mark.asyncio - async def test_object_storage_runs_index_creation(self, temp_db_path): - async with Store(temp_db_path, create=True) as store: - with patch.object(get_config().lancedb, "uri", "s3://bucket/path"): - with patch.object( - store.chunks_table, - "count_rows", - new_callable=AsyncMock, - return_value=0, - ) as mock_count: - await store._ensure_vector_index() - mock_count.assert_called() - - -class TestStoreSkipsPathValidationForRemote: - @pytest.mark.asyncio - async def test_skips_path_check_for_cloud(self, tmp_path): - nonexistent = tmp_path / "does_not_exist" / "db.lancedb" - config = AppConfig( - lancedb=LanceDBConfig( - uri="db://test-database", api_key="key", region="us-east-1" - ) - ) - with patch( - "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + async def test_cloud_skips_index_creation(self): + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), ): - with patch.object(Store, "_init_tables", new_callable=AsyncMock): - async with Store( - nonexistent, - config=config, - create=True, - skip_validation=True, - skip_migration_check=True, - ) as store: - assert store is not None + async with _remote_store_with_mock_tables("db://test-database") as store: + await store._ensure_vector_index() + store.chunks_table.count_rows.assert_not_awaited() @pytest.mark.asyncio - async def test_skips_path_check_for_object_storage(self, tmp_path): - nonexistent = tmp_path / "does_not_exist" / "db.lancedb" - config = AppConfig( - lancedb=LanceDBConfig( - uri="s3://bucket/path", - storage_options={"endpoint": "http://localhost:9000"}, - ) - ) - with patch( - "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + async def test_object_storage_runs_index_creation(self): + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch.object(Store, "_init_tables", new_callable=AsyncMock), ): - with patch.object(Store, "_init_tables", new_callable=AsyncMock): - async with Store( - nonexistent, - config=config, - create=True, - skip_validation=True, - skip_migration_check=True, - ) as store: - assert store is not None + async with _remote_store_with_mock_tables("s3://bucket/path") as store: + store.chunks_table.count_rows = AsyncMock(return_value=0) + await store._ensure_vector_index() + store.chunks_table.count_rows.assert_awaited_once() + + +class TestLocationIsFixed: + @pytest.mark.asyncio + async def test_a_store_keeps_the_location_it_opened(self, temp_db_path): + """`db_path` and the connection mode derive from the location once; a + store cannot be pointed elsewhere after it is built.""" + async with Store(temp_db_path, create=True) as store: + with pytest.raises(AttributeError): + store.location = "s3://bucket/path" # type: ignore[misc] + assert store.location == temp_db_path + assert store._connection_mode == ConnectionMode.LOCAL class TestInitFailureCleanup: @@ -412,33 +408,25 @@ class TestStoreMiscellany: class TestSessionAndConsistency: @pytest.mark.asyncio async def test_session_is_shared_across_connections(self): - config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) + config = AppConfig() with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", config) + await connect_lancedb("s3://bucket/path", config) sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] assert sessions[0] is sessions[1] @pytest.mark.asyncio async def test_cache_sizes_select_distinct_sessions(self): - small = AppConfig( - lancedb=LanceDBConfig( - uri="s3://bucket/path", index_cache_size_bytes=1 << 20 - ) - ) - large = AppConfig( - lancedb=LanceDBConfig( - uri="s3://bucket/path", index_cache_size_bytes=1 << 30 - ) - ) + small = AppConfig(lancedb=LanceDBConfig(index_cache_size_bytes=1 << 20)) + large = AppConfig(lancedb=LanceDBConfig(index_cache_size_bytes=1 << 30)) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(small) - await connect_lancedb(large) + await connect_lancedb("s3://bucket/path", small) + await connect_lancedb("s3://bucket/path", large) sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] assert sessions[0] is not sessions[1] @@ -447,7 +435,6 @@ class TestSessionAndConsistency: async def test_both_cache_sizes_are_applied(self): config = AppConfig( lancedb=LanceDBConfig( - uri="s3://bucket/path", index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20, ) @@ -458,7 +445,7 @@ class TestSessionAndConsistency: ), patch("haiku.rag.store.engine.lancedb.Session") as mock_session, ): - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", config) mock_session.assert_called_once_with( index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20 @@ -466,15 +453,11 @@ class TestSessionAndConsistency: @pytest.mark.asyncio async def test_read_consistency_interval_is_forwarded(self): - config = AppConfig( - lancedb=LanceDBConfig( - uri="s3://bucket/path", read_consistency_interval_seconds=5 - ) - ) + config = AppConfig(lancedb=LanceDBConfig(read_consistency_interval_seconds=5)) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", config) assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( seconds=5 @@ -483,14 +466,12 @@ class TestSessionAndConsistency: @pytest.mark.asyncio async def test_read_consistency_interval_omitted_when_disabled(self): config = AppConfig( - lancedb=LanceDBConfig( - uri="s3://bucket/path", read_consistency_interval_seconds=None - ) + lancedb=LanceDBConfig(read_consistency_interval_seconds=None) ) with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config) + await connect_lancedb("s3://bucket/path", config) assert mock_connect.call_args.kwargs["read_consistency_interval"] is None @@ -500,7 +481,7 @@ class TestSessionAndConsistency: with patch( "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: - await connect_lancedb(config, tmp_path / "db.lancedb") + await connect_lancedb(tmp_path / "db.lancedb", config) assert mock_connect.call_args.kwargs["session"] is not None assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index bccf1546..5d814425 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -89,7 +89,7 @@ async def _remote_client(config: AppConfig): async def test_store_connect_and_create(tmp_path, config): from haiku.rag.store.info import get_database_stats - async with Store(tmp_path / "unused", config=config, create=True) as store: + async with Store(config.lancedb.uri, config=config, create=True) as store: stats = await get_database_stats(store.db) assert stats["documents"]["exists"] assert stats["chunks"]["exists"] @@ -97,7 +97,7 @@ async def test_store_connect_and_create(tmp_path, config): @pytest.mark.asyncio async def test_store_vacuum(tmp_path, config): - async with Store(tmp_path / "unused", config=config, create=True) as store: + async with Store(config.lancedb.uri, config=config, create=True) as store: await store.vacuum() @@ -106,7 +106,7 @@ async def test_store_add_document(tmp_path, config): from haiku.rag.store.info import get_database_stats from haiku.rag.store.schema import DocumentRecord - async with Store(tmp_path / "unused", config=config, create=True) as store: + async with Store(config.lancedb.uri, config=config, create=True) as store: doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.") await store.documents_table.add([doc])