From e8390ca7475999c69ec1bc9f9a09ab8eb02bd06d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 25 Aug 2026 08:56:09 +0300 Subject: [PATCH] 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 | 155 +++++++++++--------- haiku_rag_slim/haiku/rag/client/session.py | 125 +++++++++++++++- tests/test_client.py | 13 ++ tests/test_database_scope.py | 16 ++ tests/test_multi_db.py | 108 ++++++++++++-- 5 files changed, 327 insertions(+), 90 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index a6d25ad2..d397f6b7 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/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() diff --git a/haiku_rag_slim/haiku/rag/client/session.py b/haiku_rag_slim/haiku/rag/client/session.py index ce10fcfd..8eed416e 100644 --- a/haiku_rag_slim/haiku/rag/client/session.py +++ b/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() diff --git a/tests/test_client.py b/tests/test_client.py index 7e473fab..ef1bdfad 100644 --- a/tests/test_client.py +++ b/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.""" diff --git a/tests/test_database_scope.py b/tests/test_database_scope.py index bd845f6e..c1f3eba1 100644 --- a/tests/test_database_scope.py +++ b/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" diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index 63946667..428a3e35 100644 --- a/tests/test_multi_db.py +++ b/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(