Split meta document attributes into a document_meta table
This commit is contained in:
parent
4bc52f0710
commit
3366a6d383
18 changed files with 756 additions and 182 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
98
haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py
Normal file
98
haiku_rag_slim/haiku/rag/store/upgrades/v0_57_0.py
Normal 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
0
tests/store/__init__.py
Normal file
56
tests/store/legacy_documents.py
Normal file
56
tests/store/legacy_documents.py
Normal 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)
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
108
tests/store/test_document_meta_split.py
Normal file
108
tests/store/test_document_meta_split.py
Normal 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"
|
||||
|
|
@ -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'")
|
||||
|
|
|
|||
|
|
@ -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 == {
|
||||
|
|
|
|||
98
tests/store/test_v0_57_0_migration.py
Normal file
98
tests/store/test_v0_57_0_migration.py
Normal 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"
|
||||
|
|
@ -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)."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue