Reuse the database a borrowed client already holds

`clients_for` hands back a client over a database the covering one owns, so
`async with` on it opened a second session and assigned it, and teardown
declined to close what this client did not open. Entry returns the client as
it stands.
This commit is contained in:
Yiorgis Gozadinos 2026-08-26 13:03:47 +03:00
parent de8075fcd1
commit 2b2de5a61f
No known key found for this signature in database
2 changed files with 33 additions and 0 deletions

View file

@ -295,11 +295,19 @@ class HaikuRAG:
async def __aenter__(self):
"""Async context manager entry — initializes store and repositories.
A client borrowing a database is already open and returns itself.
Opening a second session would leak it: teardown declines to close what
this client did not open.
A client covering several databases opens none of them here: which are
searched is a per-query choice, so they open on first use. `store` and the
repositories stay unset in that case, since they have no unambiguous
meaning across a set.
"""
if not self._owns_session:
assert self._session is not None
return self
scope = self._resolve_scope()
if scope.covers_multiple:
if self._create:
@ -338,6 +346,10 @@ 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.
The clients returned borrow their databases from this one and are valid
only while it is open. Closing one, or entering it as a context manager,
leaves the database alone; this client closes them all on teardown.
"""
assert isinstance(self._session, FederatedSession)
names = _without_repeats(names)

View file

@ -598,6 +598,27 @@ class TestBorrowedDatabases:
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_entering_a_borrowed_client_reuses_its_database(self, tmp_path):
"""`async with` on a borrowed client is a plausible thing to write.
Opening a second session would leak it, since teardown declines to close
what this client did not open."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(config=config) as rag:
(alpha,) = await rag.clients_for(["alpha"])
borrowed = alpha.store
async with alpha as entered:
assert entered is alpha
assert alpha.store is borrowed, "entry opened a second database"
assert borrowed.db.is_open(), "exit closed a database it borrowed"
assert alpha.store is borrowed
assert not borrowed.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."""