Split meta document attributes into a document_meta table

This commit is contained in:
Yiorgis Gozadinos 2026-06-10 20:31:43 +03:00
parent 4bc52f0710
commit 3366a6d383
No known key found for this signature in database
18 changed files with 756 additions and 182 deletions

View file

@ -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=<name>` URI; non-PDF attachments (e.g. `.joboptions`) are no longer misrouted to docling's PDF backend, and unsupported extensions are skipped.

View file

@ -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:

View file

@ -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:

View file

@ -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,

View file

@ -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 = []

View file

@ -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
)

View file

@ -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)

View file

@ -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",
)

0
tests/store/__init__.py Normal file
View file

View file

@ -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)

View file

@ -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()

View file

@ -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"

View file

@ -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'")

View file

@ -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 == {

View file

@ -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"

View file

@ -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)."""

View file

@ -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)

View file

@ -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)