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
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
### 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.
|
- 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
|
updated = True
|
||||||
|
|
||||||
if updated:
|
if updated:
|
||||||
return await client.document_repository.update(doc)
|
return await client.document_repository.update_meta(doc)
|
||||||
return doc
|
return doc
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -730,7 +730,7 @@ async def update_document(
|
||||||
existing_doc.metadata = metadata
|
existing_doc.metadata = metadata
|
||||||
|
|
||||||
if content is None and chunks is None and docling_document is None:
|
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 chunks is not None:
|
||||||
if docling_document is not None:
|
if docling_document is not None:
|
||||||
|
|
|
||||||
|
|
@ -192,7 +192,7 @@ async def _rebuild_title_only(
|
||||||
continue
|
continue
|
||||||
if title is not None:
|
if title is not None:
|
||||||
doc.title = title
|
doc.title = title
|
||||||
await client.document_repository.update(doc)
|
await client.document_repository.update_meta(doc)
|
||||||
assert doc.id is not None
|
assert doc.id is not None
|
||||||
yield doc.id
|
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
|
document. Used by RECHUNK and FULL modes after the chunks table has been
|
||||||
cleared.
|
cleared.
|
||||||
"""
|
"""
|
||||||
from haiku.rag.store.engine import DocumentRecord
|
from haiku.rag.store.engine import DocumentMetaRecord, DocumentRecord
|
||||||
|
|
||||||
if not documents:
|
if not documents:
|
||||||
return
|
return
|
||||||
|
|
||||||
now = datetime.now().isoformat()
|
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 = []
|
doc_records = []
|
||||||
|
meta_records = []
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
assert doc.id is not None
|
assert doc.id is not None
|
||||||
doc_records.append(
|
doc_records.append(
|
||||||
DocumentRecord(
|
DocumentRecord(
|
||||||
id=doc.id,
|
id=doc.id,
|
||||||
content=doc.content,
|
content=doc.content,
|
||||||
uri=doc.uri,
|
|
||||||
title=doc.title,
|
|
||||||
metadata=json.dumps(doc.metadata),
|
|
||||||
docling_document=doc.docling_document,
|
docling_document=doc.docling_document,
|
||||||
docling_pages=doc.docling_pages,
|
docling_pages=doc.docling_pages,
|
||||||
docling_version=doc.docling_version,
|
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,
|
created_at=doc.created_at.isoformat() if doc.created_at else now,
|
||||||
updated_at=now,
|
updated_at=now,
|
||||||
)
|
)
|
||||||
|
|
@ -521,6 +529,12 @@ async def _flush_rebuild_batch(
|
||||||
.when_matched_update_all()
|
.when_matched_update_all()
|
||||||
.execute(doc_records)
|
.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)
|
# Batch create all chunks (single LanceDB version)
|
||||||
if chunks:
|
if chunks:
|
||||||
|
|
|
||||||
|
|
@ -77,12 +77,20 @@ async def connect_lancedb(
|
||||||
class DocumentRecord(LanceModel):
|
class DocumentRecord(LanceModel):
|
||||||
id: str = Field(default_factory=lambda: str(uuid4()))
|
id: str = Field(default_factory=lambda: str(uuid4()))
|
||||||
content: str
|
content: str
|
||||||
uri: str | None = None
|
|
||||||
title: str | None = None
|
|
||||||
metadata: str = Field(default="{}")
|
|
||||||
docling_document: bytes | None = None
|
docling_document: bytes | None = None
|
||||||
docling_pages: bytes | None = None
|
docling_pages: bytes | None = None
|
||||||
docling_version: str | 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: "")
|
created_at: str = Field(default_factory=lambda: "")
|
||||||
updated_at: str = Field(default_factory=lambda: "")
|
updated_at: str = Field(default_factory=lambda: "")
|
||||||
|
|
||||||
|
|
@ -171,7 +179,13 @@ class SettingsRecord(LanceModel):
|
||||||
settings: str = Field(default="{}")
|
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:
|
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
|
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
|
||||||
self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
|
self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
|
||||||
|
|
||||||
# Initialize tables (creates them if they don't exist)
|
# Initialize tables (creates them if they don't exist). For an existing
|
||||||
await self._init_tables()
|
# 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
|
# Checkout tables to historical state if before is specified
|
||||||
if self._before is not None:
|
if self._before is not None:
|
||||||
await self._checkout_tables_before(self._before)
|
await self._checkout_tables_before(self._before)
|
||||||
|
|
||||||
# Set version for new databases, check migrations for existing ones
|
# Set version for new databases.
|
||||||
if is_new_db:
|
if is_new_db and not self._read_only:
|
||||||
if not self._read_only:
|
await self._set_initial_version()
|
||||||
await self._set_initial_version()
|
|
||||||
elif not self._skip_migration_check:
|
|
||||||
await self._check_migrations()
|
|
||||||
|
|
||||||
# Validate config compatibility after connection is established
|
# Validate config compatibility after connection is established
|
||||||
if not self._skip_validation:
|
if not self._skip_validation:
|
||||||
|
|
@ -493,6 +506,7 @@ class Store:
|
||||||
retention = timedelta(seconds=retention_seconds)
|
retention = timedelta(seconds=retention_seconds)
|
||||||
for table in [
|
for table in [
|
||||||
self.documents_table,
|
self.documents_table,
|
||||||
|
self.document_meta_table,
|
||||||
self.chunks_table,
|
self.chunks_table,
|
||||||
self.document_items_table,
|
self.document_items_table,
|
||||||
self.settings_table,
|
self.settings_table,
|
||||||
|
|
@ -553,9 +567,23 @@ class Store:
|
||||||
settings_repo = SettingsRepository(self)
|
settings_repo = SettingsRepository(self)
|
||||||
await settings_repo.validate_config_compatibility()
|
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)."""
|
"""Initialize database tables (create if they don't exist)."""
|
||||||
existing_tables = (await self.db.list_tables()).tables
|
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)
|
missing_tables = set(REQUIRED_TABLES) - set(existing_tables)
|
||||||
|
|
||||||
if missing_tables and self._read_only:
|
if missing_tables and self._read_only:
|
||||||
|
|
@ -572,6 +600,22 @@ class Store:
|
||||||
"documents", schema=get_documents_arrow_schema()
|
"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
|
# Create or open chunks table
|
||||||
if "chunks" in existing_tables:
|
if "chunks" in existing_tables:
|
||||||
self.chunks_table = await self.db.open_table("chunks")
|
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."""
|
"""Capture current versions of key tables for rollback using LanceDB's API."""
|
||||||
return {
|
return {
|
||||||
"documents": await self.documents_table.version(),
|
"documents": await self.documents_table.version(),
|
||||||
|
"document_meta": await self.document_meta_table.version(),
|
||||||
"chunks": await self.chunks_table.version(),
|
"chunks": await self.chunks_table.version(),
|
||||||
"document_items": await self.document_items_table.version(),
|
"document_items": await self.document_items_table.version(),
|
||||||
"settings": await self.settings_table.version(),
|
"settings": await self.settings_table.version(),
|
||||||
|
|
@ -761,6 +806,7 @@ class Store:
|
||||||
"""
|
"""
|
||||||
self._assert_writable()
|
self._assert_writable()
|
||||||
await self.documents_table.restore(int(versions["documents"]))
|
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.chunks_table.restore(int(versions["chunks"]))
|
||||||
await self.document_items_table.restore(int(versions["document_items"]))
|
await self.document_items_table.restore(int(versions["document_items"]))
|
||||||
await self.settings_table.restore(int(versions["settings"]))
|
await self.settings_table.restore(int(versions["settings"]))
|
||||||
|
|
@ -786,6 +832,7 @@ class Store:
|
||||||
|
|
||||||
tables = [
|
tables = [
|
||||||
("documents", self.documents_table),
|
("documents", self.documents_table),
|
||||||
|
("document_meta", self.document_meta_table),
|
||||||
("chunks", self.chunks_table),
|
("chunks", self.chunks_table),
|
||||||
("document_items", self.document_items_table),
|
("document_items", self.document_items_table),
|
||||||
("settings", self.settings_table),
|
("settings", self.settings_table),
|
||||||
|
|
@ -838,6 +885,7 @@ class Store:
|
||||||
"""
|
"""
|
||||||
table_map = {
|
table_map = {
|
||||||
"documents": self.documents_table,
|
"documents": self.documents_table,
|
||||||
|
"document_meta": self.document_meta_table,
|
||||||
"chunks": self.chunks_table,
|
"chunks": self.chunks_table,
|
||||||
"document_items": self.document_items_table,
|
"document_items": self.document_items_table,
|
||||||
"settings": self.settings_table,
|
"settings": self.settings_table,
|
||||||
|
|
|
||||||
|
|
@ -272,14 +272,14 @@ class ChunkRepository:
|
||||||
# filter in pandas, head(limit)) silently under-returned
|
# filter in pandas, head(limit)) silently under-returned
|
||||||
# whenever the top-N window lacked `limit` matching chunks.
|
# whenever the top-N window lacked `limit` matching chunks.
|
||||||
docs_df = await (
|
docs_df = await (
|
||||||
self.store.documents_table.query()
|
self.store.document_meta_table.query()
|
||||||
.select(["id"])
|
.select(["document_id"])
|
||||||
.where(filter)
|
.where(filter)
|
||||||
.to_pandas()
|
.to_pandas()
|
||||||
)
|
)
|
||||||
if docs_df.empty:
|
if docs_df.empty:
|
||||||
return []
|
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})"
|
chunk_filter = f"document_id IN ({id_list})"
|
||||||
|
|
||||||
if query_vector is not None:
|
if query_vector is not None:
|
||||||
|
|
@ -344,11 +344,11 @@ class ChunkRepository:
|
||||||
|
|
||||||
results = await query_to_pydantic(query, self.store.ChunkRecord)
|
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 (
|
doc_rows = await (
|
||||||
self.store.documents_table.query()
|
self.store.document_meta_table.query()
|
||||||
.select(["id", "uri", "title", "metadata"])
|
.select(["document_id", "uri", "title", "metadata"])
|
||||||
.where(f"id = '{document_id}'")
|
.where(f"document_id = '{document_id}'")
|
||||||
.limit(1)
|
.limit(1)
|
||||||
.to_list()
|
.to_list()
|
||||||
)
|
)
|
||||||
|
|
@ -504,14 +504,14 @@ class ChunkRepository:
|
||||||
documents_map: dict[str, dict] = {}
|
documents_map: dict[str, dict] = {}
|
||||||
if document_ids:
|
if document_ids:
|
||||||
id_list = "', '".join(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 (
|
doc_rows = await (
|
||||||
self.store.documents_table.query()
|
self.store.document_meta_table.query()
|
||||||
.select(["id", "uri", "title", "metadata"])
|
.select(["document_id", "uri", "title", "metadata"])
|
||||||
.where(where_clause)
|
.where(where_clause)
|
||||||
.to_list()
|
.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
|
# Build final results with document info
|
||||||
chunks_with_scores = []
|
chunks_with_scores = []
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ from uuid import uuid4
|
||||||
from lancedb.index import BTree
|
from lancedb.index import BTree
|
||||||
|
|
||||||
from haiku.rag.store.engine import (
|
from haiku.rag.store.engine import (
|
||||||
|
DocumentMetaRecord,
|
||||||
DocumentRecord,
|
DocumentRecord,
|
||||||
Store,
|
Store,
|
||||||
get_documents_arrow_schema,
|
get_documents_arrow_schema,
|
||||||
|
|
@ -16,7 +17,20 @@ from haiku.rag.utils import escape_sql_string
|
||||||
|
|
||||||
|
|
||||||
class DocumentRepository:
|
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:
|
def __init__(self, store: Store) -> None:
|
||||||
self.store = store
|
self.store = store
|
||||||
|
|
@ -43,45 +57,61 @@ class DocumentRepository:
|
||||||
self._document_item_repository = DocumentItemRepository(self.store)
|
self._document_item_repository = DocumentItemRepository(self.store)
|
||||||
return self._document_item_repository
|
return self._document_item_repository
|
||||||
|
|
||||||
def _record_to_document(self, record: DocumentRecord) -> Document:
|
def _merge_to_document(
|
||||||
"""Convert a DocumentRecord to a Document model."""
|
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(
|
return Document(
|
||||||
id=record.id,
|
id=doc.id,
|
||||||
content=record.content,
|
content=doc.content,
|
||||||
uri=record.uri,
|
uri=meta.uri if meta else None,
|
||||||
title=record.title,
|
title=meta.title if meta else None,
|
||||||
metadata=json.loads(record.metadata),
|
metadata=json.loads(meta.metadata) if meta else {},
|
||||||
docling_document=record.docling_document,
|
docling_document=doc.docling_document,
|
||||||
docling_pages=record.docling_pages,
|
docling_pages=doc.docling_pages,
|
||||||
docling_version=record.docling_version,
|
docling_version=doc.docling_version,
|
||||||
created_at=datetime.fromisoformat(record.created_at)
|
created_at=datetime.fromisoformat(created) if created else datetime.now(),
|
||||||
if record.created_at
|
updated_at=datetime.fromisoformat(updated) if updated else datetime.now(),
|
||||||
else datetime.now(),
|
|
||||||
updated_at=datetime.fromisoformat(record.updated_at)
|
|
||||||
if record.updated_at
|
|
||||||
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,
|
self,
|
||||||
entity: Document,
|
entity: Document,
|
||||||
doc_id: str,
|
doc_id: str,
|
||||||
created_at: str,
|
created_at: str,
|
||||||
updated_at: str,
|
updated_at: str,
|
||||||
) -> DocumentRecord:
|
) -> DocumentMetaRecord:
|
||||||
return DocumentRecord(
|
return DocumentMetaRecord(
|
||||||
id=doc_id,
|
document_id=doc_id,
|
||||||
content=entity.content,
|
|
||||||
uri=entity.uri,
|
uri=entity.uri,
|
||||||
title=entity.title,
|
title=entity.title,
|
||||||
metadata=json.dumps(entity.metadata),
|
metadata=json.dumps(entity.metadata),
|
||||||
docling_document=entity.docling_document,
|
|
||||||
docling_pages=entity.docling_pages,
|
|
||||||
docling_version=entity.docling_version,
|
|
||||||
created_at=created_at,
|
created_at=created_at,
|
||||||
updated_at=updated_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
|
@overload
|
||||||
async def create(self, entity: Document) -> Document: ...
|
async def create(self, entity: Document) -> Document: ...
|
||||||
|
|
||||||
|
|
@ -97,12 +127,29 @@ class DocumentRepository:
|
||||||
"""
|
"""
|
||||||
self.store._assert_writable()
|
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):
|
if isinstance(entity, Document):
|
||||||
doc_id = str(uuid4())
|
doc_id = str(uuid4())
|
||||||
now = datetime.now().isoformat()
|
now = datetime.now().isoformat()
|
||||||
await self.store.documents_table.add(
|
await self.store.document_meta_table.add(
|
||||||
[self._to_record(entity, doc_id, now, now)]
|
[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.id = doc_id
|
||||||
entity.created_at = datetime.fromisoformat(now)
|
entity.created_at = datetime.fromisoformat(now)
|
||||||
entity.updated_at = datetime.fromisoformat(now)
|
entity.updated_at = datetime.fromisoformat(now)
|
||||||
|
|
@ -114,15 +161,25 @@ class DocumentRepository:
|
||||||
|
|
||||||
now = datetime.now().isoformat()
|
now = datetime.now().isoformat()
|
||||||
created_at = datetime.fromisoformat(now)
|
created_at = datetime.fromisoformat(now)
|
||||||
records = []
|
doc_records = []
|
||||||
|
meta_records = []
|
||||||
|
doc_ids = []
|
||||||
for document in documents:
|
for document in documents:
|
||||||
doc_id = str(uuid4())
|
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.id = doc_id
|
||||||
document.created_at = created_at
|
document.created_at = created_at
|
||||||
document.updated_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
|
return documents
|
||||||
|
|
||||||
async def get_by_id(self, entity_id: str) -> Document | None:
|
async def get_by_id(self, entity_id: str) -> Document | None:
|
||||||
|
|
@ -136,7 +193,8 @@ class DocumentRepository:
|
||||||
if not results:
|
if not results:
|
||||||
return None
|
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:
|
async def get_content(self, entity_id: str) -> str | None:
|
||||||
"""Get only the text content of a document (skips docling blobs)."""
|
"""Get only the text content of a document (skips docling blobs)."""
|
||||||
|
|
@ -197,28 +255,43 @@ class DocumentRepository:
|
||||||
docling_pages=row.get("docling_pages"),
|
docling_pages=row.get("docling_pages"),
|
||||||
)
|
)
|
||||||
|
|
||||||
async def update(self, entity: Document) -> Document:
|
async def update_meta(self, entity: Document) -> Document:
|
||||||
"""Update an existing 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()
|
self.store._assert_writable()
|
||||||
|
|
||||||
assert entity.id, "Document ID is required for update"
|
assert entity.id, "Document ID is required for update"
|
||||||
|
|
||||||
# Update timestamp
|
|
||||||
now = datetime.now().isoformat()
|
now = datetime.now().isoformat()
|
||||||
entity.updated_at = datetime.fromisoformat(now)
|
entity.updated_at = datetime.fromisoformat(now)
|
||||||
|
created = entity.created_at.isoformat() if entity.created_at else now
|
||||||
record = self._to_record(
|
record = self._to_meta_record(entity, entity.id, created, now)
|
||||||
entity,
|
# Update only — no insert. Every real document has a document_meta row
|
||||||
entity.id,
|
# from create()/migration; inserting on no-match would manufacture a
|
||||||
entity.created_at.isoformat() if entity.created_at else now,
|
# ghost row (visible to list_all/count) for an id with no documents row.
|
||||||
now,
|
|
||||||
)
|
|
||||||
await (
|
await (
|
||||||
self.store.documents_table.merge_insert("id")
|
self.store.document_meta_table.merge_insert("document_id")
|
||||||
.when_matched_update_all()
|
.when_matched_update_all()
|
||||||
.execute([record])
|
.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
|
return entity
|
||||||
|
|
||||||
async def delete(self, entity_id: str) -> bool:
|
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.chunk_repository.delete_by_document_id(entity_id)
|
||||||
await self.document_item_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)
|
safe_id = escape_sql_string(entity_id)
|
||||||
await self.store.documents_table.delete(f"id = '{safe_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
|
return True
|
||||||
|
|
||||||
_LISTING_COLUMNS = ["id", "title", "uri", "metadata", "created_at", "updated_at"]
|
|
||||||
|
|
||||||
async def list_all(
|
async def list_all(
|
||||||
self,
|
self,
|
||||||
limit: int | None = None,
|
limit: int | None = None,
|
||||||
|
|
@ -250,20 +322,21 @@ class DocumentRepository:
|
||||||
) -> list[Document]:
|
) -> list[Document]:
|
||||||
"""List all documents with optional pagination and filtering.
|
"""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:
|
Args:
|
||||||
limit: Maximum number of documents to return.
|
limit: Maximum number of documents to return.
|
||||||
offset: Number of documents to skip.
|
offset: Number of documents to skip.
|
||||||
filter: Optional SQL WHERE clause to filter documents.
|
filter: Optional SQL WHERE clause over document_meta columns.
|
||||||
include_content: Whether to load content and docling_document.
|
include_content: Whether to also load content and docling blobs.
|
||||||
Defaults to False to avoid loading large blobs for listing.
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of Document instances matching the criteria.
|
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:
|
if filter is not None:
|
||||||
query = query.where(filter)
|
query = query.where(filter)
|
||||||
if offset is not None:
|
if offset is not None:
|
||||||
|
|
@ -271,50 +344,56 @@ class DocumentRepository:
|
||||||
if limit is not None:
|
if limit is not None:
|
||||||
query = query.limit(limit)
|
query = query.limit(limit)
|
||||||
|
|
||||||
if include_content:
|
meta_records = await query_to_pydantic(query, DocumentMetaRecord)
|
||||||
results = await query_to_pydantic(query, DocumentRecord)
|
|
||||||
return [self._record_to_document(doc) for doc in results]
|
|
||||||
|
|
||||||
return [
|
if not include_content:
|
||||||
Document(
|
return [
|
||||||
id=row["id"],
|
self._merge_to_document(DocumentRecord(id=m.document_id, content=""), m)
|
||||||
content="",
|
for m in meta_records
|
||||||
title=row.get("title"),
|
]
|
||||||
uri=row.get("uri"),
|
|
||||||
metadata=json.loads(row.get("metadata", "{}")),
|
documents: list[Document] = []
|
||||||
created_at=datetime.fromisoformat(row["created_at"])
|
for meta in meta_records:
|
||||||
if row.get("created_at")
|
safe_id = escape_sql_string(meta.document_id)
|
||||||
else datetime.now(),
|
doc_results = await query_to_pydantic(
|
||||||
updated_at=datetime.fromisoformat(row["updated_at"])
|
self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
|
||||||
if row.get("updated_at")
|
DocumentRecord,
|
||||||
else datetime.now(),
|
|
||||||
)
|
)
|
||||||
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:
|
async def count(self, filter: str | None = None) -> int:
|
||||||
"""Count documents with optional filtering.
|
"""Count documents with optional filtering (over document_meta columns)."""
|
||||||
|
return await self.store.document_meta_table.count_rows(filter=filter)
|
||||||
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)
|
|
||||||
|
|
||||||
async def get_by_uri(self, uri: str) -> Document | None:
|
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)
|
escaped_uri = escape_sql_string(uri)
|
||||||
results = await query_to_pydantic(
|
meta_results = await query_to_pydantic(
|
||||||
self.store.documents_table.query().where(f"uri = '{escaped_uri}'").limit(1),
|
self.store.document_meta_table.query()
|
||||||
DocumentRecord,
|
.where(f"uri = '{escaped_uri}'")
|
||||||
|
.limit(1),
|
||||||
|
DocumentMetaRecord,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not results:
|
if not meta_results:
|
||||||
return None
|
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:
|
async def delete_all(self) -> None:
|
||||||
"""Delete all documents from the database."""
|
"""Delete all documents from the database."""
|
||||||
|
|
@ -344,8 +423,18 @@ class DocumentRepository:
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
if count > 0:
|
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")
|
await self.store.db.drop_table("documents")
|
||||||
self.store.documents_table = await self.store.db.create_table(
|
self.store.documents_table = await self.store.db.create_table(
|
||||||
"documents", schema=get_documents_arrow_schema()
|
"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 (
|
from haiku.rag.store.upgrades.v0_50_0 import (
|
||||||
upgrade_canonical_metadata_keys as upgrade_0_50_0_canonical_metadata_keys,
|
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_20_0_docling)
|
||||||
upgrades.append(upgrade_0_23_1_contextualize)
|
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_45_0_extract_picture_bytes)
|
||||||
upgrades.append(upgrade_0_48_0_heading_hierarchy)
|
upgrades.append(upgrade_0_48_0_heading_hierarchy)
|
||||||
upgrades.append(upgrade_0_50_0_canonical_metadata_keys)
|
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."""
|
"""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.compression import compress_docling_split
|
||||||
from haiku.rag.store.engine import DocumentRecord
|
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()
|
docling_doc = _make_docling_doc()
|
||||||
json_str = docling_doc.model_dump_json()
|
json_str = docling_doc.model_dump_json()
|
||||||
|
|
@ -405,7 +406,6 @@ class TestDocumentItemMigration:
|
||||||
doc_record = DocumentRecord(
|
doc_record = DocumentRecord(
|
||||||
id="test-doc-1",
|
id="test-doc-1",
|
||||||
content="test content",
|
content="test content",
|
||||||
uri="test://doc",
|
|
||||||
docling_document=structure,
|
docling_document=structure,
|
||||||
docling_pages=pages,
|
docling_pages=pages,
|
||||||
docling_version=docling_doc.version,
|
docling_version=docling_doc.version,
|
||||||
|
|
@ -415,12 +415,11 @@ class TestDocumentItemMigration:
|
||||||
# Verify no items exist yet
|
# Verify no items exist yet
|
||||||
assert await store.document_items_table.count_rows() == 0
|
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:
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||||
applied = await store.migrate()
|
await _apply_populate_document_items(store)
|
||||||
|
|
||||||
# Should have applied the v0.40.0 migration
|
|
||||||
assert any("document_items" in desc for desc in applied)
|
|
||||||
|
|
||||||
# Items should now exist
|
# Items should now exist
|
||||||
item_count = await store.document_items_table.count_rows(
|
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):
|
async def test_migration_skips_documents_without_docling(self, temp_db_path):
|
||||||
"""Test that migration handles documents without docling data."""
|
"""Test that migration handles documents without docling data."""
|
||||||
from haiku.rag.store.engine import DocumentRecord
|
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:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.39.0")
|
await store.set_haiku_version("0.39.0")
|
||||||
|
|
@ -452,7 +452,7 @@ class TestDocumentItemMigration:
|
||||||
await store.documents_table.add([doc_record])
|
await store.documents_table.add([doc_record])
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
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
|
# No items should have been created
|
||||||
assert await store.document_items_table.count_rows() == 0
|
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.compression import compress_json, decompress_json
|
||||||
from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord
|
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"
|
fake_png = b"\x89PNG\r\n\x1a\nlegacy-picture-bytes-for-test"
|
||||||
data_uri = "data:image/png;base64," + base64.b64encode(fake_png).decode("ascii")
|
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}
|
assert "picture_data" not in {f.name for f in schema_before}
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||||
applied = await store.migrate()
|
# Apply the v0.45.0 migration in isolation (the full chain would also
|
||||||
assert any("picture" in d.lower() for d in applied)
|
# run later migrations that touch documents.metadata, absent here).
|
||||||
|
await _apply_extract_picture_bytes(store)
|
||||||
|
|
||||||
# Column was added by the migration
|
# Column was added by the migration
|
||||||
schema_after = await store.document_items_table.schema()
|
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.compression import compress_docling_split
|
||||||
from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store
|
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():
|
def _docling_with_levels():
|
||||||
|
|
@ -58,8 +59,10 @@ class TestV0_48_0Migration:
|
||||||
)
|
)
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||||
applied = await store.migrate()
|
# Apply the migration in isolation: store.migrate() would run the
|
||||||
assert any("0.48.0" in d for d in applied)
|
# 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 (
|
rows = await (
|
||||||
store.document_items_table.query()
|
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:
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||||
await store.migrate()
|
# Apply twice to prove idempotency (in isolation from the chain).
|
||||||
|
await _apply_backfill_heading_hierarchy(store)
|
||||||
from haiku.rag.store.upgrades.v0_48_0 import (
|
|
||||||
_apply_backfill_heading_hierarchy,
|
|
||||||
)
|
|
||||||
|
|
||||||
await _apply_backfill_heading_hierarchy(store)
|
await _apply_backfill_heading_hierarchy(store)
|
||||||
|
|
||||||
rows = await (
|
rows = await (
|
||||||
|
|
@ -144,7 +143,7 @@ class TestV0_48_0Migration:
|
||||||
)
|
)
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||||
await store.migrate()
|
await _apply_backfill_heading_hierarchy(store)
|
||||||
rows = await (
|
rows = await (
|
||||||
store.document_items_table.query()
|
store.document_items_table.query()
|
||||||
.where("document_id = 'plain'")
|
.where("document_id = 'plain'")
|
||||||
|
|
|
||||||
|
|
@ -2,19 +2,30 @@ import json
|
||||||
|
|
||||||
import pytest
|
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
|
@pytest.mark.asyncio
|
||||||
class TestV0_50_0Migration:
|
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 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:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.48.1")
|
await seed_legacy_documents(
|
||||||
await store.documents_table.add(
|
store,
|
||||||
[
|
[
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="doc-s3",
|
id="doc-s3",
|
||||||
content="x",
|
content="x",
|
||||||
uri="s3://b/k",
|
uri="s3://b/k",
|
||||||
|
|
@ -26,7 +37,7 @@ class TestV0_50_0Migration:
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
),
|
),
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="doc-fs",
|
id="doc-fs",
|
||||||
content="y",
|
content="y",
|
||||||
uri="file:///tmp/x.md",
|
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:
|
await _apply_canonical_metadata_keys(store)
|
||||||
applied = await store.migrate()
|
|
||||||
assert any("0.50.0" in d for d in applied)
|
|
||||||
|
|
||||||
rows = await store.documents_table.query().to_list()
|
rows = await store.documents_table.query().to_list()
|
||||||
by_id = {r["id"]: json.loads(r["metadata"]) for r in rows}
|
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 def test_idempotent_on_already_migrated(self, temp_db_path):
|
||||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.48.1")
|
await seed_legacy_documents(
|
||||||
await store.documents_table.add(
|
store,
|
||||||
[
|
[
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="d",
|
id="d",
|
||||||
content="x",
|
content="x",
|
||||||
uri="s3://b/k",
|
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 _apply_canonical_metadata_keys(store)
|
||||||
await store.migrate()
|
|
||||||
rows = await store.documents_table.query().to_list()
|
rows = await store.documents_table.query().to_list()
|
||||||
assert json.loads(rows[0]["metadata"]) == {
|
assert json.loads(rows[0]["metadata"]) == {
|
||||||
"source_revision": "abc",
|
"source_revision": "abc",
|
||||||
|
|
@ -91,14 +99,11 @@ class TestV0_50_0Migration:
|
||||||
quoted-key form. Values that happen to contain the substring `etag` and
|
quoted-key form. Values that happen to contain the substring `etag` and
|
||||||
composite key names like `my_etag_key` must not get rewritten."""
|
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:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.48.1")
|
await seed_legacy_documents(
|
||||||
await store.documents_table.add(
|
store,
|
||||||
[
|
[
|
||||||
# `etag` only appears as a VALUE — no `etag` key. Either
|
# `etag` only appears as a VALUE — no `etag` key.
|
||||||
# the LIKE excludes it (no work), or it pulls it in and
|
LegacyDocumentRecord(
|
||||||
# _normalize_metadata leaves it alone. Either way the row
|
|
||||||
# must end up unchanged.
|
|
||||||
DocumentRecord(
|
|
||||||
id="value-only",
|
id="value-only",
|
||||||
content="x",
|
content="x",
|
||||||
uri="u1",
|
uri="u1",
|
||||||
|
|
@ -110,8 +115,7 @@ class TestV0_50_0Migration:
|
||||||
),
|
),
|
||||||
),
|
),
|
||||||
# A composite key containing `etag` but not equal to it.
|
# A composite key containing `etag` but not equal to it.
|
||||||
# Must not be rewritten.
|
LegacyDocumentRecord(
|
||||||
DocumentRecord(
|
|
||||||
id="composite-key",
|
id="composite-key",
|
||||||
content="x",
|
content="x",
|
||||||
uri="u2",
|
uri="u2",
|
||||||
|
|
@ -122,11 +126,10 @@ class TestV0_50_0Migration:
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
),
|
),
|
||||||
]
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
await _apply_canonical_metadata_keys(store)
|
||||||
await store.migrate()
|
|
||||||
rows = await store.documents_table.query().to_list()
|
rows = await store.documents_table.query().to_list()
|
||||||
by_id = {r["id"]: json.loads(r["metadata"]) for r in rows}
|
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
|
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."""
|
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:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.48.1")
|
await seed_legacy_documents(
|
||||||
await store.documents_table.add(
|
store,
|
||||||
[
|
[
|
||||||
# Contains the substring `"etag"` so the WHERE LIKE pulls
|
# Contains the substring `"etag"` so the WHERE LIKE pulls it
|
||||||
# it in, but it's not valid JSON — json.loads fails.
|
# in, but it's not valid JSON — json.loads fails.
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="bad",
|
id="bad",
|
||||||
content="x",
|
content="x",
|
||||||
uri="u1",
|
uri="u1",
|
||||||
metadata='{"etag": broken',
|
metadata='{"etag": broken',
|
||||||
),
|
),
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="good",
|
id="good",
|
||||||
content="x",
|
content="x",
|
||||||
uri="u2",
|
uri="u2",
|
||||||
metadata=json.dumps({"etag": "abc"}),
|
metadata=json.dumps({"etag": "abc"}),
|
||||||
),
|
),
|
||||||
]
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
|
||||||
# Must not raise.
|
# Must not raise.
|
||||||
applied = await store.migrate()
|
await _apply_canonical_metadata_keys(store)
|
||||||
assert any("0.50.0" in d for d in applied)
|
|
||||||
|
|
||||||
rows = await store.documents_table.query().to_list()
|
rows = await store.documents_table.query().to_list()
|
||||||
by_id = {r["id"]: r["metadata"] for r in rows}
|
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
|
"""If both legacy and canonical keys are present, the canonical wins
|
||||||
and the legacy is dropped — defends against partial-migration states."""
|
and the legacy is dropped — defends against partial-migration states."""
|
||||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||||
await store.set_haiku_version("0.48.1")
|
await seed_legacy_documents(
|
||||||
await store.documents_table.add(
|
store,
|
||||||
[
|
[
|
||||||
DocumentRecord(
|
LegacyDocumentRecord(
|
||||||
id="d",
|
id="d",
|
||||||
content="x",
|
content="x",
|
||||||
uri="s3://b/k",
|
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 _apply_canonical_metadata_keys(store)
|
||||||
await store.migrate()
|
|
||||||
rows = await store.documents_table.query().to_list()
|
rows = await store.documents_table.query().to_list()
|
||||||
meta = json.loads(rows[0]["metadata"])
|
meta = json.loads(rows[0]["metadata"])
|
||||||
assert meta == {
|
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"
|
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()
|
@pytest.mark.vcr()
|
||||||
async def test_client_ask(allow_model_requests, temp_db_path):
|
async def test_client_ask(allow_model_requests, temp_db_path):
|
||||||
"""Test asking questions returns answer and citations (VCR recorded)."""
|
"""Test asking questions returns answer and citations (VCR recorded)."""
|
||||||
|
|
|
||||||
|
|
@ -267,7 +267,7 @@ class TestInitFailureCleanup:
|
||||||
async def fake_connect(*args, **kwargs):
|
async def fake_connect(*args, **kwargs):
|
||||||
return mock_conn
|
return mock_conn
|
||||||
|
|
||||||
async def failing_init_tables(self):
|
async def failing_init_tables(self, is_new_db):
|
||||||
raise RuntimeError("simulated table init failure")
|
raise RuntimeError("simulated table init failure")
|
||||||
|
|
||||||
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
|
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ from typing import TypedDict
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
|
||||||
class ChunkData(TypedDict):
|
class ChunkData(TypedDict):
|
||||||
|
|
@ -184,7 +185,14 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash(
|
||||||
_StagingMarkerRecord,
|
_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"])
|
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
||||||
assert doc.id is not None
|
assert doc.id is not None
|
||||||
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue