Merge pull request #57 from ggozad/feat/rollback-on-fail

Rollback version if a failure occurs when creating/updating documents

docs/cli.md

@@ -36,8 +36,10 @@ haiku-rag add-src https://example.com/article.html
 ```
 
 !!! note
-    As you add documents to `haiku.rag` the database keeps growing. By default, `lanceDB` supports versioning
+    As you add documents to `haiku.rag` the database keeps growing. By default, LanceDB supports versioning
-    of your data. You can optimize and compact the database by running the [vaccum](#vacuum-optimize-and-cleanup) command.
+    of your data. Create/update operations are atomic‑feeling: if anything fails during chunking or embedding,
+    the database rolls back to the pre‑operation snapshot using LanceDB table versioning. You can optimize and
+    compact the database by running the [vacuum](#vacuum-optimize-and-cleanup) command.
 
 ### Get Document
 

docs/python.md

@@ -109,6 +109,14 @@ await client.vacuum()
 
 This compacts tables and removes historical versions to keep disk usage in check. It’s safe to run anytime, for example after bulk imports or periodically in long‑running apps.
 
+### Atomic Writes and Rollback
+
+Document create and update operations take a snapshot of table versions before any write and automatically roll back to that snapshot if something fails (for example, during chunking or embedding). This restores both the `documents` and `chunks` tables to their pre‑operation state using LanceDB’s table versioning.
+
+- Applies to: `create_document(...)`, `create_document_from_source(...)`, `update_document(...)`, and internal rebuild/update flows.
+- Scope: Both document rows and all associated chunks are rolled back together.
+- Vacuum: Running `vacuum()` later prunes old versions for disk efficiency; rollbacks occur immediately during the failing operation and are not impacted.
+
 ## Searching Documents
 
 The search method performs native hybrid search (vector + full-text) using LanceDB with optional reranking for improved relevance:

src/haiku/rag/store/engine.py

@@ -209,6 +209,21 @@ class Store:
         # LanceDB connections are automatically managed
         pass
 
+    def current_table_versions(self) -> dict[str, int]:
+        """Capture current versions of key tables for rollback using LanceDB's API."""
+        return {
+            "documents": int(self.documents_table.version),
+            "chunks": int(self.chunks_table.version),
+            "settings": int(self.settings_table.version),
+        }
+
+    def restore_table_versions(self, versions: dict[str, int]) -> bool:
+        """Restore tables to the provided versions using LanceDB's API."""
+        self.documents_table.restore(int(versions["documents"]))
+        self.chunks_table.restore(int(versions["chunks"]))
+        self.settings_table.restore(int(versions["settings"]))
+        return True
+
     @property
     def _connection(self):
         """Compatibility property for repositories expecting _connection."""

src/haiku/rag/store/repositories/document.py

@@ -171,44 +171,64 @@ class DocumentRepository:
         chunks: list["Chunk"] | None = None,
     ) -> Document:
         """Create a document with its chunks and embeddings."""
+        # Snapshot table versions for versioned rollback (if supported)
+        versions = self.store.current_table_versions()
+
         # Create the document
         created_doc = await self.create(entity)
 
-        # Create chunks if not provided
+        # Attempt to create chunks; on failure, prefer version rollback
-        if chunks is None:
+        try:
-            assert created_doc.id is not None, (
+            # Create chunks if not provided
-                "Document ID should not be None after creation"
+            if chunks is None:
-            )
+                assert created_doc.id is not None, (
-            await self.chunk_repository.create_chunks_for_document(
+                    "Document ID should not be None after creation"
-                created_doc.id, docling_document
+                )
-            )
+                await self.chunk_repository.create_chunks_for_document(
-        else:
+                    created_doc.id, docling_document
-            # Use provided chunks, set order from list position
+                )
-            assert created_doc.id is not None, (
+            else:
-                "Document ID should not be None after creation"
+                # Use provided chunks, set order from list position
-            )
+                assert created_doc.id is not None, (
-            for order, chunk in enumerate(chunks):
+                    "Document ID should not be None after creation"
-                chunk.document_id = created_doc.id
+                )
-                chunk.metadata["order"] = order
+                for order, chunk in enumerate(chunks):
-                await self.chunk_repository.create(chunk)
+                    chunk.document_id = created_doc.id
+                    chunk.metadata["order"] = order
+                    await self.chunk_repository.create(chunk)
 
-        return created_doc
+            return created_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise
 
     async def _update_with_docling(
         self, entity: Document, docling_document: DoclingDocument
     ) -> Document:
         """Update a document and regenerate its chunks."""
-        # Delete existing chunks
         assert entity.id is not None, "Document ID is required for update"
+
+        # Snapshot table versions for versioned rollback
+        versions = self.store.current_table_versions()
+
+        # Delete existing chunks before writing new ones
         await self.chunk_repository.delete_by_document_id(entity.id)
 
-        # Update the document
+        try:
-        updated_doc = await self.update(entity)
+            # Update the document
+            updated_doc = await self.update(entity)
 
-        # Create new chunks
+            # Create new chunks
-        assert updated_doc.id is not None, "Document ID should not be None after update"
+            assert updated_doc.id is not None, (
-        await self.chunk_repository.create_chunks_for_document(
+                "Document ID should not be None after update"
-            updated_doc.id, docling_document
+            )
-        )
+            await self.chunk_repository.create_chunks_for_document(
+                updated_doc.id, docling_document
+            )
 
-        return updated_doc
+            return updated_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise

tests/test_versioning.py

@@ -0,0 +1,94 @@
+import pytest
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.utils import text_to_docling_document
+
+
+@pytest.mark.asyncio
+async def test_version_rollback_on_create_failure(temp_db_path):
+    store = Store(temp_db_path)
+    repo = DocumentRepository(store)
+
+    # Ensure chunk repository is instantiated and stub embeddings to avoid network
+    dim = repo.chunk_repository.embedder._vector_dim
+
+    async def fake_embed(x):  # type: ignore[no-redef]
+        if isinstance(x, list):
+            return [[0.0] * dim for _ in x]
+        return [0.0] * dim
+
+    repo.chunk_repository.embedder.embed = fake_embed  # type: ignore[assignment]
+
+    # Patch create_chunks_for_document to succeed then fail, triggering rollback
+    orig = repo.chunk_repository.create_chunks_for_document
+
+    async def succeed_then_fail(document_id, dl_doc):  # noqa: ARG001
+        await orig(document_id, dl_doc)
+        raise RuntimeError("boom")
+
+    repo.chunk_repository.create_chunks_for_document = succeed_then_fail  # type: ignore[assignment]
+
+    # Attempt to create document with chunks; expect failure and rollback
+    content = "Hello, rollback!"
+    doc = Document(content=content)
+    dl_doc = text_to_docling_document(content, name="test.md")
+
+    with pytest.raises(RuntimeError):
+        await repo._create_with_docling(doc, dl_doc)
+
+    # State should be restored (no documents/chunks)
+    docs = await repo.list_all()
+    assert len(docs) == 0
+    chunks_repo = ChunkRepository(store)
+    all_chunks = await chunks_repo.list_all()
+    assert len(all_chunks) == 0
+
+
+@pytest.mark.asyncio
+async def test_version_rollback_on_update_failure(temp_db_path):
+    store = Store(temp_db_path)
+    repo = DocumentRepository(store)
+
+    # Stub embeddings to avoid network
+    dim = repo.chunk_repository.embedder._vector_dim
+
+    async def fake_embed(x):  # type: ignore[no-redef]
+        if isinstance(x, list):
+            return [[0.0] * dim for _ in x]
+        return [0.0] * dim
+
+    repo.chunk_repository.embedder.embed = fake_embed  # type: ignore[assignment]
+
+    # Create a valid document first (with real chunking and stubbed embeddings)
+    base_content = "Base content"
+    base_doc = Document(content=base_content)
+    base_dl = text_to_docling_document(base_content, name="base.md")
+    created = await repo._create_with_docling(base_doc, base_dl)
+
+    # Force new chunk creation to fail during update after writing
+    orig = repo.chunk_repository.create_chunks_for_document
+
+    async def succeed_then_fail(document_id, dl_doc):  # noqa: ARG001
+        await orig(document_id, dl_doc)
+        raise RuntimeError("update fail")
+
+    repo.chunk_repository.create_chunks_for_document = succeed_then_fail  # type: ignore[assignment]
+
+    # Attempt update
+    updated_content = "Updated content"
+    created.content = updated_content
+    updated_dl = text_to_docling_document(updated_content, name="updated.md")
+
+    with pytest.raises(RuntimeError):
+        await repo._update_with_docling(created, updated_dl)
+
+    # Content and chunks should remain the original
+    persisted = await repo.get_by_id(created.id)  # type: ignore[arg-type]
+    assert persisted is not None
+    assert persisted.content == base_content
+    chunks_repo = ChunkRepository(store)
+    original_chunks = await chunks_repo.get_by_document_id(created.id)  # type: ignore[arg-type]
+    assert len(original_chunks) > 0