Merge pull request #559 from ggozad/fix/store-write-transaction

Put multi-table writes behind one transaction boundary

CHANGELOG.md

@@ -14,6 +14,7 @@
 
 ### Changed
 
+- Multi-table writes (document create, update, batch import, cascade delete) go through `Store.write_transaction()`. Rollback restores in `RESTORE_TABLE_ORDER` and is shielded from cancellation, so a cancelled write rolls back instead of committing part of itself. `Store.restore_table_versions()` is removed.
 - Search enrichment, the multimodal reranker's picture fetch, and context expansion each issue a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document. `expand_with_items` takes the items to expand from rather than fetching them, and `DocumentItemRepository` gains `resolve_refs_grouped`, `get_items_in_ranges`, `get_pictures_grouped` and `get_caption_picture_refs_grouped`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns each picture's text alongside its bytes under `with_text`, off by default so the reranker's blob fetch does not read a column it discards.
 - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
 - `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -391,7 +391,7 @@ class HaikuRAG:
         """
         from haiku.rag.client.documents import parent_uri_filter
 
-        async with self.store._write_lock:
+        async with self.store.write_transaction():
             # Resolve existence and collect the subtree under the lock so two
             # concurrent deletes of the same id can't both proceed, and children
             # can't appear or move between collection and deletion. parent_uri
@@ -414,13 +414,8 @@ class HaikuRAG:
             if not ids_to_delete:
                 return False
 
-            versions = await self.store.current_table_versions()
-            try:
-                for doc_id in ids_to_delete:
-                    await self.document_repository.delete(doc_id)
-            except Exception:
-                await self.store.restore_table_versions(versions)
-                raise
+            for doc_id in ids_to_delete:
+                await self.document_repository.delete(doc_id)
 
         if self._config.storage.auto_vacuum:
             self._schedule_vacuum()

haiku_rag_slim/haiku/rag/client/documents.py

@@ -120,9 +120,7 @@ async def _store_document_with_chunks(
     chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
     items = await asyncio.to_thread(extract_items, "", docling_document)
 
-    async with client.store._write_lock:
-        versions = await client.store.current_table_versions()
-
+    async with client.store.write_transaction():
         # A concurrent ingestion of the same URI may have created the document
         # while this one was converting/embedding outside the lock. LanceDB has
         # no unique constraint on `uri`, so re-check under the lock and update in
@@ -133,41 +131,33 @@ async def _store_document_with_chunks(
             else None
         )
 
-        try:
-            if existing is not None:
-                document.id = existing.id
-                document.created_at = existing.created_at
-                stored_doc = await client.document_repository.update(document)
-            else:
-                stored_doc = await client.document_repository.create(document)
+        if existing is not None:
+            document.id = existing.id
+            document.created_at = existing.created_at
+            stored_doc = await client.document_repository.update(document)
+        else:
+            stored_doc = await client.document_repository.create(document)
 
-            assert stored_doc.id is not None, (
-                "Document ID should not be None after storing"
+        assert stored_doc.id is not None, "Document ID should not be None after storing"
+        for order, chunk in enumerate(chunks):
+            chunk.document_id = stored_doc.id
+            chunk.order = order
+        for item in items:
+            item.document_id = stored_doc.id
+
+        if existing is not None:
+            await client.chunk_repository.replace_for_document(stored_doc.id, chunks)
+            await client.document_item_repository.replace_for_document(
+                stored_doc.id, items
             )
-            for order, chunk in enumerate(chunks):
-                chunk.document_id = stored_doc.id
-                chunk.order = order
-            for item in items:
-                item.document_id = stored_doc.id
+        else:
+            await client.chunk_repository.create(chunks)
+            await client.document_item_repository.create_items(stored_doc.id, items)
 
-            if existing is not None:
-                await client.chunk_repository.replace_for_document(
-                    stored_doc.id, chunks
-                )
-                await client.document_item_repository.replace_for_document(
-                    stored_doc.id, items
-                )
-            else:
-                await client.chunk_repository.create(chunks)
-                await client.document_item_repository.create_items(stored_doc.id, items)
+    if client._config.storage.auto_vacuum:
+        client._schedule_vacuum()
 
-            if client._config.storage.auto_vacuum:
-                client._schedule_vacuum()
-
-            return stored_doc
-        except Exception:
-            await client.store.restore_table_versions(versions)
-            raise
+    return stored_doc
 
 
 async def _update_document_with_chunks(
@@ -200,31 +190,25 @@ async def _update_document_with_chunks(
             extract_items, document.id, docling_document, existing_picture_data
         )
 
-    async with client.store._write_lock:
-        versions = await client.store.current_table_versions()
+    async with client.store.write_transaction():
+        updated_doc = await client.document_repository.update(document)
 
-        try:
-            updated_doc = await client.document_repository.update(document)
+        assert updated_doc.id is not None
+        for order, chunk in enumerate(chunks):
+            chunk.document_id = updated_doc.id
+            chunk.order = order
 
-            assert updated_doc.id is not None
-            for order, chunk in enumerate(chunks):
-                chunk.document_id = updated_doc.id
-                chunk.order = order
+        await client.chunk_repository.replace_for_document(updated_doc.id, chunks)
 
-            await client.chunk_repository.replace_for_document(updated_doc.id, chunks)
+        if items is not None:
+            await client.document_item_repository.replace_for_document(
+                updated_doc.id, items
+            )
 
-            if items is not None:
-                await client.document_item_repository.replace_for_document(
-                    updated_doc.id, items
-                )
+    if client._config.storage.auto_vacuum:
+        client._schedule_vacuum()
 
-            if client._config.storage.auto_vacuum:
-                client._schedule_vacuum()
-
-            return updated_doc
-        except Exception:
-            await client.store.restore_table_versions(versions)
-            raise
+    return updated_doc
 
 
 async def create_document(
@@ -315,36 +299,30 @@ async def _store_documents_with_chunks(
 
     all_item_lists = await asyncio.to_thread(_extract_all_items)
 
-    async with client.store._write_lock:
-        versions = await client.store.current_table_versions()
-
+    async with client.store.write_transaction():
         created = await client.document_repository.create(
             [doc for doc, _, _ in prepared]
         )
 
-        try:
-            all_chunks: list[Chunk] = []
-            all_items = []
-            for doc, doc_chunks, item_list in zip(created, embedded, all_item_lists):
-                assert doc.id is not None
-                for order, chunk in enumerate(doc_chunks):
-                    chunk.document_id = doc.id
-                    chunk.order = order
-                all_chunks.extend(doc_chunks)
-                for item in item_list:
-                    item.document_id = doc.id
-                all_items.extend(item_list)
+        all_chunks: list[Chunk] = []
+        all_items = []
+        for doc, doc_chunks, item_list in zip(created, embedded, all_item_lists):
+            assert doc.id is not None
+            for order, chunk in enumerate(doc_chunks):
+                chunk.document_id = doc.id
+                chunk.order = order
+            all_chunks.extend(doc_chunks)
+            for item in item_list:
+                item.document_id = doc.id
+            all_items.extend(item_list)
 
-            await client.chunk_repository.create(all_chunks)
-            await client.document_item_repository.create_all(all_items)
+        await client.chunk_repository.create(all_chunks)
+        await client.document_item_repository.create_all(all_items)
 
-            if client._config.storage.auto_vacuum:
-                client._schedule_vacuum()
+    if client._config.storage.auto_vacuum:
+        client._schedule_vacuum()
 
-            return created
-        except Exception:
-            await client.store.restore_table_versions(versions)
-            raise
+    return created
 
 
 async def import_documents(

haiku_rag_slim/haiku/rag/store/engine.py

@@ -1,7 +1,8 @@
 import asyncio
 import json
 import logging
-from collections.abc import Coroutine
+from collections.abc import AsyncIterator, Coroutine
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
 from enum import Enum
@@ -955,16 +956,37 @@ class Store:
         """Capture current versions of key tables for rollback using LanceDB's API."""
         return {name: await table.version() for name, table in self._tables().items()}
 
-    async def restore_table_versions(self, versions: dict[str, int]) -> bool:
-        """Restore tables to the provided versions using LanceDB's API.
+    @asynccontextmanager
+    async def write_transaction(self) -> AsyncIterator[None]:
+        """Hold the write lock for a multi-table mutation, restoring every table
+        to its pre-mutation version if the mutation fails.
+
+        Rollback follows RESTORE_TABLE_ORDER and a cancellation cannot interrupt
+        it; a cancellation absorbed during rollback is re-delivered. A rollback
+        that itself fails raises with the original failure as its cause.
+
+        In-process coordination only: a writer in another process can commit
+        between the version snapshot and the mutation.
 
         Raises:
             ReadOnlyError: If the store is in read-only mode.
         """
         self._assert_writable()
-        for name, table in self._tables().items():
-            await table.restore(int(versions[name]))
-        return True
+        async with self._write_lock:
+            versions = await self.current_table_versions()
+            try:
+                yield
+            except BaseException as exc:
+                failures, cancelled = await self._rollback_to_snapshot(versions)
+                if failures:
+                    raise RuntimeError(
+                        f"Write failed ({exc!r}) and rollback failed on: "
+                        f"{', '.join(name for name, _ in failures)}. Tables may "
+                        "be left inconsistent."
+                    ) from exc
+                if cancelled and not isinstance(exc, asyncio.CancelledError):
+                    raise asyncio.CancelledError()
+                raise
 
     async def create_tag(self, name: str) -> None:
         """Tag the current version of every table with the given name.

tests/store/test_read_only.py

@@ -110,14 +110,15 @@ class TestStoreReadOnly:
                 await store.recreate_embeddings_table()
 
     @pytest.mark.asyncio
-    async def test_restore_table_versions_raises_when_read_only(self, temp_db_path):
-        """restore_table_versions() raises ReadOnlyError when read_only=True."""
-        async with Store(temp_db_path, create=True) as store:
-            versions = await store.current_table_versions()
+    async def test_write_transaction_raises_when_read_only(self, temp_db_path):
+        """write_transaction() raises ReadOnlyError when read_only=True."""
+        async with Store(temp_db_path, create=True):
+            pass
 
         async with Store(temp_db_path, read_only=True) as store:
             with pytest.raises(ReadOnlyError):
-                await store.restore_table_versions(versions)
+                async with store.write_transaction():
+                    pass  # pragma: no cover - entering already raised
 
 
 class TestDocumentRepositoryReadOnly:

tests/test_versioning.py

@@ -62,6 +62,109 @@ async def test_version_rollback_on_update_failure(temp_db_path):
         assert len(original_chunks) > 0
 
 
+@pytest.mark.vcr()
+async def test_cancellation_mid_write_rolls_back(temp_db_path):
+    """A cancellation between two table writes must roll back, not leave the
+    chunks write committed without its document."""
+    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
+        orig_create = client.chunk_repository.create
+
+        async def succeed_then_cancel(chunks):
+            await orig_create(chunks)
+            raise asyncio.CancelledError()
+
+        client.chunk_repository.create = succeed_then_cancel
+
+        with pytest.raises(asyncio.CancelledError):
+            await client.create_document(content="cancelled mid-write")
+
+        assert await client.list_documents() == []
+        assert await client.chunk_repository.list_all() == []
+
+
+@pytest.mark.vcr()
+async def test_rollback_failure_keeps_the_original_cause(temp_db_path):
+    """A failed rollback must not hide what failed first."""
+    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
+
+        async def failing_restore(versions, *, best_effort=False):
+            return [("chunks", RuntimeError("restore refused"))]
+
+        client.store._restore_tables = failing_restore
+
+        async def boom(chunks):
+            raise RuntimeError("original failure")
+
+        client.chunk_repository.create = boom
+
+        with pytest.raises(RuntimeError, match="rollback failed on: chunks") as excinfo:
+            await client.create_document(content="doomed")
+
+        assert isinstance(excinfo.value.__cause__, RuntimeError)
+        assert str(excinfo.value.__cause__) == "original failure"
+
+
+@pytest.mark.vcr()
+async def test_cancellation_during_rollback_is_redelivered(temp_db_path):
+    """A cancellation arriving while rollback runs cannot cut it short, and is
+    delivered to the caller once the restore has completed."""
+    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
+        rollback_started = asyncio.Event()
+        rollback_finished = asyncio.Event()
+
+        async def slow_restore(versions, *, best_effort=False):
+            rollback_started.set()
+            await asyncio.sleep(0.05)
+            rollback_finished.set()
+            return []
+
+        client.store._restore_tables = slow_restore
+
+        async def boom(chunks):
+            raise RuntimeError("first failure")
+
+        client.chunk_repository.create = boom
+
+        task = asyncio.create_task(client.create_document(content="cancel in rollback"))
+        await rollback_started.wait()
+        task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert rollback_finished.is_set()
+
+
+@pytest.mark.vcr()
+async def test_batch_import_rolls_back_the_documents_write(temp_db_path):
+    """The batch document write is inside the guarded body, so a failure after
+    it lands restores the documents table too."""
+    from haiku.rag.client.documents import DocumentImport
+    from tests.store.test_document_items import _docling_doc_with_picture
+
+    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
+        orig_create = client.document_repository.create
+
+        async def succeed_then_fail(documents):
+            await orig_create(documents)
+            raise RuntimeError("after the documents write")
+
+        client.document_repository.create = succeed_then_fail
+
+        with pytest.raises(RuntimeError, match="after the documents write"):
+            await client.import_documents(
+                [
+                    DocumentImport(
+                        docling_document=_docling_doc_with_picture(),
+                        chunks=[],
+                        uri="test://batch-rollback",
+                    )
+                ]
+            )
+
+        assert await client.store.documents_table.count_rows() == 0
+
+
 async def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
     def fail_if_called(*_args, **_kwargs):
         raise AssertionError("run_pending_upgrades should not be called for new DB")