From 4bc52f0710c7e59afc510bee9b968c60dc25ec46 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Jun 2026 15:56:26 +0300 Subject: [PATCH] Serialize multi-table document writes and bound update version churn --- haiku_rag_slim/haiku/rag/client/documents.py | 132 +++++++++--------- haiku_rag_slim/haiku/rag/store/engine.py | 1 + .../haiku/rag/store/repositories/chunk.py | 75 ++++++---- .../haiku/rag/store/repositories/document.py | 42 +++--- .../rag/store/repositories/document_item.py | 25 ++++ tests/test_client.py | 46 ++++++ tests/test_versioning.py | 10 +- 7 files changed, 218 insertions(+), 113 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 02175f7e..50c87aa4 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -74,30 +74,31 @@ async def _store_document_with_chunks( """ chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder) - versions = await client.store.current_table_versions() + async with client.store._write_lock: + versions = await client.store.current_table_versions() - created_doc = await client.document_repository.create(document) + created_doc = await client.document_repository.create(document) - try: - assert created_doc.id is not None, ( - "Document ID should not be None after creation" - ) - for order, chunk in enumerate(chunks): - chunk.document_id = created_doc.id - chunk.order = order + try: + assert created_doc.id is not None, ( + "Document ID should not be None after creation" + ) + for order, chunk in enumerate(chunks): + chunk.document_id = created_doc.id + chunk.order = order - await client.chunk_repository.create(chunks) + await client.chunk_repository.create(chunks) - items = extract_items(created_doc.id, docling_document) - await client.document_item_repository.create_items(created_doc.id, items) + items = extract_items(created_doc.id, docling_document) + await client.document_item_repository.create_items(created_doc.id, items) - if client._config.storage.auto_vacuum: - client._schedule_vacuum() + if client._config.storage.auto_vacuum: + client._schedule_vacuum() - return created_doc - except Exception: - await client.store.restore_table_versions(versions) - raise + return created_doc + except Exception: + await client.store.restore_table_versions(versions) + raise async def _update_document_with_chunks( @@ -124,36 +125,36 @@ async def _update_document_with_chunks( chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder) - versions = await client.store.current_table_versions() + async with client.store._write_lock: + versions = await client.store.current_table_versions() - await client.chunk_repository.delete_by_document_id(document.id) + try: + 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.create(chunks) + if docling_document is not None: + items = extract_items( + updated_doc.id, + docling_document, + existing_picture_data=existing_picture_data, + ) + await client.document_item_repository.replace_for_document( + updated_doc.id, items + ) - if docling_document is not None: - await client.document_item_repository.delete_by_document_id(updated_doc.id) - items = extract_items( - updated_doc.id, - docling_document, - existing_picture_data=existing_picture_data, - ) - await client.document_item_repository.create_items(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 + except Exception: + await client.store.restore_table_versions(versions) + raise async def create_document( @@ -235,33 +236,36 @@ async def _store_documents_with_chunks( for _, chunks, _ in prepared ] - versions = await client.store.current_table_versions() + async with client.store._write_lock: + versions = await client.store.current_table_versions() - created = await client.document_repository.create([doc for doc, _, _ in prepared]) + created = await client.document_repository.create( + [doc for doc, _, _ in prepared] + ) - try: - all_chunks: list[Chunk] = [] - all_items = [] - for doc, doc_chunks, docling_document in zip( - created, embedded, (d for _, _, d in prepared) - ): - 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) - all_items.extend(extract_items(doc.id, docling_document)) + try: + all_chunks: list[Chunk] = [] + all_items = [] + for doc, doc_chunks, docling_document in zip( + created, embedded, (d for _, _, d in prepared) + ): + 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) + all_items.extend(extract_items(doc.id, docling_document)) - 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 + except Exception: + await client.store.restore_table_versions(versions) + raise async def import_documents( diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 98bcfac4..1c613745 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -345,6 +345,7 @@ class Store: self._skip_validation = skip_validation self._skip_migration_check = skip_migration_check self._vacuum_lock = asyncio.Lock() + self._write_lock = asyncio.Lock() self._is_new_db = False # Check if database exists (for local filesystem only) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 7044066d..d2797d08 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -12,6 +12,7 @@ from lancedb.rerankers import RRFReranker from haiku.rag.store.engine import Store, query_to_pydantic from haiku.rag.store.models.chunk import Chunk, SearchType +from haiku.rag.utils import escape_sql_string logger = logging.getLogger(__name__) @@ -42,6 +43,21 @@ class ChunkRepository: return "\n".join(meta.headings) + "\n" + chunk.content return chunk.content + def _to_record(self, chunk: Chunk, chunk_id: str): + assert chunk.document_id is not None + assert chunk.embedding is not None + return self.store.ChunkRecord( + id=chunk_id, + document_id=chunk.document_id, + content=chunk.content, + content_fts=self._contextualize_content(chunk), + metadata=json.dumps( + {k: v for k, v in chunk.metadata.items() if k != "order"} + ), + order=int(chunk.order), + vector=chunk.embedding, + ) + async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]: """Create one or more chunks in the database. @@ -55,18 +71,7 @@ class ChunkRepository: assert entity.embedding is not None, "Chunk must have an embedding" chunk_id = str(uuid4()) - - chunk_record = self.store.ChunkRecord( - id=chunk_id, - document_id=entity.document_id, - content=entity.content, - content_fts=self._contextualize_content(entity), - metadata=json.dumps( - {k: v for k, v in entity.metadata.items() if k != "order"} - ), - order=int(entity.order), - vector=entity.embedding, - ) + chunk_record = self._to_record(entity, chunk_id) await self.store.chunks_table.add([chunk_record]) @@ -88,19 +93,7 @@ class ChunkRepository: for chunk in chunks: chunk_id = str(uuid4()) - assert chunk.document_id is not None - assert chunk.embedding is not None - chunk_record = self.store.ChunkRecord( - id=chunk_id, - document_id=chunk.document_id, - content=chunk.content, - content_fts=self._contextualize_content(chunk), - metadata=json.dumps( - {k: v for k, v in chunk.metadata.items() if k != "order"} - ), - order=int(chunk.order), - vector=chunk.embedding, - ) + chunk_record = self._to_record(chunk, chunk_id) chunk_records.append(chunk_record) chunk.id = chunk_id @@ -109,6 +102,38 @@ class ChunkRepository: return chunks + async def replace_for_document( + self, document_id: str, chunks: list[Chunk] + ) -> list[Chunk]: + """Replace all chunks for a document with one scoped merge operation.""" + self.store._assert_writable() + + if not chunks: + await self.delete_by_document_id(document_id) + return [] + + for chunk in chunks: + assert chunk.document_id == document_id, ( + "All chunks must belong to the replaced document" + ) + assert chunk.embedding is not None, "All chunks must have embeddings" + + records = [] + for chunk in chunks: + chunk_id = str(uuid4()) + records.append(self._to_record(chunk, chunk_id)) + chunk.id = chunk_id + + safe_id = escape_sql_string(document_id) + await ( + self.store.chunks_table.merge_insert(["document_id", "order"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .when_not_matched_by_source_delete(f"document_id = '{safe_id}'") + .execute(records) + ) + return chunks + async def get_by_id(self, entity_id: str) -> Chunk | None: """Get a chunk by its ID.""" results = await query_to_pydantic( diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index fb06435b..1fec4b72 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -62,7 +62,13 @@ class DocumentRepository: else datetime.now(), ) - def _to_record(self, entity: Document, doc_id: str, now: str) -> DocumentRecord: + def _to_record( + self, + entity: Document, + doc_id: str, + created_at: str, + updated_at: str, + ) -> DocumentRecord: return DocumentRecord( id=doc_id, content=entity.content, @@ -72,8 +78,8 @@ class DocumentRepository: docling_document=entity.docling_document, docling_pages=entity.docling_pages, docling_version=entity.docling_version, - created_at=now, - updated_at=now, + created_at=created_at, + updated_at=updated_at, ) @overload @@ -94,7 +100,9 @@ class DocumentRepository: if isinstance(entity, Document): doc_id = str(uuid4()) now = datetime.now().isoformat() - await self.store.documents_table.add([self._to_record(entity, doc_id, now)]) + await self.store.documents_table.add( + [self._to_record(entity, doc_id, now, now)] + ) entity.id = doc_id entity.created_at = datetime.fromisoformat(now) entity.updated_at = datetime.fromisoformat(now) @@ -109,7 +117,7 @@ class DocumentRepository: records = [] for document in documents: doc_id = str(uuid4()) - records.append(self._to_record(document, doc_id, now)) + records.append(self._to_record(document, doc_id, now, now)) document.id = doc_id document.created_at = created_at document.updated_at = created_at @@ -199,20 +207,16 @@ class DocumentRepository: now = datetime.now().isoformat() entity.updated_at = datetime.fromisoformat(now) - # Update the record - safe_id = escape_sql_string(entity.id) - await self.store.documents_table.update( - { - "content": entity.content, - "uri": entity.uri, - "title": entity.title, - "metadata": json.dumps(entity.metadata), - "docling_document": entity.docling_document, - "docling_pages": entity.docling_pages, - "docling_version": entity.docling_version, - "updated_at": now, - }, - where=f"id = '{safe_id}'", + record = self._to_record( + entity, + entity.id, + entity.created_at.isoformat() if entity.created_at else now, + now, + ) + await ( + self.store.documents_table.merge_insert("id") + .when_matched_update_all() + .execute([record]) ) return entity diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 4e1f5d49..897745cc 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -69,6 +69,31 @@ class DocumentItemRepository: records = [self._to_record(item.document_id, item) for item in items] await self.store.document_items_table.add(records) + async def replace_for_document( + self, document_id: str, items: list[DocumentItem] + ) -> None: + """Replace all items for a document with one scoped merge operation.""" + self.store._assert_writable() + + if not items: + await self.delete_by_document_id(document_id) + return + + for item in items: + assert item.document_id == document_id, ( + "All items must belong to the replaced document" + ) + + safe_id = escape_sql_string(document_id) + records = [self._to_record(document_id, item) for item in items] + await ( + self.store.document_items_table.merge_insert(["document_id", "self_ref"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .when_not_matched_by_source_delete(f"document_id = '{safe_id}'") + .execute(records) + ) + async def get_all_items(self, document_id: str) -> list[DocumentItem]: """Get all items for a document, sorted by position.""" safe_id = escape_sql_string(document_id) diff --git a/tests/test_client.py b/tests/test_client.py index 571a9c6f..83d4595b 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -861,6 +861,52 @@ async def test_client_import_documents_empty(temp_db_path): assert after == before +async def test_client_update_document_replaces_rows_with_bounded_versions( + temp_db_path, +): + """Updating one document should replace stale rows with bounded versions.""" + dim = Config.embeddings.model.vector_dim + + async with HaikuRAG(temp_db_path, create=True) as client: + created = await client.import_document( + _docling_doc("original", "Original body"), + [Chunk(content="Original body", embedding=[0.1] * dim, order=0)], + uri="mem://replace", + title="Replace", + ) + assert created.id is not None + + updated_docling = _docling_doc("updated", "Updated body") + updated_chunks = [ + Chunk(content="Updated body A", embedding=[0.2] * dim, order=0), + Chunk(content="Updated body B", embedding=[0.3] * dim, order=1), + ] + + before = await client.store.current_table_versions() + updated = await client.update_document( + created.id, + docling_document=updated_docling, + chunks=updated_chunks, + ) + after = await client.store.current_table_versions() + + assert updated.id == created.id + assert after["documents"] - before["documents"] == 1 + # Indexed LanceDB tables record one additional physical version for + # merge replacement in 0.30.x. + assert after["chunks"] - before["chunks"] <= 2 + assert after["document_items"] - before["document_items"] <= 2 + + stored_chunks = await client.chunk_repository.get_by_document_id(created.id) + assert [chunk.content for chunk in stored_chunks] == [ + "Updated body A", + "Updated body B", + ] + stored_items = await client.document_item_repository.get_all_items(created.id) + assert len(stored_items) == 1 + assert stored_items[0].text == "Updated body" + + @pytest.mark.vcr() async def test_client_ask(allow_model_requests, temp_db_path): """Test asking questions returns answer and citations (VCR recorded).""" diff --git a/tests/test_versioning.py b/tests/test_versioning.py index a23520ac..8dcda789 100644 --- a/tests/test_versioning.py +++ b/tests/test_versioning.py @@ -38,14 +38,14 @@ async def test_version_rollback_on_update_failure(temp_db_path): base_content = "Base content" created = await client.create_document(content=base_content) - # Patch chunk_repository.create to succeed then fail during update - orig_create = client.chunk_repository.create + # Patch chunk replacement to succeed then fail during update + orig_replace = client.chunk_repository.replace_for_document - async def succeed_then_fail(chunks): - await orig_create(chunks) + async def succeed_then_fail(document_id, chunks): + await orig_replace(document_id, chunks) raise RuntimeError("update fail") - client.chunk_repository.create = succeed_then_fail + client.chunk_repository.replace_for_document = succeed_then_fail # Attempt update with pytest.raises(RuntimeError):