diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 09f475f1..3c86667a 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/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, diff --git a/tests/store/legacy_documents.py b/tests/store/legacy_documents.py index d8653464..f6ba2781 100644 --- a/tests/store/legacy_documents.py +++ b/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") diff --git a/tests/store/test_document_meta_split.py b/tests/store/test_document_meta_split.py index 4d84c85b..69038899 100644 --- a/tests/store/test_document_meta_split.py +++ b/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") diff --git a/tests/store/test_v0_48_0_migration.py b/tests/store/test_v0_48_0_migration.py index 8cb594be..94cdd142 100644 --- a/tests/store/test_v0_48_0_migration.py +++ b/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) diff --git a/tests/store/test_v0_50_0_migration.py b/tests/store/test_v0_50_0_migration.py index f2c8ea3b..f53904d6 100644 --- a/tests/store/test_v0_50_0_migration.py +++ b/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): diff --git a/tests/test_client.py b/tests/test_client.py index b83aed91..b2ebc8a2 100644 --- a/tests/test_client.py +++ b/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)."""