Serialize and roll back the multi-table document delete

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -373,18 +373,49 @@ class HaikuRAG:
 
     async def delete_document(self, document_id: str) -> bool:
         """Delete a document by its ID. Cascades to children linked via
-        ``metadata.parent_uri``."""
+        ``metadata.parent_uri``.
+
+        The whole subtree (root + transitive children) is deleted under a single
+        write lock and a single version snapshot, so the cascade is atomic: any
+        failure restores every table to the pre-delete state, and no other write
+        can interleave between deleting a child and its parent.
+        """
         from haiku.rag.client.documents import parent_uri_filter
 
-        doc = await self.get_document_by_id(document_id)
-        if doc is None:
-            return False
-        if doc.uri:
-            children = await self.list_documents(filter=parent_uri_filter(doc.uri))
-            for child in children:
-                if child.id and child.id != document_id:
-                    await self.delete_document(child.id)
-        return await self.document_repository.delete(document_id)
+        async with self.store._write_lock:
+            # 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
+            # links a child to its parent's uri; walk transitively, guarding
+            # against cycles.
+            ids_to_delete: list[str] = []
+            seen: set[str] = set()
+            queue = [await self.get_document_by_id(document_id)]
+            while queue:
+                doc = queue.pop()
+                if doc is None or doc.id is None or doc.id in seen:
+                    continue
+                seen.add(doc.id)
+                ids_to_delete.append(doc.id)
+                if doc.uri:
+                    queue.extend(
+                        await self.list_documents(filter=parent_uri_filter(doc.uri))
+                    )
+
+            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
+
+        if self._config.storage.auto_vacuum:
+            self._schedule_vacuum()
+        return True
 
     async def list_documents(
         self,

tests/store/legacy_documents.py

@@ -1,9 +1,9 @@
-"""Helpers to seed a `documents` table in its pre-0.57 shape.
+"""Helpers to seed a `documents` table in its pre-0.58 shape.
 
-Before the document_meta split (v0.57.0), the `documents` table carried
+Before the document_meta split (v0.58.0), the `documents` table carried
 `uri/title/metadata/created_at/updated_at` alongside the content+blobs. The
 migration-chain tests need to reproduce that legacy layout so the older
-migrations (which read `documents.metadata`, etc.) and v0.57.0 itself have the
+migrations (which read `documents.metadata`, etc.) and v0.58.0 itself have the
 columns they operate on.
 """
 
@@ -17,7 +17,7 @@ from haiku.rag.store.engine import Store
 
 
 class LegacyDocumentRecord(LanceModel):
-    """The pre-0.57 `documents` record (mutable attributes still inline)."""
+    """The pre-0.58 `documents` record (mutable attributes still inline)."""
 
     id: str = Field(default_factory=lambda: str(uuid4()))
     content: str
@@ -45,7 +45,7 @@ def legacy_documents_schema() -> pa.Schema:
 async def seed_legacy_documents(
     store: Store, records: list[LegacyDocumentRecord]
 ) -> None:
-    """Recreate the `documents` table with the pre-0.57 schema and add records,
+    """Recreate the `documents` table with the pre-0.58 schema and add records,
     simulating a database created before the document_meta split."""
     if "documents" in (await store.db.list_tables()).tables:
         await store.db.drop_table("documents")

tests/store/test_document_meta_split.py

@@ -61,7 +61,7 @@ async def test_update_missing_id_does_not_create_ghost(temp_db_path):
 
 @pytest.mark.asyncio
 async def test_opening_legacy_db_raises_migration_without_mutating(temp_db_path):
-    """Opening a pre-0.57 DB (no document_meta) must raise MigrationRequiredError
+    """Opening a pre-0.58 DB (no document_meta) must raise MigrationRequiredError
     up front — in both writable and read-only mode — and must not mutate the DB
     by creating the new table on open."""
     async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
@@ -69,7 +69,7 @@ async def test_opening_legacy_db_raises_migration_without_mutating(temp_db_path)
             store,
             [LegacyDocumentRecord(id="d", content="x", uri="u", metadata="{}")],
         )
-        # A real pre-0.57 DB has no document_meta table.
+        # A real pre-0.58 DB has no document_meta table.
         await store.db.drop_table("document_meta")
         await store.set_haiku_version("0.56.0")
 

tests/store/test_v0_48_0_migration.py

@@ -60,7 +60,7 @@ class TestV0_48_0Migration:
 
         async with Store(temp_db_path, skip_migration_check=True) as store:
             # Apply the migration in isolation: store.migrate() would run the
-            # whole chain (incl. v0.50.0/v0.57.0 which touch documents.metadata,
+            # whole chain (incl. v0.50.0/v0.58.0 which touch documents.metadata,
             # absent from this docling-only fixture).
             await _apply_backfill_heading_hierarchy(store)
 

