From 3366a6d383e191e27dbe4847dd3d9fdd13a76ee5 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 10 Jun 2026 20:31:43 +0300 Subject: [PATCH] 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)