diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 1a65c714..a6d25ad2 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -16,6 +16,7 @@ from urllib.parse import urlparse import httpx from haiku.rag.client.documents import DocumentImport +from haiku.rag.client.session import SingleDatabaseSession from haiku.rag.config import AppConfig, get_config from haiku.rag.converters import get_converter from haiku.rag.embeddings import get_embedder @@ -50,17 +51,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# Throttle for the background auto-vacuum: under sustained ingestion, scheduling -# a compaction on every write degenerates into back-to-back optimize() passes -# that churn the blob-bearing documents table. Fire at most one per interval; a -# final vacuum on close collapses anything throttled here. -_VACUUM_MIN_INTERVAL_S = 300.0 - - -# Failures whose message names the remedy and never the location, so the failing -# database is named alongside it instead of in place of it. -_NAMEABLE_FAILURES = (MigrationRequiredError, ConfigMismatchError, ReadOnlyError) - async def first_found( clients: "list[HaikuRAG]", @@ -154,14 +144,41 @@ class HaikuRAG: self._skip_validation = skip_validation self._create = create self._read_only = read_only - self._vacuum_tasks: set[asyncio.Task] = set() - self._last_vacuum_at: float | None = None - self._vacuum_dirty = False self._requested_sources = sources self._clients: dict[str, HaikuRAG] = {} self._federated: dict[str, str] = {} self._clients_lock = asyncio.Lock() self._source: str | None = None + self._session: SingleDatabaseSession | None = None + + @property + def store(self) -> Store: + """The store of the database this client opened. + + Absent while covering a set: a store has no unambiguous meaning across + several, and `clients_for` reaches the one holding a given database. + """ + if self._session is None: + raise AttributeError("store") + return self._session.store + + @property + def document_repository(self) -> DocumentRepository: + if self._session is None: + raise AttributeError("document_repository") + return self._session.document_repository + + @property + def chunk_repository(self) -> ChunkRepository: + if self._session is None: + raise AttributeError("chunk_repository") + return self._session.chunk_repository + + @property + def document_item_repository(self) -> DocumentItemRepository: + if self._session is None: + raise AttributeError("document_item_repository") + return self._session.document_item_repository @property def is_read_only(self) -> bool: @@ -247,47 +264,14 @@ class HaikuRAG: if db_path is not None: self._db_path = db_path - failure: str | None = None - try: - self.store = Store( - self._db_path, - config=self._config, - skip_validation=self._skip_validation, - create=self._create, - 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 — - # __aexit__ won't run because the `async with` never entered. - try: - await self.store._initialize() - except BaseException: - self.store.close() - raise - except _NAMEABLE_FAILURES as error: - # These say what to run and never where the database is, so the name - # is added to the message rather than replacing it: the operator needs - # both which database failed and what to do about it. - if self._source is None: - raise - raise type(error)(f"database {self._source!r}: {error}") from error - except Exception as error: - # A legacy `uri` or `db_path` client has no name to report instead, so - # its error passes through as it always has. - if self._source is None: - raise - failure = type(error).__name__ - if failure is not None: - # Raised outside the except block on purpose. A database named in - # config is reported by name, and the original spells out the path or - # the bucket: `from None` would only stop it being *printed*, leaving - # it on `__context__` for anything that walks the chain. - raise SourceUnavailableError( - f"database {self._source!r} could not be opened: {failure}" - ) - self.document_repository = DocumentRepository(self.store) - self.chunk_repository = ChunkRepository(self.store) - self.document_item_repository = DocumentItemRepository(self.store) + self._session = await SingleDatabaseSession( + self._db_path, + self._config, + skip_validation=self._skip_validation, + create=self._create, + read_only=self._read_only, + source=self._source, + ).open() return self async def clients_for(self, names: list[str]) -> list["HaikuRAG"]: @@ -410,44 +394,12 @@ class HaikuRAG: await _aclose_quietly(cached, name) async def _await_vacuum_tasks(self) -> None: - """Drain background vacuum work and run a final collapse before teardown. - - Writes schedule a throttled background vacuum; many are debounced or skip - because another vacuum holds the lock. The final pass collapses the - versions those left behind. It runs whenever writes happened - (``_vacuum_dirty``) — not gated on in-flight tasks remaining, since a - debounced run may have scheduled none — but never when nothing was - written (so opening + closing a store still never writes). - """ - if self._vacuum_tasks: - await asyncio.gather(*self._vacuum_tasks, return_exceptions=True) - if not self._vacuum_dirty: - return - self._vacuum_dirty = False - # __aexit__ runs during exception unwinding; a raising vacuum here would - # mask the original exception, so the drain stays best-effort. - try: - await self.store.vacuum() - except Exception: - logger.debug("Final vacuum on close failed", exc_info=True) + if self._session is not None: + await self._session.drain_vacuum() def _schedule_vacuum(self) -> None: - """Schedule a background vacuum, throttled to at most one per - ``_VACUUM_MIN_INTERVAL_S``. Sustained writes would otherwise trigger - back-to-back compaction of the blob-bearing documents table. The throttle - only skips the background task — ``_vacuum_dirty`` still marks that a - final vacuum on close is owed.""" - self._vacuum_dirty = True - now = monotonic() - if ( - self._last_vacuum_at is not None - and now - self._last_vacuum_at < _VACUUM_MIN_INTERVAL_S - ): - return - self._last_vacuum_at = now - task = asyncio.create_task(self.store.vacuum()) - self._vacuum_tasks.add(task) - task.add_done_callback(self._vacuum_tasks.discard) + if self._session is not None: + self._session.schedule_vacuum() # ========================================================================= # Processing Primitives @@ -935,4 +887,5 @@ class HaikuRAG: def close(self): """Close the underlying store connection.""" self._require_one_database("close") - self.store.close() + 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 new file mode 100644 index 00000000..ce10fcfd --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/session.py @@ -0,0 +1,147 @@ +import asyncio +import logging +from pathlib import Path +from time import monotonic + +from haiku.rag.config import AppConfig +from haiku.rag.store.engine import Store +from haiku.rag.store.exceptions import ( + ConfigMismatchError, + MigrationRequiredError, + ReadOnlyError, + SourceUnavailableError, +) +from haiku.rag.store.repositories.chunk import ChunkRepository +from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.store.repositories.document_item import DocumentItemRepository + +logger = logging.getLogger(__name__) + +# Throttle for the background auto-vacuum: under sustained ingestion, scheduling +# a compaction on every write degenerates into back-to-back optimize() passes +# that churn the blob-bearing documents table. Fire at most one per interval; a +# final vacuum on close collapses anything throttled here. +_VACUUM_MIN_INTERVAL_S = 300.0 + + +# Failures whose message names the remedy and never the location, so the failing +# database is named alongside it instead of in place of it. +_NAMEABLE_FAILURES = (MigrationRequiredError, ConfigMismatchError, ReadOnlyError) + + +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. + """ + + def __init__( + self, + db_path: Path | str, + config: AppConfig, + *, + skip_validation: bool = False, + create: bool = False, + read_only: bool = False, + source: str | None = None, + ) -> None: + self._db_path = db_path + self._config = config + 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 + + async def open(self) -> "SingleDatabaseSession": + """Connect, validate, and build the repositories.""" + failure: str | None = None + try: + self.store = Store( + self._db_path, + config=self._config, + skip_validation=self._skip_validation, + create=self._create, + 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 — + # the caller's `async with` never entered, so its exit won't run. + try: + await self.store._initialize() + except BaseException: + self.store.close() + raise + except _NAMEABLE_FAILURES as error: + # These say what to run and never where the database is, so the name + # is added to the message rather than replacing it: the operator needs + # both which database failed and what to do about it. + if self.source is None: + raise + raise type(error)(f"database {self.source!r}: {error}") from error + except Exception as error: + # A legacy `uri` or `db_path` session has no name to report instead, + # so its error passes through as it always has. + if self.source is None: + raise + failure = type(error).__name__ + if failure is not None: + # Raised outside the except block on purpose. A database named in + # config is reported by name, and the original spells out the path or + # the bucket: `from None` would only stop it being *printed*, leaving + # it on `__context__` for anything that walks the chain. + raise SourceUnavailableError( + f"database {self.source!r} could not be opened: {failure}" + ) + self.document_repository = DocumentRepository(self.store) + self.chunk_repository = ChunkRepository(self.store) + self.document_item_repository = DocumentItemRepository(self.store) + return self + + async def drain_vacuum(self) -> None: + """Drain background vacuum work and run a final collapse before teardown. + + Writes schedule a throttled background vacuum; many are debounced or skip + because another vacuum holds the lock. The final pass collapses the + versions those left behind. It runs whenever writes happened + (``_vacuum_dirty``) — not gated on in-flight tasks remaining, since a + debounced run may have scheduled none — but never when nothing was + written (so opening + closing a store still never writes). + """ + if self._vacuum_tasks: + await asyncio.gather(*self._vacuum_tasks, return_exceptions=True) + if not self._vacuum_dirty: + return + self._vacuum_dirty = False + # Teardown runs during exception unwinding; a raising vacuum here would + # mask the original exception, so the drain stays best-effort. + try: + await self.store.vacuum() + except Exception: + logger.debug("Final vacuum on close failed", exc_info=True) + + def schedule_vacuum(self) -> None: + """Schedule a background vacuum, throttled to at most one per + ``_VACUUM_MIN_INTERVAL_S``. Sustained writes would otherwise trigger + back-to-back compaction of the blob-bearing documents table. The throttle + only skips the background task — ``_vacuum_dirty`` still marks that a + final vacuum on close is owed.""" + self._vacuum_dirty = True + now = monotonic() + if ( + self._last_vacuum_at is not None + and now - self._last_vacuum_at < _VACUUM_MIN_INTERVAL_S + ): + return + self._last_vacuum_at = now + task = asyncio.create_task(self.store.vacuum()) + self._vacuum_tasks.add(task) + task.add_done_callback(self._vacuum_tasks.discard) + + def close(self) -> None: + """Close the underlying store connection.""" + self.store.close() diff --git a/tests/test_client.py b/tests/test_client.py index 8636a99a..7e473fab 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1196,10 +1196,10 @@ async def test_delete_marks_vacuum_dirty(temp_db_path): uri="mem://del", ) assert doc.id is not None - client._vacuum_dirty = False # isolate the delete + client._session._vacuum_dirty = False # isolate the delete assert await client.delete_document(doc.id) is True - assert client._vacuum_dirty is True + assert client._session._vacuum_dirty is True async def test_delete_rolls_back_on_partial_failure(temp_db_path, monkeypatch): @@ -1284,9 +1284,9 @@ async def test_delete_missing_id_returns_false_without_vacuum(temp_db_path): """Deleting an id that doesn't exist returns False and owes no vacuum (the existence check is inside the lock, so a no-op delete stays a no-op).""" async with HaikuRAG(temp_db_path, create=True) as client: - client._vacuum_dirty = False + client._session._vacuum_dirty = False assert await client.delete_document("does-not-exist") is False - assert client._vacuum_dirty is False + assert client._session._vacuum_dirty is False @pytest.mark.vcr() @@ -2843,7 +2843,7 @@ async def test_import_documents_schedules_vacuum_per_config(temp_db_path, auto_v ) ] ) - await asyncio.gather(*client._vacuum_tasks) + await asyncio.gather(*client._session._vacuum_tasks) assert vacuum.await_count == (1 if auto_vacuum else 0) diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index f0741747..63946667 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -485,6 +485,23 @@ class TestOperationsThatNeedOneDatabase: with pytest.raises(AmbiguousDatabaseError, match="clients_for"): await rag.vacuum() + @pytest.mark.asyncio + async def test_a_set_has_no_store_of_its_own(self, tmp_path): + """A store and its repositories belong to one database. `clients_for` + reaches the one holding a given database.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + + async with HaikuRAG(config=config) as rag: + for name in ( + "store", + "document_repository", + "chunk_repository", + "document_item_repository", + ): + with pytest.raises(AttributeError, match=name): + getattr(rag, name) + @pytest.mark.asyncio async def test_a_selected_database_is_still_writable(self, tmp_path): """Naming one of the set is how a write picks its database.""" diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py index 6e6de526..57c2f3b4 100644 --- a/tests/test_vacuum_debounce.py +++ b/tests/test_vacuum_debounce.py @@ -2,7 +2,7 @@ import asyncio import pytest -import haiku.rag.client as client_mod +import haiku.rag.client.session as session_mod from haiku.rag.client import HaikuRAG from haiku.rag.client.documents import _refresh_doc_metadata from haiku.rag.config import get_config @@ -23,7 +23,7 @@ async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch): """Rapid writes within the throttle window schedule only one background vacuum; once the interval elapses, a new one is scheduled.""" t = {"now": 1000.0} - monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"]) + monkeypatch.setattr(session_mod, "monotonic", lambda: t["now"]) async with HaikuRAG(temp_db_path, create=True) as client: calls: list[int] = [] @@ -34,13 +34,13 @@ async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch): monkeypatch.setattr(client.store, "vacuum", fake_vacuum) for _ in range(3): - client._schedule_vacuum() - await asyncio.gather(*client._vacuum_tasks) + client._session.schedule_vacuum() + await asyncio.gather(*client._session._vacuum_tasks) assert len(calls) == 1 # debounced within the interval - t["now"] += client_mod._VACUUM_MIN_INTERVAL_S + 1 - client._schedule_vacuum() - await asyncio.gather(*client._vacuum_tasks) + t["now"] += session_mod._VACUUM_MIN_INTERVAL_S + 1 + client._session.schedule_vacuum() + await asyncio.gather(*client._session._vacuum_tasks) assert len(calls) == 2 # interval elapsed -> a new vacuum scheduled @@ -49,7 +49,7 @@ async def test_debounced_writes_still_collapse_on_close(temp_db_path, monkeypatc """Even when scheduled vacuums after the first are debounced, the writes are marked dirty so the close-time drain runs a final collapse.""" t = {"now": 1000.0} - monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"]) + monkeypatch.setattr(session_mod, "monotonic", lambda: t["now"]) calls: list[int] = [] async with HaikuRAG(temp_db_path, create=True) as client: @@ -59,13 +59,13 @@ async def test_debounced_writes_still_collapse_on_close(temp_db_path, monkeypatc monkeypatch.setattr(client.store, "vacuum", fake_vacuum) - client._schedule_vacuum() # schedules the first background pass - client._schedule_vacuum() # debounced (no task) + client._session.schedule_vacuum() # schedules the first background pass + client._session.schedule_vacuum() # debounced (no task) - await client._await_vacuum_tasks() + await client._session.drain_vacuum() # one scheduled background pass + one final collapse on drain assert len(calls) == 2 - assert client._vacuum_dirty is False + assert client._session._vacuum_dirty is False @pytest.mark.asyncio @@ -82,7 +82,7 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path): metadata={"source_revision": "r1"}, ) # Isolate the refresh: the import already scheduled a vacuum. - client._vacuum_dirty = False + client._session._vacuum_dirty = False await _refresh_doc_metadata( client, @@ -91,7 +91,7 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path): user_metadata={}, source_metadata={"source_revision": "r2", "md5": "same"}, ) - assert client._vacuum_dirty is True + assert client._session._vacuum_dirty is True @pytest.mark.asyncio diff --git a/tests/test_versioning.py b/tests/test_versioning.py index 9e8e7953..fe5a0c18 100644 --- a/tests/test_versioning.py +++ b/tests/test_versioning.py @@ -205,7 +205,7 @@ async def test_existing_database_checks_migrations(monkeypatch, temp_db_path): async def _wait_for_background_vacuum(client): """Wait for any in-flight background vacuum tasks to complete.""" - await client._await_vacuum_tasks() + await client._session.drain_vacuum() @pytest.mark.vcr() @@ -438,7 +438,8 @@ async def test_close_suppresses_failing_drain_vacuum(temp_db_path, monkeypatch): raise RuntimeError("vacuum boom") # Writes happened, so close owes a final vacuum — force that drain branch. - client._vacuum_dirty = True + assert client._session is not None + client._session._vacuum_dirty = True monkeypatch.setattr(client.store, "vacuum", boom) # Must not raise despite the drain vacuum erroring.