diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 5dd47305..b1264882 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -96,7 +96,14 @@ class HaikuRAG: read_only=self._read_only, before=self._before, ) - await self.store._initialize() + # 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 self.document_repository = DocumentRepository(self.store) self.chunk_repository = ChunkRepository(self.store) self.document_item_repository = DocumentItemRepository(self.store) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 955c7bd7..73efce18 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -29,7 +29,10 @@ async def rebuild_database( if mode is None: mode = RebuildMode.FULL - # Wait for any background vacuum before destructive table operations + # Wait for any background vacuum before destructive table operations. + # Rebuild drops and recreates tables (+ creates indices); a concurrent + # optimize on the same table fails with "CreateIndex transaction was + # preempted" from lance. await client._await_vacuum_tasks() # Update settings to current config diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index bd1eeb5a..f57dc94d 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -274,7 +274,14 @@ class Store: await self._validate_configuration() async def __aenter__(self): - await self._initialize() + # If _initialize connects to LanceDB but then fails (e.g. migration + # check, config validation), close the connection so it doesn't + # leak — __aexit__ won't run because the `async with` never entered. + try: + await self._initialize() + except BaseException: + self.close() + raise return self async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 9b048472..39301411 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -238,3 +238,62 @@ class TestStoreSkipsPathValidationForRemote: skip_migration_check=True, ) as store: assert store is not None + + +class TestInitFailureCleanup: + @pytest.mark.asyncio + async def test_store_aenter_closes_connection_on_init_failure( + self, temp_db_path, monkeypatch + ): + """If _initialize raises after connect, __aenter__ must close the + AsyncConnection so it doesn't leak (no __aexit__ runs in that case).""" + mock_conn = AsyncMock() + mock_conn.close = lambda: mock_conn.close_calls.append(True) # type: ignore[attr-defined] + mock_conn.close_calls = [] # type: ignore[attr-defined] + + async def fake_connect(*args, **kwargs): + return mock_conn + + async def failing_init_tables(self): + raise RuntimeError("simulated table init failure") + + monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect) + monkeypatch.setattr(Store, "_init_tables", failing_init_tables) + + with pytest.raises(RuntimeError, match="simulated table init failure"): + async with Store(temp_db_path, create=True) as store: + assert store is not None + + assert mock_conn.close_calls == [True], ( + "AsyncConnection.close() was not called on init failure" + ) + + @pytest.mark.asyncio + async def test_client_aenter_closes_store_on_init_failure( + self, temp_db_path, monkeypatch + ): + """HaikuRAG.__aenter__ must close the store if _initialize fails.""" + from haiku.rag.client import HaikuRAG + + close_calls: list[bool] = [] + + original_close = Store.close + + def tracking_close(self): + close_calls.append(True) + original_close(self) + + async def failing_init(self): + # Set db so close() has something to close + self.db = AsyncMock() + self.db.close = lambda: None + raise RuntimeError("simulated initialize failure") + + monkeypatch.setattr(Store, "_initialize", failing_init) + monkeypatch.setattr(Store, "close", tracking_close) + + with pytest.raises(RuntimeError, match="simulated initialize failure"): + async with HaikuRAG(temp_db_path, create=True): + pass + + assert close_calls, "Store.close() was not called when _initialize raised"