Compose a set out of single-database sessions

FederatedSession opens the databases a query covers and owns their
teardown; the client keeps the wrappers it hands out. A wrapper releases
what it built and never closes the database it borrowed.

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -16,7 +16,12 @@ from urllib.parse import urlparse
 import httpx
 
 from haiku.rag.client.documents import DocumentImport
-from haiku.rag.client.session import SingleDatabaseSession
+from haiku.rag.client.scope import DatabaseRef, DatabaseScope
+from haiku.rag.client.session import (
+    FederatedSession,
+    SingleDatabaseSession,
+    aclose_quietly,
+)
 from haiku.rag.config import AppConfig, get_config
 from haiku.rag.converters import get_converter
 from haiku.rag.embeddings import get_embedder
@@ -70,18 +75,6 @@ async def first_found(
     return None
 
 
-async def _aclose_quietly(closeable: Any, what: str) -> None:
-    """Close, reporting failure to the log rather than raising.
-
-    Teardown can run while an exception unwinds, so a raising close must
-    neither mask that exception nor stop a sibling from being closed.
-    """
-    try:
-        await closeable.aclose()
-    except Exception:
-        logger.debug("Closing the %s failed on teardown", what, exc_info=True)
-
-
 def _spell(embedding: tuple[str | None, str | None, int | None]) -> str:
     """An embedder identity, for an error message."""
     provider, name, vector_dim = embedding
@@ -150,6 +143,8 @@ class HaikuRAG:
         self._clients_lock = asyncio.Lock()
         self._source: str | None = None
         self._session: SingleDatabaseSession | None = None
+        self._federated_session: FederatedSession | None = None
+        self._owns_session = True
 
     @property
     def store(self) -> Store:
@@ -254,6 +249,17 @@ class HaikuRAG:
                     "sources=[name]"
                 )
             self._federated = selected
+            self._federated_session = FederatedSession(
+                DatabaseScope(
+                    tuple(
+                        DatabaseRef.configured(name, location)
+                        for name, location in selected.items()
+                    )
+                ),
+                self._config,
+                skip_validation=self._skip_validation,
+                read_only=self._read_only,
+            )
             return self
         if selected:
             [(self._source, location)] = selected.items()
@@ -280,36 +286,38 @@ class HaikuRAG:
         Opening is per query rather than at entry: a set of 25 configured
         databases is typically queried a few at a time, and a database nobody
         asked for must not be able to fail a query, or be opened for nothing.
-
-        Missing ones open together: on object storage a serial loop makes the
-        first query cost the sum of the opens.
         """
+        assert self._federated_session is not None
         names = _without_repeats(names)
-        unknown = [n for n in names if n not in self._federated]
-        if unknown:
-            raise KeyError(
-                f"unknown database(s) {', '.join(sorted(unknown))}; configured: "
-                f"{', '.join(sorted(self._federated))}"
-            )
-        async with self._clients_lock:
-            missing = [n for n in names if n not in self._clients]
-            if missing:
-                opened = await asyncio.gather(
-                    *(self._open_client(n, self._federated[n]) for n in missing),
-                    return_exceptions=True,
-                )
-                # Whatever opened is tracked before the failure is reported, so
-                # `__aexit__` closes it: `gather` does not cancel the siblings of
-                # the one that raised, and an untracked connection leaks.
-                failure: BaseException | None = None
-                for name, result in zip(missing, opened, strict=True):
-                    if isinstance(result, BaseException):
-                        failure = failure or result
-                    else:
-                        self._clients[name] = result
-                if failure is not None:
-                    raise failure
-        return [self._clients[n] for n in names]
+        sessions = await self._federated_session.sessions_for(names)
+        return [
+            self._facade_for(name, session)
+            for name, session in zip(names, sessions, strict=True)
+        ]
+
+    def _facade_for(self, name: str, session: SingleDatabaseSession) -> "HaikuRAG":
+        """The client for one covered database, made once and kept.
+
+        The facade is public and the session is not, so the wrapper lives here
+        while the federated session keeps the database it wraps. A borrowed facade
+        never closes what it did not open.
+        """
+        facade = self._clients.get(name)
+        if facade is None:
+            facade = HaikuRAG._from_session(session)
+            self._clients[name] = facade
+        return facade
+
+    @classmethod
+    def _from_session(cls, session: SingleDatabaseSession) -> "HaikuRAG":
+        """A client over a database another session opened and will close."""
+        client = cls(
+            session.db_path, config=session.config, read_only=session.read_only
+        )
+        client._session = session
+        client._source = session.source
+        client._owns_session = False
+        return client
 
     def _require_one_embedder(self, clients: "list[HaikuRAG]") -> None:
         """Fail when two of these databases were written with different embedders.
