56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""Helpers to seed a `documents` table in its pre-0.58 shape.
|
|
|
|
Before the document_meta split (v0.58.0), the `documents` table carried
|
|
`uri/title/metadata/created_at/updated_at` alongside the content+blobs. The
|
|
migration-chain tests need to reproduce that legacy layout so the older
|
|
migrations (which read `documents.metadata`, etc.) and v0.58.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.58 `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.58 schema and add records,
|
|
simulating a database created before the document_meta split."""
|
|
if "documents" in (await store.db.list_tables()).tables:
|
|
await store.db.drop_table("documents")
|
|
store.documents_table = await store.db.create_table(
|
|
"documents", schema=legacy_documents_schema()
|
|
)
|
|
if records:
|
|
await store.documents_table.add(records)
|