tests/store/test_v0_50_0_migration.py

@@ -14,10 +14,10 @@ from tests.store.legacy_documents import (
 class TestV0_50_0Migration:
     """v0.50.0 normalises document.metadata to source-agnostic keys.
 
-    Applied in isolation against the pre-0.57 documents schema (metadata still
+    Applied in isolation against the pre-0.58 documents schema (metadata still
     inline), so the assertions read documents.metadata directly. The full chain
-    (where v0.57.0 later relocates metadata to document_meta) is covered by the
-    v0.57.0 migration test.
+    (where v0.58.0 later relocates metadata to document_meta) is covered by the
+    v0.58.0 migration test.
     """
 
     async def test_renames_etag_and_content_type(self, temp_db_path):

tests/test_client.py

@@ -953,6 +953,110 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
         assert fetched.get_docling_document() is not None
 
 
+async def test_delete_marks_vacuum_dirty(temp_db_path):
+    """A delete adds tombstone/table versions, so it must enter the auto-vacuum
+    lifecycle — otherwise a delete-only run closes without a final vacuum."""
+    dim = Config.embeddings.model.vector_dim
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(
+            _docling_doc("d", "body"),
+            [Chunk(content="body", embedding=[0.1] * dim, order=0)],
+            uri="mem://del",
+        )
+        assert doc.id is not None
+        client._vacuum_dirty = False  # isolate the delete
+
+        assert await client.delete_document(doc.id) is True
+        assert client._vacuum_dirty is True
+
+
+async def test_delete_rolls_back_on_partial_failure(temp_db_path, monkeypatch):
+    """A multi-table delete is atomic: if a later table delete fails, the write
+    lock + version restore bring every table back, leaving no orphaned rows."""
+    dim = Config.embeddings.model.vector_dim
+    config = Config.model_copy(deep=True)
+    config.storage.auto_vacuum = False
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
+        doc = await client.import_document(
+            _docling_doc("d", "body"),
+            [Chunk(content="body", embedding=[0.1] * dim, order=0)],
+            uri="mem://del",
+            title="T",
+            metadata={"k": "v"},
+        )
+        assert doc.id is not None
+
+        async def boom(*_a, **_k):
+            raise RuntimeError("meta delete failed")
+
+        # Fail the final step (document_meta) after chunks/items/documents deleted.
+        monkeypatch.setattr(client.store.document_meta_table, "delete", boom)
+        with pytest.raises(RuntimeError, match="meta delete failed"):
+            await client.delete_document(doc.id)
+        monkeypatch.undo()
+
+        # Rollback restored every table — the document is fully intact.
+        restored = await client.get_document_by_id(doc.id)
+        assert restored is not None
+        assert restored.title == "T"
+        assert restored.metadata["k"] == "v"
+        assert restored.content == "body"
+        chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        assert len(chunks) == 1
+
+
+async def test_cascade_delete_is_atomic(temp_db_path, monkeypatch):
+    """Deleting a parent cascades to children under one lock + snapshot. If any
+    delete in the subtree fails, the whole subtree is restored — a child isn't
+    left deleted while its parent survives."""
+    dim = Config.embeddings.model.vector_dim
+    config = Config.model_copy(deep=True)
+    config.storage.auto_vacuum = False
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
+        parent = await client.import_document(
+            _docling_doc("p", "parent"),
+            [Chunk(content="parent", embedding=[0.1] * dim, order=0)],
+            uri="mem://parent",
+        )
+        child = await client.import_document(
+            _docling_doc("c", "child"),
+            [Chunk(content="child", embedding=[0.2] * dim, order=0)],
+            uri="mem://child",
+            metadata={"parent_uri": "mem://parent"},
+        )
+        assert parent.id is not None and child.id is not None
+
+        orig_delete = client.document_repository.delete
+
+        async def delete_failing_on_child(doc_id):
+            if doc_id == child.id:
+                raise RuntimeError("child delete failed")
+            return await orig_delete(doc_id)
+
+        monkeypatch.setattr(
+            client.document_repository, "delete", delete_failing_on_child
+        )
+        with pytest.raises(RuntimeError, match="child delete failed"):
+            await client.delete_document(parent.id)
+        monkeypatch.undo()
+
+        # Atomic: the parent delete was rolled back too — both survive.
+        assert await client.get_document_by_id(parent.id) is not None
+        assert await client.get_document_by_id(child.id) is not None
+        assert await client.count_documents() == 2
+
+
+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
+        assert await client.delete_document("does-not-exist") is False
+        assert client._vacuum_dirty is False
+
+
 @pytest.mark.vcr()
 async def test_client_ask(allow_model_requests, temp_db_path):
     """Test asking questions returns answer and citations (VCR recorded)."""