@@ -340,49 +348,44 @@ class HaikuRAG:
                     "query once, so their vectors are not comparable"
                 )
 
-    async def _open_client(self, name: str, location: str) -> "HaikuRAG":
-        uri, db_path = locate_database(location)
-        config = self._config.model_copy(deep=True)
-        config.lancedb.databases = {}
-        config.lancedb.uri = uri
-        client = HaikuRAG(
-            db_path,
-            config=config,
-            skip_validation=self._skip_validation,
-            read_only=self._read_only,
-        )
-        client._source = name
-        return await client.__aenter__()
-
-    async def _close_clients(self) -> None:
-        for client in self._clients.values():
-            try:
-                await client.__aexit__(None, None, None)
-            except Exception:
-                logger.debug("Closing a database failed on teardown", exc_info=True)
-        self._clients.clear()
-
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
         """Async context manager exit."""
         # Branch on what this client covers, not on what it happened to open:
         # a federating client that answered no query has nothing open and no
         # store either.
         if self._federated:
-            await self._close_clients()
+            assert self._federated_session is not None
+            # The wrappers built over covered databases hold rerankers of their
+            # own; the databases themselves are the federated session's to close.
+            for facade in self._clients.values():
+                await facade._release_own()
+            self._clients.clear()
+            await self._federated_session.aclose()
             # The set shares this client's embedder and reranker, so this is the
             # only place they are closed — and only if anything built them.
             await self._aclose_cached("embedder")
             await self._aclose_cached("reranker")
             return False
-        await self._await_vacuum_tasks()
-        # Accessed so the store's embedder is closed even where nothing used it;
-        # `cached_property` stores it, which is what `_aclose_cached` discards.
-        _ = self.embedder
-        await self._aclose_cached("embedder")
-        await self._aclose_cached("reranker")
-        self.close()
+        if not self._owns_session:
+            await self._release_own()
+            return False
+        assert self._session is not None
+        # The session drains, releases its store's embedder and closes; the
+        # cached reference here is only discarded, never closed twice.
+        await self._release_own()
+        await self._session.aclose()
         return False
 
+    async def _release_own(self) -> None:
+        """Release what this client built, leaving the database to its owner.
+
+        The embedder belongs to the store and is closed with it, so the cached
+        reference is only discarded. The reranker is this client's own, built on
+        its first text query.
+        """
+        self.__dict__.pop("embedder", None)
+        await self._aclose_cached("reranker")
+
     async def _aclose_cached(self, name: str) -> None:
         """Close a cached_property this client materialized, and discard it.
 
@@ -391,7 +394,7 @@ class HaikuRAG:
         """
         cached = self.__dict__.pop(name, None)
         if cached is not None:
-            await _aclose_quietly(cached, name)
+            await aclose_quietly(cached, name)
 
     async def _await_vacuum_tasks(self) -> None:
         if self._session is not None:
@@ -885,7 +888,13 @@ class HaikuRAG:
         await self.store.vacuum()
 
     def close(self):
-        """Close the underlying store connection."""
+        """Close the underlying store connection.
+
+        A client covering one of a set borrows that database and never closes
+        it: the set opened it and the set closes it.
+        """
         self._require_one_database("close")
+        if not self._owns_session:
+            return
         assert self._session is not None
         self._session.close()

haiku_rag_slim/haiku/rag/client/session.py

@@ -2,7 +2,9 @@ import asyncio
 import logging
 from pathlib import Path
 from time import monotonic
+from typing import Any
 
