From 4bc52f0710c7e59afc510bee9b968c60dc25ec46 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Jun 2026 15:56:26 +0300 Subject: [PATCH 1/9] 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): From 3366a6d383e191e27dbe4847dd3d9fdd13a76ee5 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 10 Jun 2026 20:31:43 +0300 Subject: [PATCH 2/9] Split meta document attributes into a document_meta table --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/client/documents.py | 4 +- haiku_rag_slim/haiku/rag/client/rebuild.py | 26 +- haiku_rag_slim/haiku/rag/store/engine.py | 74 ++++- .../haiku/rag/store/repositories/chunk.py | 22 +- .../haiku/rag/store/repositories/document.py | 261 ++++++++++++------ .../haiku/rag/store/upgrades/__init__.py | 4 + .../haiku/rag/store/upgrades/v0_57_0.py | 98 +++++++ tests/store/__init__.py | 0 tests/store/legacy_documents.py | 56 ++++ tests/store/test_document_items.py | 20 +- tests/store/test_document_meta_split.py | 108 ++++++++ tests/store/test_v0_48_0_migration.py | 17 +- tests/store/test_v0_50_0_migration.py | 88 +++--- tests/store/test_v0_57_0_migration.py | 98 +++++++ tests/test_client.py | 46 +++ tests/test_lancedb_connection.py | 2 +- tests/test_rebuild.py | 10 +- 18 files changed, 756 insertions(+), 182 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py create mode 100644 tests/store/__init__.py create mode 100644 tests/store/legacy_documents.py create mode 100644 tests/store/test_document_meta_split.py create mode 100644 tests/store/test_v0_57_0_migration.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 753f3b07..4dfadb23 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Changed + +- Mutable document attributes (`uri`, `title`, `metadata`, `created_at`, `updated_at`) moved from the `documents` table into a new `document_meta` table (1:1 on `document_id`); metadata/title/`source_revision` updates no longer rewrite the docling blobs. Migration `v0_58_0` relocates existing data and runs a one-time `vacuum` to reclaim prior bloat. + ### Fixed - Embedded PDF attachment extension is derived from the attachment filename, not the parent's synthetic `...#attachment=` URI; non-PDF attachments (e.g. `.joboptions`) are no longer misrouted to docling's PDF backend, and unsupported extensions are skipped. diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 50c87aa4..1676eee3 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -323,7 +323,7 @@ async def _refresh_doc_metadata( updated = True if updated: - return await client.document_repository.update(doc) + return await client.document_repository.update_meta(doc) return doc @@ -730,7 +730,7 @@ async def update_document( existing_doc.metadata = metadata if content is None and chunks is None and docling_document is None: - return await client.document_repository.update(existing_doc) + return await client.document_repository.update_meta(existing_doc) if chunks is not None: if docling_document is not None: diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 25070eb1..6ea1986c 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -192,7 +192,7 @@ async def _rebuild_title_only( continue if title is not None: doc.title = title - await client.document_repository.update(doc) + await client.document_repository.update_meta(doc) assert doc.id is not None yield doc.id @@ -490,27 +490,35 @@ async def _flush_rebuild_batch( document. Used by RECHUNK and FULL modes after the chunks table has been cleared. """ - from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.engine import DocumentMetaRecord, DocumentRecord if not documents: return now = datetime.now().isoformat() - # Batch update documents using merge_insert (single LanceDB version) + # Batch update documents and document_meta using merge_insert (one LanceDB + # version per table). Content+blobs go to documents; mutable attributes go + # to document_meta. doc_records = [] + meta_records = [] for doc in documents: assert doc.id is not None doc_records.append( DocumentRecord( id=doc.id, content=doc.content, - uri=doc.uri, - title=doc.title, - metadata=json.dumps(doc.metadata), docling_document=doc.docling_document, docling_pages=doc.docling_pages, docling_version=doc.docling_version, + ) + ) + meta_records.append( + DocumentMetaRecord( + document_id=doc.id, + uri=doc.uri, + title=doc.title, + metadata=json.dumps(doc.metadata), created_at=doc.created_at.isoformat() if doc.created_at else now, updated_at=now, ) @@ -521,6 +529,12 @@ async def _flush_rebuild_batch( .when_matched_update_all() .execute(doc_records) ) + await ( + client.store.document_meta_table.merge_insert("document_id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(meta_records) + ) # Batch create all chunks (single LanceDB version) if chunks: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 1c613745..e988b279 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -77,12 +77,20 @@ async def connect_lancedb( class DocumentRecord(LanceModel): id: str = Field(default_factory=lambda: str(uuid4())) content: str - uri: str | None = None - title: str | None = None - metadata: str = Field(default="{}") docling_document: bytes | None = None docling_pages: bytes | None = None docling_version: str | None = None + + +class DocumentMetaRecord(LanceModel): + """Mutable, lightweight document attributes, kept separate from the + write-once content/blobs in `documents`. Updating these (metadata, title, + source_revision) must not rewrite the multi-MB docling row.""" + + document_id: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") created_at: str = Field(default_factory=lambda: "") updated_at: str = Field(default_factory=lambda: "") @@ -171,7 +179,13 @@ class SettingsRecord(LanceModel): settings: str = Field(default="{}") -REQUIRED_TABLES: tuple[str, ...] = ("documents", "chunks", "document_items", "settings") +REQUIRED_TABLES: tuple[str, ...] = ( + "documents", + "document_meta", + "chunks", + "document_items", + "settings", +) async def get_database_stats(db: lancedb.AsyncConnection) -> dict: @@ -390,19 +404,18 @@ class Store: chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim) - # Initialize tables (creates them if they don't exist) - await self._init_tables() + # Initialize tables (creates them if they don't exist). For an existing + # DB this raises MigrationRequiredError up front when migrations are + # pending, before creating any newly-introduced table. + await self._init_tables(is_new_db) # Checkout tables to historical state if before is specified if self._before is not None: await self._checkout_tables_before(self._before) - # Set version for new databases, check migrations for existing ones - if is_new_db: - if not self._read_only: - await self._set_initial_version() - elif not self._skip_migration_check: - await self._check_migrations() + # Set version for new databases. + if is_new_db and not self._read_only: + await self._set_initial_version() # Validate config compatibility after connection is established if not self._skip_validation: @@ -493,6 +506,7 @@ class Store: retention = timedelta(seconds=retention_seconds) for table in [ self.documents_table, + self.document_meta_table, self.chunks_table, self.document_items_table, self.settings_table, @@ -553,9 +567,23 @@ class Store: settings_repo = SettingsRepository(self) await settings_repo.validate_config_compatibility() - async def _init_tables(self): + async def _init_tables(self, is_new_db: bool): """Initialize database tables (create if they don't exist).""" existing_tables = (await self.db.list_tables()).tables + + # Surface pending migrations BEFORE creating any newly-introduced table. + # Otherwise opening a legacy DB would either mutate it (creating an empty + # document_meta on open) or raise the wrong ReadOnlyError instead of + # telling the user to run `haiku-rag migrate`. The settings table exists + # on any non-new DB, which is all _check_migrations needs. + if ( + not is_new_db + and not self._skip_migration_check + and "settings" in existing_tables + ): + self.settings_table = await self.db.open_table("settings") + await self._check_migrations() + missing_tables = set(REQUIRED_TABLES) - set(existing_tables) if missing_tables and self._read_only: @@ -572,6 +600,22 @@ class Store: "documents", schema=get_documents_arrow_schema() ) + # Create or open document_meta table (mutable attributes kept out of the + # blob-bearing documents row). Indexed by document_id and uri — both are + # hot look-up keys (get_by_id, get_by_uri). + if "document_meta" in existing_tables: + self.document_meta_table = await self.db.open_table("document_meta") + else: + self.document_meta_table = await self.db.create_table( + "document_meta", schema=DocumentMetaRecord + ) + await self.document_meta_table.create_index( + "document_id", config=BTree(), replace=True + ) + await self.document_meta_table.create_index( + "uri", config=BTree(), replace=True + ) + # Create or open chunks table if "chunks" in existing_tables: self.chunks_table = await self.db.open_table("chunks") @@ -748,6 +792,7 @@ class Store: """Capture current versions of key tables for rollback using LanceDB's API.""" return { "documents": await self.documents_table.version(), + "document_meta": await self.document_meta_table.version(), "chunks": await self.chunks_table.version(), "document_items": await self.document_items_table.version(), "settings": await self.settings_table.version(), @@ -761,6 +806,7 @@ class Store: """ self._assert_writable() await self.documents_table.restore(int(versions["documents"])) + await self.document_meta_table.restore(int(versions["document_meta"])) await self.chunks_table.restore(int(versions["chunks"])) await self.document_items_table.restore(int(versions["document_items"])) await self.settings_table.restore(int(versions["settings"])) @@ -786,6 +832,7 @@ class Store: tables = [ ("documents", self.documents_table), + ("document_meta", self.document_meta_table), ("chunks", self.chunks_table), ("document_items", self.document_items_table), ("settings", self.settings_table), @@ -838,6 +885,7 @@ class Store: """ table_map = { "documents": self.documents_table, + "document_meta": self.document_meta_table, "chunks": self.chunks_table, "document_items": self.document_items_table, "settings": self.settings_table, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index d2797d08..09a0f40a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -272,14 +272,14 @@ class ChunkRepository: # filter in pandas, head(limit)) silently under-returned # whenever the top-N window lacked `limit` matching chunks. docs_df = await ( - self.store.documents_table.query() - .select(["id"]) + self.store.document_meta_table.query() + .select(["document_id"]) .where(filter) .to_pandas() ) if docs_df.empty: return [] - id_list = ", ".join(f"'{d}'" for d in docs_df["id"]) + id_list = ", ".join(f"'{d}'" for d in docs_df["document_id"]) chunk_filter = f"document_id IN ({id_list})" if query_vector is not None: @@ -344,11 +344,11 @@ class ChunkRepository: results = await query_to_pydantic(query, self.store.ChunkRecord) - # Get document info (only metadata columns, skip content/docling blobs) + # Get document info from the mutable attributes table doc_rows = await ( - self.store.documents_table.query() - .select(["id", "uri", "title", "metadata"]) - .where(f"id = '{document_id}'") + self.store.document_meta_table.query() + .select(["document_id", "uri", "title", "metadata"]) + .where(f"document_id = '{document_id}'") .limit(1) .to_list() ) @@ -504,14 +504,14 @@ class ChunkRepository: documents_map: dict[str, dict] = {} if document_ids: id_list = "', '".join(document_ids) - where_clause = f"id IN ('{id_list}')" + where_clause = f"document_id IN ('{id_list}')" doc_rows = await ( - self.store.documents_table.query() - .select(["id", "uri", "title", "metadata"]) + self.store.document_meta_table.query() + .select(["document_id", "uri", "title", "metadata"]) .where(where_clause) .to_list() ) - documents_map = {str(row["id"]): row for row in doc_rows} + documents_map = {str(row["document_id"]): row for row in doc_rows} # Build final results with document info chunks_with_scores = [] diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 1fec4b72..8e18e028 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -6,6 +6,7 @@ from uuid import uuid4 from lancedb.index import BTree from haiku.rag.store.engine import ( + DocumentMetaRecord, DocumentRecord, Store, get_documents_arrow_schema, @@ -16,7 +17,20 @@ from haiku.rag.utils import escape_sql_string class DocumentRepository: - """Repository for Document operations.""" + """Repository for Document operations. + + A document is stored across two tables with a strict invariant: every + `documents` row (id, content, docling blobs — write-once) has exactly one + matching `document_meta` row (uri, title, metadata, timestamps — mutable), + keyed by `document_id`. The mutable attributes never live in `documents`, + so metadata/title/source_revision updates (`update_meta`) rewrite only the + small meta row and never the multi-MB docling blob. + + To preserve the invariant: `create` writes meta then documents and deletes + the meta row if the documents write fails; `update_meta` updates matched + rows only (no insert — an insert on a missing id would create a ghost + surfaced by `list_all`/`count`); `delete` removes both rows. + """ def __init__(self, store: Store) -> None: self.store = store @@ -43,45 +57,61 @@ class DocumentRepository: self._document_item_repository = DocumentItemRepository(self.store) return self._document_item_repository - def _record_to_document(self, record: DocumentRecord) -> Document: - """Convert a DocumentRecord to a Document model.""" + def _merge_to_document( + self, doc: DocumentRecord, meta: DocumentMetaRecord | None + ) -> Document: + """Merge a `documents` record (content+blobs) with its `document_meta` + record (uri/title/metadata/timestamps) into a Document.""" + created = meta.created_at if meta else "" + updated = meta.updated_at if meta else "" return Document( - id=record.id, - content=record.content, - uri=record.uri, - title=record.title, - metadata=json.loads(record.metadata), - docling_document=record.docling_document, - docling_pages=record.docling_pages, - docling_version=record.docling_version, - created_at=datetime.fromisoformat(record.created_at) - if record.created_at - else datetime.now(), - updated_at=datetime.fromisoformat(record.updated_at) - if record.updated_at - else datetime.now(), + id=doc.id, + content=doc.content, + uri=meta.uri if meta else None, + title=meta.title if meta else None, + metadata=json.loads(meta.metadata) if meta else {}, + docling_document=doc.docling_document, + docling_pages=doc.docling_pages, + docling_version=doc.docling_version, + created_at=datetime.fromisoformat(created) if created else datetime.now(), + updated_at=datetime.fromisoformat(updated) if updated else datetime.now(), ) - def _to_record( + def _to_documents_record(self, entity: Document, doc_id: str) -> DocumentRecord: + return DocumentRecord( + id=doc_id, + content=entity.content, + docling_document=entity.docling_document, + docling_pages=entity.docling_pages, + docling_version=entity.docling_version, + ) + + def _to_meta_record( self, entity: Document, doc_id: str, created_at: str, updated_at: str, - ) -> DocumentRecord: - return DocumentRecord( - id=doc_id, - content=entity.content, + ) -> DocumentMetaRecord: + return DocumentMetaRecord( + document_id=doc_id, 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, created_at=created_at, updated_at=updated_at, ) + async def _meta_by_id(self, doc_id: str) -> DocumentMetaRecord | None: + safe_id = escape_sql_string(doc_id) + results = await query_to_pydantic( + self.store.document_meta_table.query() + .where(f"document_id = '{safe_id}'") + .limit(1), + DocumentMetaRecord, + ) + return results[0] if results else None + @overload async def create(self, entity: Document) -> Document: ... @@ -97,12 +127,29 @@ class DocumentRepository: """ self.store._assert_writable() + # document_meta is written before documents so the documents row write + # is the commit point: time-travel to any documents version always sees + # the matching (earlier-written) document_meta row. If the documents + # write then fails, delete just the rows we added (not a table-version + # restore, which would clobber a concurrent writer's meta write) so a + # failed create can't leave a ghost row that list_all/count (which read + # document_meta) would surface. 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, now)] + await self.store.document_meta_table.add( + [self._to_meta_record(entity, doc_id, now, now)] ) + try: + await self.store.documents_table.add( + [self._to_documents_record(entity, doc_id)] + ) + except Exception: + safe_id = escape_sql_string(doc_id) + await self.store.document_meta_table.delete( + f"document_id = '{safe_id}'" + ) + raise entity.id = doc_id entity.created_at = datetime.fromisoformat(now) entity.updated_at = datetime.fromisoformat(now) @@ -114,15 +161,25 @@ class DocumentRepository: now = datetime.now().isoformat() created_at = datetime.fromisoformat(now) - records = [] + doc_records = [] + meta_records = [] + doc_ids = [] for document in documents: doc_id = str(uuid4()) - records.append(self._to_record(document, doc_id, now, now)) + doc_ids.append(doc_id) + doc_records.append(self._to_documents_record(document, doc_id)) + meta_records.append(self._to_meta_record(document, doc_id, now, now)) document.id = doc_id document.created_at = created_at document.updated_at = created_at - await self.store.documents_table.add(records) + await self.store.document_meta_table.add(meta_records) + try: + await self.store.documents_table.add(doc_records) + except Exception: + ids = ", ".join(f"'{escape_sql_string(d)}'" for d in doc_ids) + await self.store.document_meta_table.delete(f"document_id IN ({ids})") + raise return documents async def get_by_id(self, entity_id: str) -> Document | None: @@ -136,7 +193,8 @@ class DocumentRepository: if not results: return None - return self._record_to_document(results[0]) + meta = await self._meta_by_id(entity_id) + return self._merge_to_document(results[0], meta) async def get_content(self, entity_id: str) -> str | None: """Get only the text content of a document (skips docling blobs).""" @@ -197,28 +255,43 @@ class DocumentRepository: docling_pages=row.get("docling_pages"), ) - async def update(self, entity: Document) -> Document: - """Update an existing document.""" + async def update_meta(self, entity: Document) -> Document: + """Update only the mutable attributes (uri/title/metadata/updated_at) in + `document_meta`. Does NOT touch the `documents` row, so the multi-MB + docling blob is never rewritten — this is the blob-bloat fix for + metadata/title/source_revision changes.""" self.store._assert_writable() - assert entity.id, "Document ID is required for update" - # Update timestamp now = datetime.now().isoformat() entity.updated_at = datetime.fromisoformat(now) - - record = self._to_record( - entity, - entity.id, - entity.created_at.isoformat() if entity.created_at else now, - now, - ) + created = entity.created_at.isoformat() if entity.created_at else now + record = self._to_meta_record(entity, entity.id, created, now) + # Update only — no insert. Every real document has a document_meta row + # from create()/migration; inserting on no-match would manufacture a + # ghost row (visible to list_all/count) for an id with no documents row. await ( - self.store.documents_table.merge_insert("id") + self.store.document_meta_table.merge_insert("document_id") .when_matched_update_all() .execute([record]) ) + return entity + async def update(self, entity: Document) -> Document: + """Update a document's content+blobs (genuine re-conversion) and its + mutable attributes. Rewrites the `documents` row, so use only when the + docling content actually changed; for metadata/title-only changes use + `update_meta`.""" + self.store._assert_writable() + assert entity.id, "Document ID is required for update" + + doc_record = self._to_documents_record(entity, entity.id) + await ( + self.store.documents_table.merge_insert("id") + .when_matched_update_all() + .execute([doc_record]) + ) + await self.update_meta(entity) return entity async def delete(self, entity_id: str) -> bool: @@ -234,13 +307,12 @@ class DocumentRepository: await self.chunk_repository.delete_by_document_id(entity_id) await self.document_item_repository.delete_by_document_id(entity_id) - # Delete the document + # Delete the document row, its mutable attributes safe_id = escape_sql_string(entity_id) await self.store.documents_table.delete(f"id = '{safe_id}'") + await self.store.document_meta_table.delete(f"document_id = '{safe_id}'") return True - _LISTING_COLUMNS = ["id", "title", "uri", "metadata", "created_at", "updated_at"] - async def list_all( self, limit: int | None = None, @@ -250,20 +322,21 @@ class DocumentRepository: ) -> list[Document]: """List all documents with optional pagination and filtering. + Listing reads `document_meta` (uri/title/metadata/timestamps); the + SQL `filter` is evaluated against those columns. When `include_content` + is set, the content+blob row is loaded from `documents` and merged in. + Args: limit: Maximum number of documents to return. offset: Number of documents to skip. - filter: Optional SQL WHERE clause to filter documents. - include_content: Whether to load content and docling_document. - Defaults to False to avoid loading large blobs for listing. + filter: Optional SQL WHERE clause over document_meta columns. + include_content: Whether to also load content and docling blobs. Returns: List of Document instances matching the criteria. """ - query = self.store.documents_table.query() + query = self.store.document_meta_table.query() - if not include_content: - query = query.select(self._LISTING_COLUMNS) if filter is not None: query = query.where(filter) if offset is not None: @@ -271,50 +344,56 @@ class DocumentRepository: if limit is not None: query = query.limit(limit) - if include_content: - results = await query_to_pydantic(query, DocumentRecord) - return [self._record_to_document(doc) for doc in results] + meta_records = await query_to_pydantic(query, DocumentMetaRecord) - return [ - Document( - id=row["id"], - content="", - title=row.get("title"), - uri=row.get("uri"), - metadata=json.loads(row.get("metadata", "{}")), - created_at=datetime.fromisoformat(row["created_at"]) - if row.get("created_at") - else datetime.now(), - updated_at=datetime.fromisoformat(row["updated_at"]) - if row.get("updated_at") - else datetime.now(), + if not include_content: + return [ + self._merge_to_document(DocumentRecord(id=m.document_id, content=""), m) + for m in meta_records + ] + + documents: list[Document] = [] + for meta in meta_records: + safe_id = escape_sql_string(meta.document_id) + doc_results = await query_to_pydantic( + self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1), + DocumentRecord, ) - for row in await query.to_list() - ] + doc_record = ( + doc_results[0] + if doc_results + else DocumentRecord(id=meta.document_id, content="") + ) + documents.append(self._merge_to_document(doc_record, meta)) + return documents async def count(self, filter: str | None = None) -> int: - """Count documents with optional filtering. - - Args: - filter: Optional SQL WHERE clause to filter documents. - - Returns: - Number of documents matching the criteria. - """ - return await self.store.documents_table.count_rows(filter=filter) + """Count documents with optional filtering (over document_meta columns).""" + return await self.store.document_meta_table.count_rows(filter=filter) async def get_by_uri(self, uri: str) -> Document | None: - """Get a document by its URI.""" + """Get a document by its URI (resolved via document_meta).""" escaped_uri = escape_sql_string(uri) - results = await query_to_pydantic( - self.store.documents_table.query().where(f"uri = '{escaped_uri}'").limit(1), - DocumentRecord, + meta_results = await query_to_pydantic( + self.store.document_meta_table.query() + .where(f"uri = '{escaped_uri}'") + .limit(1), + DocumentMetaRecord, ) - if not results: + if not meta_results: return None - return self._record_to_document(results[0]) + meta = meta_results[0] + safe_id = escape_sql_string(meta.document_id) + doc_results = await query_to_pydantic( + self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1), + DocumentRecord, + ) + if not doc_results: + return None + + return self._merge_to_document(doc_results[0], meta) async def delete_all(self) -> None: """Delete all documents from the database.""" @@ -344,8 +423,18 @@ class DocumentRepository: ) ) if count > 0: - # Drop and recreate table to clear all data + # Drop and recreate tables to clear all data await self.store.db.drop_table("documents") self.store.documents_table = await self.store.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await self.store.db.drop_table("document_meta") + self.store.document_meta_table = await self.store.db.create_table( + "document_meta", schema=DocumentMetaRecord + ) + await self.store.document_meta_table.create_index( + "document_id", config=BTree(), replace=True + ) + await self.store.document_meta_table.create_index( + "uri", config=BTree(), replace=True + ) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 8907f5a2..3c742849 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -93,6 +93,9 @@ from haiku.rag.store.upgrades.v0_48_0 import ( from haiku.rag.store.upgrades.v0_50_0 import ( upgrade_canonical_metadata_keys as upgrade_0_50_0_canonical_metadata_keys, ) +from haiku.rag.store.upgrades.v0_57_0 import ( + upgrade_split_document_meta as upgrade_0_57_0_split_document_meta, +) upgrades.append(upgrade_0_20_0_docling) upgrades.append(upgrade_0_23_1_contextualize) @@ -102,3 +105,4 @@ upgrades.append(upgrade_0_40_0_document_items) upgrades.append(upgrade_0_45_0_extract_picture_bytes) upgrades.append(upgrade_0_48_0_heading_hierarchy) upgrades.append(upgrade_0_50_0_canonical_metadata_keys) +upgrades.append(upgrade_0_57_0_split_document_meta) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py new file mode 100644 index 00000000..2e98c473 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py @@ -0,0 +1,98 @@ +import logging +import shutil + +from haiku.rag.store.engine import DocumentMetaRecord, Store +from haiku.rag.store.upgrades import Upgrade + +logger = logging.getLogger(__name__) + +_LEGACY_COLUMNS = ["uri", "title", "metadata", "created_at", "updated_at"] + + +async def _apply_split_document_meta(store: Store) -> None: + """Move the mutable document attributes (uri/title/metadata/created_at/ + updated_at) out of the blob-bearing `documents` row into `document_meta`. + + After this, a metadata/title/source_revision update writes only the small + `document_meta` row instead of rewriting the multi-MB docling blob. The + blobs (content, docling_document, docling_pages) stay in `documents`. + + The `document_meta` table itself is created on open by `_init_tables`; this + migration populates it and drops the now-relocated columns from `documents`. + Idempotent: a re-run after a partial failure skips already-moved rows and + skips the column drop if the columns are already gone. + """ + schema = await store.documents_table.schema() + present = [c for c in _LEGACY_COLUMNS if c in schema.names] + + if not present: + logger.info("documents already split; nothing to move") + return + + # Resume support: skip documents whose meta row already exists. + existing_meta = { + row["document_id"] + for row in await store.document_meta_table.query() + .select(["document_id"]) + .to_list() + } + + rows = await store.documents_table.query().select(["id", *present]).to_list() + records = [] + for row in rows: + if row["id"] in existing_meta: + continue + meta = row.get("metadata") + records.append( + DocumentMetaRecord( + document_id=row["id"], + uri=row.get("uri"), + title=row.get("title"), + metadata=meta if isinstance(meta, str) and meta else "{}", + created_at=row.get("created_at") or "", + updated_at=row.get("updated_at") or "", + ) + ) + if records: + logger.info( + "Moving attributes for %d document(s) into document_meta", len(records) + ) + await store.document_meta_table.add(records) + + # Drop the relocated columns from documents (a metadata operation — no row + # rewrite, so it is cheap and safe even on a near-full disk). + logger.info("Dropping %s from documents", ", ".join(present)) + await store.documents_table.drop_columns(present) + + # Reclaim the bloat accumulated before the fix (superseded docling rows from + # past metadata churn). retention=0 is safe ONLY because migrate is + # exclusive/single-writer; this must never become normal ingester behaviour. + # Compaction rewrites the live blobs once (transient peak ~= current size + + # one compacted copy), so skip it (best-effort) when free disk cannot cover + # that — the split is already done; the user can run `haiku-rag vacuum`. + try: + # lancedb's .stats() stub claims TableStatistics but returns a plain dict. + stats: dict = await store.documents_table.stats() # type: ignore[assignment] # ty: ignore[invalid-assignment] + live_bytes = int(stats.get("total_bytes", 0)) + except Exception: + live_bytes = 0 + free_bytes = shutil.disk_usage(store.db_path).free + if live_bytes and free_bytes < live_bytes: + logger.warning( + "Skipping post-migration vacuum: need ~%.2f GB free to compact the " + "documents table, have %.2f GB. Run `haiku-rag vacuum` once you have " + "space to reclaim the accumulated bloat.", + live_bytes / 1e9, + free_bytes / 1e9, + ) + return + + logger.info("Vacuuming to reclaim accumulated document bloat") + await store.vacuum(retention_seconds=0) + + +upgrade_split_document_meta = Upgrade( + version="0.57.0", + apply=_apply_split_document_meta, + description="Move mutable document attributes into the document_meta table", +) diff --git a/tests/store/__init__.py b/tests/store/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/store/legacy_documents.py b/tests/store/legacy_documents.py new file mode 100644 index 00000000..d8653464 --- /dev/null +++ b/tests/store/legacy_documents.py @@ -0,0 +1,56 @@ +"""Helpers to seed a `documents` table in its pre-0.57 shape. + +Before the document_meta split (v0.57.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 +columns they operate on. +""" + +from uuid import uuid4 + +import pyarrow as pa +from lancedb.pydantic import LanceModel +from pydantic import Field + +from haiku.rag.store.engine import Store + + +class LegacyDocumentRecord(LanceModel): + """The pre-0.57 `documents` record (mutable attributes still inline).""" + + id: str = Field(default_factory=lambda: str(uuid4())) + content: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + docling_document: bytes | None = None + docling_pages: bytes | None = None + docling_version: str | None = None + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + +def legacy_documents_schema() -> pa.Schema: + base = LegacyDocumentRecord.to_arrow_schema() + large_binary_columns = {"docling_document", "docling_pages"} + return pa.schema( + [ + pa.field(f.name, pa.large_binary()) if f.name in large_binary_columns else f + for f in base + ] + ) + + +async def seed_legacy_documents( + store: Store, records: list[LegacyDocumentRecord] +) -> None: + """Recreate the `documents` table with the pre-0.57 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") + store.documents_table = await store.db.create_table( + "documents", schema=legacy_documents_schema() + ) + if records: + await store.documents_table.add(records) diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index f0164146..45a7c6a4 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -394,6 +394,7 @@ class TestDocumentItemMigration: """Test that the v0.40.0 migration populates items for pre-existing documents.""" from haiku.rag.store.compression import compress_docling_split from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items docling_doc = _make_docling_doc() json_str = docling_doc.model_dump_json() @@ -405,7 +406,6 @@ class TestDocumentItemMigration: doc_record = DocumentRecord( id="test-doc-1", content="test content", - uri="test://doc", docling_document=structure, docling_pages=pages, docling_version=docling_doc.version, @@ -415,12 +415,11 @@ class TestDocumentItemMigration: # Verify no items exist yet assert await store.document_items_table.count_rows() == 0 - # Re-open with skip_migration_check and run migration + # Re-open and apply the v0.40.0 migration in isolation (the full chain + # would also run later migrations that touch documents.metadata, absent + # from this docling-only fixture). async with Store(temp_db_path, skip_migration_check=True) as store: - applied = await store.migrate() - - # Should have applied the v0.40.0 migration - assert any("document_items" in desc for desc in applied) + await _apply_populate_document_items(store) # Items should now exist item_count = await store.document_items_table.count_rows( @@ -442,6 +441,7 @@ class TestDocumentItemMigration: async def test_migration_skips_documents_without_docling(self, temp_db_path): """Test that migration handles documents without docling data.""" from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items async with Store(temp_db_path, create=True, skip_migration_check=True) as store: await store.set_haiku_version("0.39.0") @@ -452,7 +452,7 @@ class TestDocumentItemMigration: await store.documents_table.add([doc_record]) async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() + await _apply_populate_document_items(store) # No items should have been created assert await store.document_items_table.count_rows() == 0 @@ -705,6 +705,7 @@ class TestPictureDataMigrationBackfill: from haiku.rag.store.compression import compress_json, decompress_json from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord + from haiku.rag.store.upgrades.v0_45_0 import _apply_extract_picture_bytes fake_png = b"\x89PNG\r\n\x1a\nlegacy-picture-bytes-for-test" data_uri = "data:image/png;base64," + base64.b64encode(fake_png).decode("ascii") @@ -756,8 +757,9 @@ class TestPictureDataMigrationBackfill: assert "picture_data" not in {f.name for f in schema_before} async with Store(temp_db_path, skip_migration_check=True) as store: - applied = await store.migrate() - assert any("picture" in d.lower() for d in applied) + # Apply the v0.45.0 migration in isolation (the full chain would also + # run later migrations that touch documents.metadata, absent here). + await _apply_extract_picture_bytes(store) # Column was added by the migration schema_after = await store.document_items_table.schema() diff --git a/tests/store/test_document_meta_split.py b/tests/store/test_document_meta_split.py new file mode 100644 index 00000000..4d84c85b --- /dev/null +++ b/tests/store/test_document_meta_split.py @@ -0,0 +1,108 @@ +import lancedb +import pytest + +from haiku.rag.store.engine import Store +from haiku.rag.store.exceptions import MigrationRequiredError +from haiku.rag.store.models.document import Document +from haiku.rag.store.repositories.document import DocumentRepository +from tests.store.legacy_documents import ( + LegacyDocumentRecord, + seed_legacy_documents, +) + + +@pytest.mark.asyncio +async def test_create_rolls_back_meta_when_documents_write_fails(temp_db_path): + """A failed documents write must not leave an orphan document_meta row that + list_all/count would surface (they read document_meta). The rollback must be + targeted — it deletes only the failed row, leaving other documents intact.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + + # A pre-existing document that must survive the failed create's rollback. + good = await repo.create(Document(content="keep", uri="mem://keep")) + assert good.id is not None + + original_add = store.documents_table.add + + async def boom(*_args, **_kwargs): + raise RuntimeError("documents write failed") + + store.documents_table.add = boom + with pytest.raises(RuntimeError, match="documents write failed"): + await repo.create(Document(content="x", uri="mem://ghost")) + store.documents_table.add = original_add + + # The ghost's meta row was deleted; the good document is untouched. + assert await repo.count() == 1 + assert [d.id for d in await repo.list_all()] == [good.id] + assert await store.document_meta_table.count_rows() == 1 + assert await repo.get_by_uri("mem://ghost") is None + fetched = await repo.get_by_id(good.id) + assert fetched is not None and fetched.uri == "mem://keep" + + +@pytest.mark.asyncio +async def test_update_missing_id_does_not_create_ghost(temp_db_path): + """update()/update_meta() for an id with no documents row must not insert a + document_meta row — otherwise it would show up in list_all/count while + get_by_id returns None.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + + await repo.update(Document(id="missing", content="x", uri="u")) + + assert await repo.count() == 0 + assert await repo.list_all() == [] + assert await store.documents_table.count_rows() == 0 + assert await store.document_meta_table.count_rows() == 0 + assert await repo.get_by_id("missing") is None + + +@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 + 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: + await seed_legacy_documents( + store, + [LegacyDocumentRecord(id="d", content="x", uri="u", metadata="{}")], + ) + # A real pre-0.57 DB has no document_meta table. + await store.db.drop_table("document_meta") + await store.set_haiku_version("0.56.0") + + # Writable open: pending migration surfaces before any table creation. + with pytest.raises(MigrationRequiredError): + async with Store(temp_db_path): + pass + + # Read-only open: must also be MigrationRequiredError (not ReadOnlyError). + with pytest.raises(MigrationRequiredError): + async with Store(temp_db_path, read_only=True): + pass + + # The failed opens did not create document_meta. + raw = await lancedb.connect_async(str(temp_db_path)) + assert "document_meta" not in (await raw.list_tables()).tables + + +@pytest.mark.asyncio +async def test_migrate_creates_and_populates_document_meta(temp_db_path): + """The migrate path (skip_migration_check) still creates and fills + document_meta for a legacy DB.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [LegacyDocumentRecord(id="d", content="x", uri="u", metadata="{}")], + ) + await store.db.drop_table("document_meta") + await store.set_haiku_version("0.56.0") + + async with Store(temp_db_path, skip_migration_check=True) as store: + await store.migrate() + assert "document_meta" in (await store.db.list_tables()).tables + repo = DocumentRepository(store) + doc = await repo.get_by_id("d") + assert doc is not None and doc.uri == "u" diff --git a/tests/store/test_v0_48_0_migration.py b/tests/store/test_v0_48_0_migration.py index 7710fa10..8cb594be 100644 --- a/tests/store/test_v0_48_0_migration.py +++ b/tests/store/test_v0_48_0_migration.py @@ -4,6 +4,7 @@ import pytest from haiku.rag.store.compression import compress_docling_split from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store +from haiku.rag.store.upgrades.v0_48_0 import _apply_backfill_heading_hierarchy def _docling_with_levels(): @@ -58,8 +59,10 @@ class TestV0_48_0Migration: ) async with Store(temp_db_path, skip_migration_check=True) as store: - applied = await store.migrate() - assert any("0.48.0" in d for d in applied) + # Apply the migration in isolation: store.migrate() would run the + # whole chain (incl. v0.50.0/v0.57.0 which touch documents.metadata, + # absent from this docling-only fixture). + await _apply_backfill_heading_hierarchy(store) rows = await ( store.document_items_table.query() @@ -107,12 +110,8 @@ class TestV0_48_0Migration: ) async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() - - from haiku.rag.store.upgrades.v0_48_0 import ( - _apply_backfill_heading_hierarchy, - ) - + # Apply twice to prove idempotency (in isolation from the chain). + await _apply_backfill_heading_hierarchy(store) await _apply_backfill_heading_hierarchy(store) rows = await ( @@ -144,7 +143,7 @@ class TestV0_48_0Migration: ) async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() + await _apply_backfill_heading_hierarchy(store) rows = await ( store.document_items_table.query() .where("document_id = 'plain'") diff --git a/tests/store/test_v0_50_0_migration.py b/tests/store/test_v0_50_0_migration.py index d446f691..f2c8ea3b 100644 --- a/tests/store/test_v0_50_0_migration.py +++ b/tests/store/test_v0_50_0_migration.py @@ -2,19 +2,30 @@ import json import pytest -from haiku.rag.store.engine import DocumentRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades.v0_50_0 import _apply_canonical_metadata_keys +from tests.store.legacy_documents import ( + LegacyDocumentRecord, + seed_legacy_documents, +) @pytest.mark.asyncio class TestV0_50_0Migration: - """v0.50.0 normalises document.metadata to source-agnostic keys.""" + """v0.50.0 normalises document.metadata to source-agnostic keys. + + Applied in isolation against the pre-0.57 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. + """ async def test_renames_etag_and_content_type(self, temp_db_path): async with Store(temp_db_path, create=True, skip_migration_check=True) as store: - await store.set_haiku_version("0.48.1") - await store.documents_table.add( + await seed_legacy_documents( + store, [ - DocumentRecord( + LegacyDocumentRecord( id="doc-s3", content="x", uri="s3://b/k", @@ -26,7 +37,7 @@ class TestV0_50_0Migration: } ), ), - DocumentRecord( + LegacyDocumentRecord( id="doc-fs", content="y", uri="file:///tmp/x.md", @@ -37,12 +48,10 @@ class TestV0_50_0Migration: } ), ), - ] + ], ) - async with Store(temp_db_path, skip_migration_check=True) as store: - applied = await store.migrate() - assert any("0.50.0" in d for d in applied) + await _apply_canonical_metadata_keys(store) rows = await store.documents_table.query().to_list() by_id = {r["id"]: json.loads(r["metadata"]) for r in rows} @@ -59,10 +68,10 @@ class TestV0_50_0Migration: async def test_idempotent_on_already_migrated(self, temp_db_path): async with Store(temp_db_path, create=True, skip_migration_check=True) as store: - await store.set_haiku_version("0.48.1") - await store.documents_table.add( + await seed_legacy_documents( + store, [ - DocumentRecord( + LegacyDocumentRecord( id="d", content="x", uri="s3://b/k", @@ -74,11 +83,10 @@ class TestV0_50_0Migration: } ), ) - ] + ], ) - async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() + await _apply_canonical_metadata_keys(store) rows = await store.documents_table.query().to_list() assert json.loads(rows[0]["metadata"]) == { "source_revision": "abc", @@ -91,14 +99,11 @@ class TestV0_50_0Migration: quoted-key form. Values that happen to contain the substring `etag` and composite key names like `my_etag_key` must not get rewritten.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: - await store.set_haiku_version("0.48.1") - await store.documents_table.add( + await seed_legacy_documents( + store, [ - # `etag` only appears as a VALUE — no `etag` key. Either - # the LIKE excludes it (no work), or it pulls it in and - # _normalize_metadata leaves it alone. Either way the row - # must end up unchanged. - DocumentRecord( + # `etag` only appears as a VALUE — no `etag` key. + LegacyDocumentRecord( id="value-only", content="x", uri="u1", @@ -110,8 +115,7 @@ class TestV0_50_0Migration: ), ), # A composite key containing `etag` but not equal to it. - # Must not be rewritten. - DocumentRecord( + LegacyDocumentRecord( id="composite-key", content="x", uri="u2", @@ -122,11 +126,10 @@ class TestV0_50_0Migration: } ), ), - ] + ], ) - async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() + await _apply_canonical_metadata_keys(store) rows = await store.documents_table.query().to_list() by_id = {r["id"]: json.loads(r["metadata"]) for r in rows} @@ -144,30 +147,28 @@ class TestV0_50_0Migration: migration: it's logged and skipped, and well-formed rows alongside it still get rewritten. The bad row's metadata is left exactly as-is.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: - await store.set_haiku_version("0.48.1") - await store.documents_table.add( + await seed_legacy_documents( + store, [ - # Contains the substring `"etag"` so the WHERE LIKE pulls - # it in, but it's not valid JSON — json.loads fails. - DocumentRecord( + # Contains the substring `"etag"` so the WHERE LIKE pulls it + # in, but it's not valid JSON — json.loads fails. + LegacyDocumentRecord( id="bad", content="x", uri="u1", metadata='{"etag": broken', ), - DocumentRecord( + LegacyDocumentRecord( id="good", content="x", uri="u2", metadata=json.dumps({"etag": "abc"}), ), - ] + ], ) - async with Store(temp_db_path, skip_migration_check=True) as store: # Must not raise. - applied = await store.migrate() - assert any("0.50.0" in d for d in applied) + await _apply_canonical_metadata_keys(store) rows = await store.documents_table.query().to_list() by_id = {r["id"]: r["metadata"] for r in rows} @@ -179,10 +180,10 @@ class TestV0_50_0Migration: """If both legacy and canonical keys are present, the canonical wins and the legacy is dropped — defends against partial-migration states.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: - await store.set_haiku_version("0.48.1") - await store.documents_table.add( + await seed_legacy_documents( + store, [ - DocumentRecord( + LegacyDocumentRecord( id="d", content="x", uri="s3://b/k", @@ -195,11 +196,10 @@ class TestV0_50_0Migration: } ), ) - ] + ], ) - async with Store(temp_db_path, skip_migration_check=True) as store: - await store.migrate() + await _apply_canonical_metadata_keys(store) rows = await store.documents_table.query().to_list() meta = json.loads(rows[0]["metadata"]) assert meta == { diff --git a/tests/store/test_v0_57_0_migration.py b/tests/store/test_v0_57_0_migration.py new file mode 100644 index 00000000..bc39122d --- /dev/null +++ b/tests/store/test_v0_57_0_migration.py @@ -0,0 +1,98 @@ +import json + +import pytest + +from haiku.rag.store.engine import Store +from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.store.upgrades.v0_57_0 import _apply_split_document_meta +from tests.store.legacy_documents import ( + LegacyDocumentRecord, + seed_legacy_documents, +) + +_LEGACY_COLUMNS = {"uri", "title", "metadata", "created_at", "updated_at"} + + +@pytest.mark.asyncio +class TestV0_57_0Migration: + """v0.57.0 moves mutable attributes out of the documents row into + document_meta so metadata/title updates stop rewriting the docling blob.""" + + async def test_moves_attributes_and_drops_columns(self, temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [ + LegacyDocumentRecord( + id="doc-1", + content="body one", + uri="s3://b/one", + title="One", + metadata=json.dumps({"source_revision": "r1", "md5": "a"}), + docling_document=b"structure-blob-1", + docling_pages=b"pages-blob-1", + docling_version="1.10.0", + created_at="2026-01-01T00:00:00", + updated_at="2026-01-02T00:00:00", + ) + ], + ) + await store.set_haiku_version("0.56.0") + + async with Store(temp_db_path, skip_migration_check=True) as store: + applied = await store.migrate() + assert any("0.57.0" in d for d in applied) + + # Legacy columns dropped from documents; blobs stay. + doc_names = {f.name for f in await store.documents_table.schema()} + assert _LEGACY_COLUMNS.isdisjoint(doc_names) + assert {"id", "content", "docling_document", "docling_pages"} <= doc_names + + # Attributes landed in document_meta. + meta_rows = await store.document_meta_table.query().to_list() + assert len(meta_rows) == 1 + row = meta_rows[0] + assert row["document_id"] == "doc-1" + assert row["uri"] == "s3://b/one" + assert row["title"] == "One" + assert json.loads(row["metadata"]) == {"source_revision": "r1", "md5": "a"} + + # Full hydration still works (content + metadata + blobs intact). + repo = DocumentRepository(store) + doc = await repo.get_by_id("doc-1") + assert doc is not None + assert doc.content == "body one" + assert doc.uri == "s3://b/one" + assert doc.title == "One" + assert doc.metadata == {"source_revision": "r1", "md5": "a"} + assert doc.docling_document == b"structure-blob-1" + assert doc.docling_pages == b"pages-blob-1" + assert doc.docling_version == "1.10.0" + + # Lookup by uri (resolved via document_meta) works too. + by_uri = await repo.get_by_uri("s3://b/one") + assert by_uri is not None and by_uri.id == "doc-1" + + async def test_idempotent_on_already_split(self, temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [ + LegacyDocumentRecord( + id="doc-1", + content="body", + uri="u1", + metadata=json.dumps({"source_revision": "r1"}), + ) + ], + ) + await store.set_haiku_version("0.56.0") + + async with Store(temp_db_path, skip_migration_check=True) as store: + await store.migrate() + # Re-applying must be a no-op (documents already split). + await _apply_split_document_meta(store) + + meta_rows = await store.document_meta_table.query().to_list() + assert len(meta_rows) == 1 + assert meta_rows[0]["document_id"] == "doc-1" diff --git a/tests/test_client.py b/tests/test_client.py index 83d4595b..b83aed91 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -907,6 +907,52 @@ async def test_client_update_document_replaces_rows_with_bounded_versions( assert stored_items[0].text == "Updated body" +async def test_metadata_only_update_does_not_advance_documents_table(temp_db_path): + """Metadata/title-only updates must not rewrite the heavy documents row. + + This is the blob-bloat fix: source_revision rolling on every ingester sweep + used to rewrite the multi-MB docling row each time. Mutable attributes now + live in document_meta, so the documents table version must stay frozen while + only metadata/title change — and reads must still hydrate the full Document. + """ + 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: + created = await client.import_document( + _docling_doc("doc", "Body text"), + [Chunk(content="Body text", embedding=[0.1] * dim, order=0)], + uri="mem://meta-bloat", + title="Original", + metadata={"source_revision": "rev-0"}, + ) + assert created.id is not None + + docs_v0 = await client.store.documents_table.version() + meta_v0 = await client.store.document_meta_table.version() + + for i in range(1, 6): + await client.update_document( + created.id, + metadata={"source_revision": f"rev-{i}"}, + title=f"Title {i}", + ) + + # The heavy documents table must not advance on metadata-only updates. + assert await client.store.documents_table.version() == docs_v0 + # The light document_meta table absorbs the updates. + assert await client.store.document_meta_table.version() > meta_v0 + + # Reads still hydrate the full document (content + blobs + metadata). + fetched = await client.get_document_by_id(created.id) + assert fetched is not None + assert fetched.metadata["source_revision"] == "rev-5" + assert fetched.title == "Title 5" + assert fetched.content == "Body text" + assert fetched.get_docling_document() is not None + + @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_lancedb_connection.py b/tests/test_lancedb_connection.py index ffc49e52..1d283557 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -267,7 +267,7 @@ class TestInitFailureCleanup: async def fake_connect(*args, **kwargs): return mock_conn - async def failing_init_tables(self): + async def failing_init_tables(self, is_new_db): raise RuntimeError("simulated table init failure") monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect) diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 1bf9ebd0..9649a625 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -6,6 +6,7 @@ from typing import TypedDict import pytest from haiku.rag.client import HaikuRAG, RebuildMode +from haiku.rag.config import Config class ChunkData(TypedDict): @@ -184,7 +185,14 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash( _StagingMarkerRecord, ) - async with HaikuRAG(temp_db_path, create=True) as client: + # auto_vacuum off: this test drops the chunks table by hand to simulate a + # crash, where no background vacuum would be in flight. Leaving it on lets + # create_document's scheduled optimize race the raw drop_table ("Directory + # not empty"). + 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.create_document(content=qa_corpus[0]["document_extracted"]) assert doc.id is not None original_chunks = await client.chunk_repository.get_by_document_id(doc.id) From df8af5429825954d964542cb351c2de967a46247 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 10:29:45 +0300 Subject: [PATCH 3/9] rebase from main --- app/backend/pyproject.toml | 2 +- evaluations/pyproject.toml | 2 +- haiku_rag_slim/haiku/rag/store/upgrades/__init__.py | 6 +++--- .../rag/store/upgrades/{v0_57_0.py => v0_58_0.py} | 2 +- haiku_rag_slim/pyproject.toml | 2 +- pyproject.toml | 10 +++++----- ...0_57_0_migration.py => test_v0_58_0_migration.py} | 12 ++++++------ uv.lock | 6 +++--- 8 files changed, 21 insertions(+), 21 deletions(-) rename haiku_rag_slim/haiku/rag/store/upgrades/{v0_57_0.py => v0_58_0.py} (99%) rename tests/store/{test_v0_57_0_migration.py => test_v0_58_0_migration.py} (92%) diff --git a/app/backend/pyproject.toml b/app/backend/pyproject.toml index 55a90c26..8d447938 100644 --- a/app/backend/pyproject.toml +++ b/app/backend/pyproject.toml @@ -8,7 +8,7 @@ dependencies = [ "uvicorn[standard]>=0.40.0", "pydantic-ai-slim[ag-ui,anthropic,openai]>=1.81.0", "python-dotenv>=1.2.1", - "haiku.rag-slim>=0.57.0", + "haiku.rag-slim>=0.58.0", "logfire[pydantic-ai]>=3.17.0", ] diff --git a/evaluations/pyproject.toml b/evaluations/pyproject.toml index 2924dfca..f983baec 100644 --- a/evaluations/pyproject.toml +++ b/evaluations/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-evals" description = "Benchmarking and evaluation scripts for haiku.rag" -version = "0.57.0" +version = "0.58.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } requires-python = ">=3.12" diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 3c742849..115f9a1c 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -93,8 +93,8 @@ from haiku.rag.store.upgrades.v0_48_0 import ( from haiku.rag.store.upgrades.v0_50_0 import ( upgrade_canonical_metadata_keys as upgrade_0_50_0_canonical_metadata_keys, ) -from haiku.rag.store.upgrades.v0_57_0 import ( - upgrade_split_document_meta as upgrade_0_57_0_split_document_meta, +from haiku.rag.store.upgrades.v0_58_0 import ( + upgrade_split_document_meta as upgrade_0_58_0_split_document_meta, ) upgrades.append(upgrade_0_20_0_docling) @@ -105,4 +105,4 @@ upgrades.append(upgrade_0_40_0_document_items) upgrades.append(upgrade_0_45_0_extract_picture_bytes) upgrades.append(upgrade_0_48_0_heading_hierarchy) upgrades.append(upgrade_0_50_0_canonical_metadata_keys) -upgrades.append(upgrade_0_57_0_split_document_meta) +upgrades.append(upgrade_0_58_0_split_document_meta) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py similarity index 99% rename from haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py rename to haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py index 2e98c473..38924db2 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py @@ -92,7 +92,7 @@ async def _apply_split_document_meta(store: Store) -> None: upgrade_split_document_meta = Upgrade( - version="0.57.0", + version="0.58.0", apply=_apply_split_document_meta, description="Move mutable document attributes into the document_meta table", ) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 83f1a316..33877d95 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-slim" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" -version = "0.57.0" +version = "0.58.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } diff --git a/pyproject.toml b/pyproject.toml index 25506992..02ce95de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling" -version = "0.57.0" +version = "0.58.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } @@ -30,7 +30,7 @@ classifiers = [ ] dependencies = [ - "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.57.0", + "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.58.0", ] [project.scripts] @@ -38,9 +38,9 @@ haiku-rag = "haiku.rag.cli:cli" [project.optional-dependencies] tui = ["textual>=8.2.4"] -s3 = ["haiku.rag-slim[s3]==0.57.0"] -cross-encoder = ["haiku.rag-slim[cross-encoder]==0.57.0"] -ingester = ["haiku.rag-slim[ingester]==0.57.0"] +s3 = ["haiku.rag-slim[s3]==0.58.0"] +cross-encoder = ["haiku.rag-slim[cross-encoder]==0.58.0"] +ingester = ["haiku.rag-slim[ingester]==0.58.0"] [build-system] requires = ["hatchling"] diff --git a/tests/store/test_v0_57_0_migration.py b/tests/store/test_v0_58_0_migration.py similarity index 92% rename from tests/store/test_v0_57_0_migration.py rename to tests/store/test_v0_58_0_migration.py index bc39122d..0d85d294 100644 --- a/tests/store/test_v0_57_0_migration.py +++ b/tests/store/test_v0_58_0_migration.py @@ -4,7 +4,7 @@ import pytest from haiku.rag.store.engine import Store from haiku.rag.store.repositories.document import DocumentRepository -from haiku.rag.store.upgrades.v0_57_0 import _apply_split_document_meta +from haiku.rag.store.upgrades.v0_58_0 import _apply_split_document_meta from tests.store.legacy_documents import ( LegacyDocumentRecord, seed_legacy_documents, @@ -14,8 +14,8 @@ _LEGACY_COLUMNS = {"uri", "title", "metadata", "created_at", "updated_at"} @pytest.mark.asyncio -class TestV0_57_0Migration: - """v0.57.0 moves mutable attributes out of the documents row into +class TestV0_58_0Migration: + """v0.58.0 moves mutable attributes out of the documents row into document_meta so metadata/title updates stop rewriting the docling blob.""" async def test_moves_attributes_and_drops_columns(self, temp_db_path): @@ -37,11 +37,11 @@ class TestV0_57_0Migration: ) ], ) - await store.set_haiku_version("0.56.0") + await store.set_haiku_version("0.57.0") async with Store(temp_db_path, skip_migration_check=True) as store: applied = await store.migrate() - assert any("0.57.0" in d for d in applied) + assert any("0.58.0" in d for d in applied) # Legacy columns dropped from documents; blobs stay. doc_names = {f.name for f in await store.documents_table.schema()} @@ -86,7 +86,7 @@ class TestV0_57_0Migration: ) ], ) - await store.set_haiku_version("0.56.0") + await store.set_haiku_version("0.57.0") async with Store(temp_db_path, skip_migration_check=True) as store: await store.migrate() diff --git a/uv.lock b/uv.lock index a3983353..7a4be853 100644 --- a/uv.lock +++ b/uv.lock @@ -1564,7 +1564,7 @@ wheels = [ [[package]] name = "haiku-rag" -version = "0.57.0" +version = "0.58.0" source = { editable = "." } dependencies = [ { name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] }, @@ -1631,7 +1631,7 @@ dev = [ [[package]] name = "haiku-rag-evals" -version = "0.57.0" +version = "0.58.0" source = { editable = "evaluations" } dependencies = [ { name = "datasets" }, @@ -1654,7 +1654,7 @@ requires-dist = [ [[package]] name = "haiku-rag-slim" -version = "0.57.0" +version = "0.58.0" source = { editable = "haiku_rag_slim" } dependencies = [ { name = "docling-core" }, From 5d5d87d44c41dbfd322b77f8eecae1df36ba0c2d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 11:01:31 +0300 Subject: [PATCH 4/9] Throttle background auto-vacuum to at most once per 5 minutes --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/client/__init__.py | 42 ++++++++++++---- tests/test_vacuum_debounce.py | 56 +++++++++++++++++++++ tests/test_versioning.py | 6 +-- 4 files changed, 92 insertions(+), 13 deletions(-) create mode 100644 tests/test_vacuum_debounce.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4dfadb23..549bee08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Changed - Mutable document attributes (`uri`, `title`, `metadata`, `created_at`, `updated_at`) moved from the `documents` table into a new `document_meta` table (1:1 on `document_id`); metadata/title/`source_revision` updates no longer rewrite the docling blobs. Migration `v0_58_0` relocates existing data and runs a one-time `vacuum` to reclaim prior bloat. +- Background auto-vacuum is throttled to at most once per 5 minutes; a final vacuum on close collapses any throttled writes. Sustained ingestion no longer triggers back-to-back compaction of the `documents` table. ### Fixed diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 33d422a5..09f475f1 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -9,6 +9,7 @@ from datetime import datetime from enum import Enum from functools import cached_property from pathlib import Path +from time import monotonic from typing import TYPE_CHECKING, overload from urllib.parse import urlparse @@ -40,6 +41,12 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# Throttle for the background auto-vacuum: under sustained ingestion, scheduling +# a compaction on every write degenerates into back-to-back optimize() passes +# that churn the blob-bearing documents table. Fire at most one per interval; a +# final vacuum on close collapses anything throttled here. +_VACUUM_MIN_INTERVAL_S = 300.0 + class RebuildMode(Enum): """Mode for rebuilding the database.""" @@ -87,6 +94,8 @@ class HaikuRAG: self._read_only = read_only self._before = before self._vacuum_tasks: set[asyncio.Task] = set() + self._last_vacuum_at: float | None = None + self._vacuum_dirty = False @property def is_read_only(self) -> bool: @@ -137,17 +146,20 @@ class HaikuRAG: return False async def _await_vacuum_tasks(self) -> None: - """Drain background vacuum work before tearing down the connection. + """Drain background vacuum work and run a final collapse before teardown. - Each create_document / update_document can schedule its own vacuum task; - all must be awaited, not just the most recently scheduled one. Vacuum - skips when another is already running, so the cleanup for the final - writes may have been a no-op. Run one more pass once the in-flight tasks - are done to collapse versions created after the last vacuum took the lock. + Writes schedule a throttled background vacuum; many are debounced or skip + because another vacuum holds the lock. The final pass collapses the + versions those left behind. It runs whenever writes happened + (``_vacuum_dirty``) — not gated on in-flight tasks remaining, since a + debounced run may have scheduled none — but never when nothing was + written (so opening + closing a store still never writes). """ - if not self._vacuum_tasks: + if self._vacuum_tasks: + await asyncio.gather(*self._vacuum_tasks, return_exceptions=True) + if not self._vacuum_dirty: return - await asyncio.gather(*self._vacuum_tasks, return_exceptions=True) + self._vacuum_dirty = False # __aexit__ runs during exception unwinding; a raising vacuum here would # mask the original exception, so the drain stays best-effort. try: @@ -156,7 +168,19 @@ class HaikuRAG: logger.debug("Final vacuum on close failed", exc_info=True) def _schedule_vacuum(self) -> None: - """Schedule a background vacuum and track the task for later awaiting.""" + """Schedule a background vacuum, throttled to at most one per + ``_VACUUM_MIN_INTERVAL_S``. Sustained writes would otherwise trigger + back-to-back compaction of the blob-bearing documents table. The throttle + only skips the background task — ``_vacuum_dirty`` still marks that a + final vacuum on close is owed.""" + self._vacuum_dirty = True + now = monotonic() + if ( + self._last_vacuum_at is not None + and now - self._last_vacuum_at < _VACUUM_MIN_INTERVAL_S + ): + return + self._last_vacuum_at = now task = asyncio.create_task(self.store.vacuum()) self._vacuum_tasks.add(task) task.add_done_callback(self._vacuum_tasks.discard) diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py new file mode 100644 index 00000000..c2c8ccc5 --- /dev/null +++ b/tests/test_vacuum_debounce.py @@ -0,0 +1,56 @@ +import asyncio + +import pytest + +import haiku.rag.client as client_mod +from haiku.rag.client import HaikuRAG + + +@pytest.mark.asyncio +async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch): + """Rapid writes within the throttle window schedule only one background + vacuum; once the interval elapses, a new one is scheduled.""" + t = {"now": 1000.0} + monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"]) + + async with HaikuRAG(temp_db_path, create=True) as client: + calls: list[int] = [] + + async def fake_vacuum(*_a, **_k): + calls.append(1) + + monkeypatch.setattr(client.store, "vacuum", fake_vacuum) + + for _ in range(3): + client._schedule_vacuum() + await asyncio.gather(*client._vacuum_tasks) + assert len(calls) == 1 # debounced within the interval + + t["now"] += client_mod._VACUUM_MIN_INTERVAL_S + 1 + client._schedule_vacuum() + await asyncio.gather(*client._vacuum_tasks) + assert len(calls) == 2 # interval elapsed -> a new vacuum scheduled + + +@pytest.mark.asyncio +async def test_debounced_writes_still_collapse_on_close(temp_db_path, monkeypatch): + """Even when scheduled vacuums after the first are debounced, the writes are + marked dirty so the close-time drain runs a final collapse.""" + t = {"now": 1000.0} + monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"]) + calls: list[int] = [] + + async with HaikuRAG(temp_db_path, create=True) as client: + + async def fake_vacuum(*_a, **_k): + calls.append(1) + + monkeypatch.setattr(client.store, "vacuum", fake_vacuum) + + client._schedule_vacuum() # schedules the first background pass + client._schedule_vacuum() # debounced (no task) + + await client._await_vacuum_tasks() + # one scheduled background pass + one final collapse on drain + assert len(calls) == 2 + assert client._vacuum_dirty is False diff --git a/tests/test_versioning.py b/tests/test_versioning.py index 8dcda789..07b64a3f 100644 --- a/tests/test_versioning.py +++ b/tests/test_versioning.py @@ -334,10 +334,8 @@ async def test_close_suppresses_failing_drain_vacuum(temp_db_path, monkeypatch): calls.append(1) raise RuntimeError("vacuum boom") - # A finished task in the set forces the drain branch to run. - task = asyncio.create_task(asyncio.sleep(0)) - await task - client._vacuum_tasks.add(task) + # Writes happened, so close owes a final vacuum — force that drain branch. + client._vacuum_dirty = True monkeypatch.setattr(client.store, "vacuum", boom) # Must not raise despite the drain vacuum erroring. From f2a5ac42467f1b5820f82d632392a5a662f527f1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 11:11:42 +0300 Subject: [PATCH 5/9] Throttle background auto-vacuum to at most once per 5 minutes --- haiku_rag_slim/haiku/rag/client/documents.py | 12 +++++-- tests/test_vacuum_debounce.py | 38 ++++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 1676eee3..67def6d9 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -323,7 +323,12 @@ async def _refresh_doc_metadata( updated = True if updated: - return await client.document_repository.update_meta(doc) + result = await client.document_repository.update_meta(doc) + # Reclaim the document_meta churn from rolling source_revision sweeps. + # The vacuum is debounced, and document_meta is tiny, so this is cheap. + if client._config.storage.auto_vacuum: + client._schedule_vacuum() + return result return doc @@ -730,7 +735,10 @@ async def update_document( existing_doc.metadata = metadata if content is None and chunks is None and docling_document is None: - return await client.document_repository.update_meta(existing_doc) + updated = await client.document_repository.update_meta(existing_doc) + if client._config.storage.auto_vacuum: + client._schedule_vacuum() + return updated if chunks is not None: if docling_document is not None: diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py index c2c8ccc5..2b4176f8 100644 --- a/tests/test_vacuum_debounce.py +++ b/tests/test_vacuum_debounce.py @@ -4,6 +4,18 @@ import pytest import haiku.rag.client as client_mod from haiku.rag.client import HaikuRAG +from haiku.rag.client.documents import _refresh_doc_metadata +from haiku.rag.config import Config +from haiku.rag.store.models.chunk import Chunk + + +def _docling_doc(name: str, text: str): + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + doc = DoclingDocument(name=name) + doc.add_text(label=DocItemLabel.TEXT, text=text) + return doc @pytest.mark.asyncio @@ -54,3 +66,29 @@ async def test_debounced_writes_still_collapse_on_close(temp_db_path, monkeypatc # one scheduled background pass + one final collapse on drain assert len(calls) == 2 assert client._vacuum_dirty is False + + +@pytest.mark.asyncio +async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path): + """A source re-sweep that only rolls source_revision (MD5/revision + short-circuit) writes document_meta and must still schedule the (debounced) + vacuum, so that tiny churn gets reclaimed instead of accumulating.""" + 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://sweep", + metadata={"source_revision": "r1"}, + ) + # Isolate the refresh: the import already scheduled a vacuum. + client._vacuum_dirty = False + + await _refresh_doc_metadata( + client, + doc, + title=None, + user_metadata={}, + source_metadata={"source_revision": "r2", "md5": "same"}, + ) + assert client._vacuum_dirty is True From 2b8b9477b7e266a79b2948b49602cf85126939dd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 11:17:43 +0300 Subject: [PATCH 6/9] Serialize and roll back the multi-table document delete --- haiku_rag_slim/haiku/rag/client/__init__.py | 51 ++++++++-- tests/store/legacy_documents.py | 10 +- tests/store/test_document_meta_split.py | 4 +- tests/store/test_v0_48_0_migration.py | 2 +- tests/store/test_v0_50_0_migration.py | 6 +- tests/test_client.py | 104 ++++++++++++++++++++ 6 files changed, 156 insertions(+), 21 deletions(-) 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).""" From bfc51b8b281e7e3570170e2712fa4e1806cfa40c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 12:27:44 +0300 Subject: [PATCH 7/9] urface document_meta in info, history, and the inspector --- haiku_rag_slim/haiku/rag/app.py | 15 +++++++++++++-- haiku_rag_slim/haiku/rag/cli.py | 2 +- .../haiku/rag/inspector/widgets/info_modal.py | 10 ++++++++++ haiku_rag_slim/haiku/rag/store/engine.py | 5 +++-- 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 98be2d3f..6d39fbb1 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -107,7 +107,7 @@ class HaikuRAGApp: # pragma: no cover # Per-table row counts and sizes. Missing required tables are # reported as "absent" rather than raising. - for name in ("documents", "chunks", "document_items"): + for name in ("documents", "document_meta", "chunks", "document_items"): entry = tables[name] if entry.exists: self.console.print( @@ -155,6 +155,11 @@ class HaikuRAGApp: # pragma: no cover f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: " f"{tables['documents'].num_versions}" ) + if tables["document_meta"].exists: + self.console.print( + f" [repr.attrib_name]versions (document_meta)[/repr.attrib_name]: " + f"{tables['document_meta'].num_versions}" + ) if tables["chunks"].exists: self.console.print( f" [repr.attrib_name]versions (chunks)[/repr.attrib_name]: " @@ -215,7 +220,13 @@ class HaikuRAGApp: # pragma: no cover skip_migration_check=True, before=self.before, ) as store: - tables = ["documents", "chunks", "settings"] + tables = [ + "documents", + "document_meta", + "chunks", + "document_items", + "settings", + ] if table: if table not in tables: self.console.print( diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 9e7e21ab..3c02ddc4 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -597,7 +597,7 @@ def history( # pragma: no cover None, "--table", "-t", - help="Specific table to show history for (documents, chunks, settings)", + help="Specific table to show history for (documents, document_meta, chunks, document_items, settings)", ), limit: int | None = typer.Option( None, diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index 0da65169..539b21ca 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -119,6 +119,10 @@ class InfoModal(ModalScreen): num_docs = stats["documents"].get("num_rows", 0) doc_bytes = stats["documents"].get("total_bytes", 0) + num_meta = stats["document_meta"].get("num_rows", 0) + meta_bytes = stats["document_meta"].get("total_bytes", 0) + meta_versions = stats["document_meta"].get("num_versions", 0) + num_chunks = stats["chunks"].get("num_rows", 0) chunk_bytes = stats["chunks"].get("total_bytes", 0) @@ -148,6 +152,9 @@ class InfoModal(ModalScreen): lines.append( f"[bold $accent]documents[/bold $accent]: {num_docs} ({format_bytes(doc_bytes)})" ) + lines.append( + f"[bold $accent]document_meta[/bold $accent]: {num_meta} ({format_bytes(meta_bytes)})" + ) lines.append( f"[bold $accent]chunks[/bold $accent]: {num_chunks} ({format_bytes(chunk_bytes)})" ) @@ -180,6 +187,9 @@ class InfoModal(ModalScreen): lines.append( f"[bold $accent]versions (documents)[/bold $accent]: {doc_versions}" ) + lines.append( + f"[bold $accent]versions (document_meta)[/bold $accent]: {meta_versions}" + ) lines.append( f"[bold $accent]versions (chunks)[/bold $accent]: {chunk_versions}" ) diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index e988b279..860afd6d 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -309,7 +309,7 @@ async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo total_bytes=stats[name].get("total_bytes", 0), num_versions=stats[name].get("num_versions", 0), ) - for name in ("documents", "chunks", "document_items") + for name in ("documents", "document_meta", "chunks", "document_items") ] vector_index = VectorIndexInfo() @@ -878,7 +878,8 @@ class Store: """List version history for a table. Args: - table_name: Name of the table ("documents", "chunks", or "settings") + table_name: Name of the table ("documents", "document_meta", + "chunks", "document_items", or "settings") Returns: List of version info dicts with "version" and "timestamp" keys From 6c04dd5881d868b60458c020c676b94311fa26c6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 12:37:32 +0300 Subject: [PATCH 8/9] docs --- docs/cli.md | 6 +++--- docs/configuration/storage.md | 2 +- docs/python.md | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 7fe283ad..037d544b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -275,9 +275,9 @@ Shows: - path to the database - stored haiku.rag version (from settings) - embeddings provider/model and vector dimension -- number of documents and chunks (with storage sizes) +- per-table row counts and storage sizes (documents, document_meta, chunks, document_items) - vector index status (exists/not created, indexed/unindexed chunks) -- table versions per table (documents, chunks) +- table versions per table (documents, document_meta, chunks) At the end, a separate "Versions" section lists runtime package versions: - haiku.rag @@ -401,7 +401,7 @@ Reduce disk usage by optimizing and pruning old table versions across all tables haiku-rag vacuum ``` -**Automatic Cleanup:** Vacuum runs automatically in the background after document operations. By default, it removes versions older than 1 day (configurable via `storage.vacuum_retention_seconds`), preserving recent versions for concurrent connections. Manual vacuum can be useful for cleanup after bulk operations or to free disk space immediately. +**Automatic Cleanup:** Vacuum runs automatically in the background after document operations, throttled to at most once every 5 minutes so sustained ingestion does not trigger continuous compaction (a final vacuum runs when the client closes). By default, it removes versions older than 1 day (configurable via `storage.vacuum_retention_seconds`), preserving recent versions for concurrent connections. Manual vacuum can be useful for cleanup after bulk operations or to free disk space immediately. ## MCP Server diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 380fdc07..40e4d66a 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -12,7 +12,7 @@ storage: ``` - **data_dir**: Directory for local database storage. When empty, uses platform-specific default locations -- **auto_vacuum**: When enabled (default), automatically runs vacuum after document create/update operations and database rebuilds. Set to `false` to disable automatic vacuuming and rely on manual `haiku-rag vacuum` commands only. Disabling can help avoid potential crashes in high-concurrency scenarios +- **auto_vacuum**: When enabled (default), automatically runs vacuum after document create/update/delete operations and database rebuilds. Background vacuums are throttled to at most one every 5 minutes, so sustained ingestion does not trigger continuous compaction, and a final vacuum runs when the client closes. Set to `false` to disable automatic vacuuming and rely on manual `haiku-rag vacuum` commands only. Disabling can help avoid potential crashes in high-concurrency scenarios - **vacuum_retention_seconds**: When vacuum runs, old table versions older than this threshold are removed. Default: 86400 seconds (1 day). Set to 0 for aggressive cleanup (removes all old versions immediately) !!! warning "Vacuum Retention Threshold" diff --git a/docs/python.md b/docs/python.md index c4cc9f0a..cc028fb1 100644 --- a/docs/python.md +++ b/docs/python.md @@ -396,7 +396,7 @@ The `docling_document` provides rich metadata for visual grounding, page numbers ### Batch Import -Each `create_document*` / `import_document` call writes new versions of the `documents`, `chunks`, and `document_items` tables. Ingesting many documents in a loop therefore creates a table version per document. Use `import_documents()` to write the whole batch in a single version per table: +Each `create_document*` / `import_document` call writes new versions of the `documents`, `document_meta`, `chunks`, and `document_items` tables. Ingesting many documents in a loop therefore creates a table version per document. Use `import_documents()` to write the whole batch in a single version per table: ```python from haiku.rag.client import DocumentImport @@ -487,8 +487,8 @@ See [Automatic Title Generation](configuration/processing.md#automatic-title-gen ### 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. +Document create, update, and delete 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 the `documents`, `document_meta`, `chunks`, and `document_items` tables to their pre‑operation state using LanceDB’s table versioning. These writes are serialized under a single lock, so the rollback is safe under concurrent ingester workers. -- 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. +- Applies to: `create_document(...)`, `create_document_from_source(...)`, `update_document(...)`, `delete_document(...)` (including the `parent_uri` cascade), and internal rebuild/update flows. +- Scope: Document rows, their mutable attributes, and all associated chunks and items 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. From 885e7b7ce7e0ee4d7a09d5f8638b543ecf81655a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 11 Jun 2026 13:03:24 +0300 Subject: [PATCH 9/9] coverage --- .../haiku/rag/store/upgrades/v0_58_0.py | 4 +- tests/store/test_document_meta_split.py | 53 ++++++++++ tests/store/test_v0_58_0_migration.py | 98 +++++++++++++++++++ 3 files changed, 154 insertions(+), 1 deletion(-) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py index 38924db2..99fc78e8 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py @@ -74,7 +74,9 @@ async def _apply_split_document_meta(store: Store) -> None: # lancedb's .stats() stub claims TableStatistics but returns a plain dict. stats: dict = await store.documents_table.stats() # type: ignore[assignment] # ty: ignore[invalid-assignment] live_bytes = int(stats.get("total_bytes", 0)) - except Exception: + except ( + Exception + ): # pragma: no cover - defensive; stats() failure shouldn't block the split live_bytes = 0 free_bytes = shutil.disk_usage(store.db_path).free if live_bytes and free_bytes < live_bytes: diff --git a/tests/store/test_document_meta_split.py b/tests/store/test_document_meta_split.py index 69038899..52321d16 100644 --- a/tests/store/test_document_meta_split.py +++ b/tests/store/test_document_meta_split.py @@ -106,3 +106,56 @@ async def test_migrate_creates_and_populates_document_meta(temp_db_path): repo = DocumentRepository(store) doc = await repo.get_by_id("d") assert doc is not None and doc.uri == "u" + + +@pytest.mark.asyncio +async def test_repository_empty_and_missing_id_paths(temp_db_path): + """Cover the repository's early-return branches for empty input and + missing ids, plus get_content.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + + assert await repo.create([]) == [] + assert await repo.get_content("nope") is None + assert await repo.get_docling_data("nope") is None + assert await repo.get_pages_data("nope") is None + assert await repo.delete("nope") is False + + doc = await repo.create(Document(content="hello body", uri="u1")) + assert doc.id is not None + assert await repo.get_content(doc.id) == "hello body" + + +@pytest.mark.asyncio +async def test_batch_create_rolls_back_meta_on_failure(temp_db_path, monkeypatch): + """The list create path also rolls back its document_meta rows if the + documents write fails.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + + async def boom(*_a, **_k): + raise RuntimeError("documents add failed") + + monkeypatch.setattr(store.documents_table, "add", boom) + with pytest.raises(RuntimeError, match="documents add failed"): + await repo.create( + [Document(content="x", uri="u1"), Document(content="y", uri="u2")] + ) + monkeypatch.undo() + + assert await store.documents_table.count_rows() == 0 + assert await store.document_meta_table.count_rows() == 0 + + +@pytest.mark.asyncio +async def test_get_by_uri_with_orphan_meta_returns_none(temp_db_path): + """Defensive: a document_meta row whose documents row is missing (an + invariant violation) resolves to None, not a half-hydrated document.""" + from haiku.rag.store.engine import DocumentMetaRecord + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + await store.document_meta_table.add( + [DocumentMetaRecord(document_id="ghost", uri="u-ghost", metadata="{}")] + ) + assert await repo.get_by_uri("u-ghost") is None diff --git a/tests/store/test_v0_58_0_migration.py b/tests/store/test_v0_58_0_migration.py index 0d85d294..dab7ac5b 100644 --- a/tests/store/test_v0_58_0_migration.py +++ b/tests/store/test_v0_58_0_migration.py @@ -96,3 +96,101 @@ class TestV0_58_0Migration: meta_rows = await store.document_meta_table.query().to_list() assert len(meta_rows) == 1 assert meta_rows[0]["document_id"] == "doc-1" + + +@pytest.mark.asyncio +class TestV0_58_0MigrationEdgeCases: + async def test_resume_skips_already_migrated_rows(self, temp_db_path): + """A half-finished prior run leaves some document_meta rows; re-running + migrates only the rest and never duplicates.""" + from haiku.rag.store.engine import DocumentMetaRecord + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [ + LegacyDocumentRecord(id="a", content="x", uri="u-a", metadata="{}"), + LegacyDocumentRecord(id="b", content="y", uri="u-b", metadata="{}"), + ], + ) + # Pretend a prior run already moved doc "a". + await store.document_meta_table.add( + [DocumentMetaRecord(document_id="a", uri="u-a", metadata="{}")] + ) + await store.set_haiku_version("0.57.0") + + async with Store(temp_db_path, skip_migration_check=True) as store: + await store.migrate() + rows = await store.document_meta_table.query().to_list() + by_id = {r["document_id"]: r for r in rows} + assert set(by_id) == {"a", "b"} # exactly one row each, no duplicates + assert len(rows) == 2 + + async def test_skips_vacuum_when_disk_is_tight(self, temp_db_path, monkeypatch): + """When free disk can't cover one compacted copy, the split still + completes but the reclaim vacuum is skipped.""" + from types import SimpleNamespace + + from haiku.rag.store.upgrades import v0_58_0 + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_legacy_documents( + store, + [LegacyDocumentRecord(id="a", content="x", uri="u", metadata="{}")], + ) + await store.set_haiku_version("0.57.0") + + # Pretend almost no free disk so the reclaim vacuum is skipped. + monkeypatch.setattr( + v0_58_0.shutil, + "disk_usage", + lambda _p: SimpleNamespace(total=1, used=1, free=1), + ) + + vacuum_calls: list[int] = [] + + async with Store(temp_db_path, skip_migration_check=True) as store: + # Force a known nonzero live size so the free