Keep the sandbox's connection open for the owners it hands out
`_documents` took its owner clients inside an ephemeral connection and used them after that connection closed, then stored them for the file reads that follow, so every standalone multi-database read went through a closed database. A sandbox covering a set now retains the client it opened until `close()`. One database has no owners and keeps its connection no longer than the read that opened it, so a write from elsewhere is still visible to the next read.
This commit is contained in:
parent
9210642f65
commit
c8a9b5df50
2 changed files with 40 additions and 0 deletions
|
|
@ -149,6 +149,7 @@ class Sandbox:
|
|||
_doc_chunk_index: dict[str, dict[str, list[str]]]
|
||||
_items_jsonl_cache: dict[str, str]
|
||||
_toc_json_cache: dict[str, str]
|
||||
_opened: "HaikuRAG | None"
|
||||
_pool: AsyncMonty | None
|
||||
_session: AsyncMontySession | None
|
||||
_vfs: OSAccess | None
|
||||
|
|
@ -206,6 +207,7 @@ class Sandbox:
|
|||
self._config = config
|
||||
self._context = context
|
||||
self._rag = rag
|
||||
self._opened = None
|
||||
self._owners = {}
|
||||
self._lock = lock
|
||||
self._search_results = []
|
||||
|
|
@ -239,11 +241,23 @@ class Sandbox:
|
|||
else:
|
||||
yield connection
|
||||
return
|
||||
if self._scope.covers_multiple:
|
||||
yield await self._open_connection()
|
||||
return
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG._covering(self._scope, self._config, read_only=True) as rag:
|
||||
yield rag
|
||||
|
||||
async def _open_connection(self) -> "HaikuRAG":
|
||||
"""Open and retain a federated client for owner-backed reads."""
|
||||
if self._opened is None:
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
self._opened = HaikuRAG._covering(self._scope, self._config, read_only=True)
|
||||
await self._opened.__aenter__()
|
||||
return self._opened
|
||||
|
||||
async def _documents(self) -> "tuple[list[Any], dict[str, HaikuRAG]]":
|
||||
"""Every document in scope, and the client holding each of them.
|
||||
|
||||
|
|
@ -338,6 +352,9 @@ class Sandbox:
|
|||
if self._pool is not None:
|
||||
await self._pool.__aexit__(None, None, None)
|
||||
self._pool = None
|
||||
if self._opened is not None:
|
||||
await self._opened.__aexit__(None, None, None)
|
||||
self._opened = None
|
||||
|
||||
def _build_external_functions(self) -> dict[str, Any]:
|
||||
"""Build async external functions for the Monty interpreter."""
|
||||
|
|
|
|||
|
|
@ -20,6 +20,29 @@ async def _mounted(rag, sources=None):
|
|||
return sandbox, docs, owners
|
||||
|
||||
|
||||
class TestStandaloneAcrossDatabases:
|
||||
"""Without a lent client the sandbox opens its own. The owners it hands out
|
||||
are stored for later file reads, so that connection has to outlive the call
|
||||
that produced them."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_owners_stay_open_for_later_reads(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"])
|
||||
|
||||
sandbox = Sandbox(db_path=None, config=config, context=AnalysisContext())
|
||||
try:
|
||||
_, owners = await sandbox._documents()
|
||||
|
||||
assert len(owners) == 2
|
||||
assert all(owner.store.db.is_open() for owner in owners.values())
|
||||
finally:
|
||||
await sandbox.close()
|
||||
|
||||
assert not any(owner.store.db.is_open() for owner in owners.values())
|
||||
|
||||
|
||||
class TestDocumentsAcrossDatabases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_corpus_covers_every_configured_database(self, tmp_path):
|
||||
|
|
|
|||
Loading…
Reference in a new issue