+from haiku.rag.client.scope import DatabaseRef, DatabaseScope
 from haiku.rag.config import AppConfig
 from haiku.rag.store.engine import Store
 from haiku.rag.store.exceptions import (
@@ -29,12 +31,33 @@ _VACUUM_MIN_INTERVAL_S = 300.0
 _NAMEABLE_FAILURES = (MigrationRequiredError, ConfigMismatchError, ReadOnlyError)
 
 
+async def aclose_quietly(closeable: Any, what: str) -> None:
+    """Close, reporting failure to the log rather than raising.
+
+    Teardown can run while an exception unwinds, so a raising close must
+    neither mask that exception nor stop a sibling from being closed.
+    """
+    try:
+        await closeable.aclose()
+    except Exception:
+        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.
+
+    ``db_path``, ``config``, ``read_only`` and ``source`` are readable because a
+    client facade is built over a session it does not own, and needs to report
+    the same things it would have reported had it opened the database itself.
     """
 
     def __init__(
@@ -47,12 +70,12 @@ class SingleDatabaseSession:
         read_only: bool = False,
         source: str | None = None,
     ) -> None:
-        self._db_path = db_path
-        self._config = config
+        self.db_path = db_path
+        self.config = config
+        self.read_only = read_only
+        self.source = source
         self._skip_validation = skip_validation
         self._create = create
-        self._read_only = read_only
-        self.source = source
         self._vacuum_tasks: set[asyncio.Task] = set()
         self._last_vacuum_at: float | None = None
         self._vacuum_dirty = False
@@ -62,11 +85,11 @@ class SingleDatabaseSession:
         failure: str | None = None
         try:
             self.store = Store(
-                self._db_path,
-                config=self._config,
+                self.db_path,
+                config=self.config,
                 skip_validation=self._skip_validation,
                 create=self._create,
-                read_only=self._read_only,
+                read_only=self.read_only,
             )
             # If _initialize fails mid-way (e.g. migration check raises after
             # connect), close the store so we don't leak the LanceDB connection —
@@ -142,6 +165,94 @@ class SingleDatabaseSession:
         self._vacuum_tasks.add(task)
         task.add_done_callback(self._vacuum_tasks.discard)
 
+    async def aclose(self) -> None:
+        """Drain, release the embedder, and close the connection.
+
+        The store owns the embedder, so releasing it belongs here rather than
+        with whoever happened to hold the session.
+        """
+        await self.drain_vacuum()
+        await aclose_quietly(self.store.embedder, "embedder")
+        self.close()
+
     def close(self) -> None:
         """Close the underlying store connection."""
         self.store.close()
+
+
+class FederatedSession:
+    """Several databases, read as one.
+
+    Composes single-database sessions and owns their teardown. They open on first
+    use rather than at entry: which databases a query covers is a per-query
+    choice, so a database nobody asked for must neither be opened for nothing nor
+    be able to fail a query.
+
+    Reads only. Writing names a database, and naming one is what
+    ``SingleDatabaseSession`` is.
+    """
+
+    def __init__(
+        self,
+        scope: DatabaseScope,
+        config: AppConfig,
+        *,
+        skip_validation: bool = False,
+        read_only: bool = False,
+    ) -> None:
+        self._refs: dict[str, DatabaseRef] = {
+            ref.name: ref for ref in scope.databases if ref.name is not None
+        }
+        self._config = config
+        self._skip_validation = skip_validation
+        self._read_only = read_only
+        self._sessions: dict[str, SingleDatabaseSession] = {}
+        self._lock = asyncio.Lock()
+
+    async def sessions_for(self, names: list[str]) -> list[SingleDatabaseSession]:
+        """The sessions for these databases, opening any not yet open.
+
+        Missing ones open together: on object storage a serial loop makes the
+        first query cost the sum of the opens.
+        """
+        unknown = [name for name in names if name not in self._refs]
+        if unknown:
+            raise KeyError(
+                f"unknown database(s) {', '.join(sorted(unknown))}; configured: "
+                f"{', '.join(sorted(self._refs))}"
+            )
+        async with self._lock:
+            missing = [name for name in names if name not in self._sessions]
+            if missing:
+                opened = await asyncio.gather(
+                    *(self._open(self._refs[name]) for name in missing),
+                    return_exceptions=True,
+                )
+                # Whatever opened is tracked before the failure is reported, so
+                # teardown closes it: `gather` does not cancel the siblings of the
+                # one that raised, and an untracked connection leaks.
+                failure: BaseException | None = None
+                for name, result in zip(missing, opened, strict=True):
+                    if isinstance(result, BaseException):
+                        failure = failure or result
+                    else:
+                        self._sessions[name] = result
+                if failure is not None:
+                    raise failure
+        return [self._sessions[name] for name in names]
+
+    async def _open(self, ref: DatabaseRef) -> SingleDatabaseSession:
+        one, db_path = ref.connection(self._config)
+        return await SingleDatabaseSession(
+            db_path if db_path is not None else default_db_path(one),
+            one,
+            skip_validation=self._skip_validation,
+            read_only=self._read_only,
+            source=ref.name,
+        ).open()
+
+    async def aclose(self) -> None:
+        """Close every database this session opened."""
+        for session in self._sessions.values():
+            await aclose_quietly(session, "database")
+        self._sessions.clear()

tests/test_client.py

@@ -1185,6 +1185,19 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
         assert docling.get_docling_document() is not None
 
 
+async def test_close_releases_the_connection(temp_db_path):
+    """A caller managing the client itself closes it directly; leaving the
+    context goes through the session instead."""
+    client = HaikuRAG(temp_db_path, create=True)
+    await client.__aenter__()
+    store = client.store
+    assert store.db.is_open()
+
+    client.close()
+
+    assert not store.db.is_open()
+
+
 async def test_delete_marks_vacuum_dirty(temp_db_path):
     """A delete adds tombstone/table versions, so it must enter the auto-vacuum
     lifecycle — otherwise a delete-only run closes without a final vacuum."""

tests/test_database_scope.py

@@ -3,6 +3,7 @@ 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
 
@@ -176,3 +177,18 @@ class TestConnectionDerivation:
         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"

tests/test_multi_db.py

@@ -6,6 +6,7 @@ from docling_core.types.doc.labels import DocItemLabel
 from pydantic import ValidationError
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.client.session import FederatedSession
 from haiku.rag.config import get_config
 from haiku.rag.config.models import AppConfig, LanceDBConfig
 from haiku.rag.store.exceptions import (
@@ -154,16 +155,17 @@ class TestOpeningDatabases:
             await _seed(config, name, [f"{name} document about cats"])
 
         async with HaikuRAG(config=config) as rag:
+            assert rag._federated_session is not None
             barrier = asyncio.Barrier(len(names))
-            open_one = rag._open_client
+            open_one = rag._federated_session._open
 
-            async def gated(name: str, location: str):
+            async def gated(ref):
                 # Every open has to be in flight before any of them finishes, so
                 # a serial loop cannot get past this and the wait times out.
                 await barrier.wait()
-                return await open_one(name, location)
+                return await open_one(ref)
 
-            rag._open_client = gated
+            rag._federated_session._open = gated
             clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
 
         assert {client._source for client in clients} == set(names)
@@ -180,7 +182,8 @@ class TestOpeningDatabases:
             with pytest.raises(SourceUnavailableError, match="beta"):
                 await rag.clients_for(["alpha", "beta"])
 
-            assert set(rag._clients) == {"alpha"}
+            assert rag._federated_session is not None
+            assert set(rag._federated_session._sessions) == {"alpha"}
 
     @pytest.mark.asyncio
     async def test_a_database_named_twice_is_opened_once(self, tmp_path):
@@ -342,6 +345,83 @@ class TestLookupByIdentifier:
             assert await rag.get_document_by_uri("test://nowhere") is None
 
 
+class TestClosingASet:
+    @pytest.mark.asyncio
+    async def test_every_database_opened_is_released(self, tmp_path):
+        """A covered database owns an embedder and may owe a vacuum. Closing only
+        its connection would leave both behind."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+        await _seed(config, "beta", ["beta one"])
+
+        released: list[str | None] = []
+        drained: list[str | None] = []
+
+        async with HaikuRAG(config=config, read_only=True) as rag:
+            assert rag._federated_session is not None
+            await rag.clients_for(["alpha", "beta"])
+            for name, session in rag._federated_session._sessions.items():
+                original = session.store.embedder.aclose
+                drain = session.drain_vacuum
+
+                async def release(_original=original, _name=name):
+                    released.append(_name)
+                    return await _original()
+
+                async def drain_it(_drain=drain, _name=name):
+                    drained.append(_name)
+                    return await _drain()
+
+                session.store.embedder.aclose = release
+                session.drain_vacuum = drain_it
+
+        assert sorted(released) == ["alpha", "beta"]
+        assert sorted(drained) == ["alpha", "beta"]
+
+
+class TestBorrowedDatabases:
+    """A client for one of a set wraps a database the set opened."""
+
+    @pytest.mark.asyncio
+    async def test_closing_a_borrowed_client_leaves_the_set_working(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            (alpha,) = await rag.clients_for(["alpha"])
+            store = alpha.store
+
+            alpha.close()
+            assert store.db.is_open(), "close() closed a database it borrowed"
+
+            await alpha.__aexit__(None, None, None)
+            assert store.db.is_open(), "exit closed a database it borrowed"
+
+            results = await rag.search("cats", search_type="fts")
+
+        assert {r.source for r in results} == {"alpha", "beta"}
+        assert not store.db.is_open(), "the set left a database open"
+
+    @pytest.mark.asyncio
+    async def test_a_borrowed_client_releases_what_it_built(self, tmp_path):
+        """Its reranker is its own; the database it wraps is not."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        closed: list[str] = []
+
+        class Reranker:
+            async def aclose(self):
+                closed.append("reranker")
+
+        async with HaikuRAG(config=config) as rag:
+            (alpha,) = await rag.clients_for(["alpha"])
+            alpha.__dict__["reranker"] = Reranker()
+
+        assert closed == ["reranker"]
+
+
 class TestDatabaseIndependentWork:
     """Converting, chunking and titling are functions of the configuration, not
     of a database, so covering a set does not stop them."""
@@ -354,11 +434,11 @@ class TestDatabaseIndependentWork:
 
         opened: list[str] = []
 
-        async def refuse(self, name, location):
-            opened.append(name)
+        async def refuse(self, ref):
+            opened.append(ref.name)
             raise AssertionError("opened a database to chunk a document")
 
-        monkeypatch.setattr(HaikuRAG, "_open_client", refuse)
+        monkeypatch.setattr(FederatedSession, "_open", refuse)
 
         doc = DoclingDocument(name="note")
         doc.add_text(
@@ -484,6 +564,8 @@ class TestOperationsThatNeedOneDatabase:
                 await rag.create_document("orphan")
             with pytest.raises(AmbiguousDatabaseError, match="clients_for"):
                 await rag.vacuum()
+            with pytest.raises(AmbiguousDatabaseError, match="close"):
+                rag.close()
 
     @pytest.mark.asyncio
     async def test_a_set_has_no_store_of_its_own(self, tmp_path):
@@ -850,15 +932,21 @@ class TestRerankerFusion:
         rag = HaikuRAG(config=config)
         await rag.__aenter__()
         await rag.clients_for(["alpha", "beta"])
+        assert rag._federated_session is not None
+        sessions = rag._federated_session._sessions
 
-        async def boom(exc_type, exc_val, exc_tb):
+        async def boom():
             raise RuntimeError("close failed")
 
-        rag._clients["alpha"].__aexit__ = boom  # ty: ignore[invalid-assignment]
+        sessions["alpha"].aclose = boom  # ty: ignore[invalid-assignment]
+        beta = sessions["beta"].store
 
         await rag.__aexit__(None, None, None)
 
+        # The failure is swallowed, and the sibling is still closed after it.
         assert rag._clients == {}
+        assert rag._federated_session._sessions == {}
+        assert not beta.db.is_open()
 
     @pytest.mark.asyncio
     async def test_multimodal_reranking_attaches_each_database_own_pictures(