629e1ba4ea
engine.py held four unrelated things: what the tables are, how to open a connection, how to read a database's state, and the Store that coordinates writes. At 1240 lines the Store's own concerns — locks, migrations, vacuum, tags — were hard to find among them. Table records, Arrow schemas, index_specs, ensure_indexes, REQUIRED_TABLES and query_to_pydantic move to store/schema.py, which imports nothing from haiku.rag: it describes the tables and never opens or mutates one. gather_database_info, get_database_stats, DatabaseInfo and its result models move to store/info.py. Nothing in Store calls them — they are read paths for the CLI, doctor, inspector and ingester API — so info depends on engine and not the reverse. engine.py keeps the Store, ConnectionMode, connect_lancedb, the tag helpers and the restore-order and retention constants. No re-exports: importers point at the new modules. test_app_info_uses_connect_lancedb_for_remote patched haiku.rag.store.engine.connect_lancedb; gather_database_info now binds that name in info.py, so the patch targets where the call is looked up.
65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
"""Tests for the v0.40.0 document_items population migration.
|
|
|
|
The migration walks every document with a docling blob, extracts items, and
|
|
populates the ``document_items`` table. It must stay independent of columns
|
|
introduced by later migrations (``picture_data`` in v0.45.0,
|
|
``heading_level`` / ``tree_depth`` in v0.48.0).
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.store.compression import compress_docling_split
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.schema import DocumentRecord
|
|
from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items
|
|
|
|
|
|
def _simple_docling_doc():
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
doc = DoclingDocument(name="simple")
|
|
doc.add_heading(text="Intro", level=1)
|
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello there.")
|
|
return doc
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_populate_handles_extra_columns_on_items_table(temp_db_path):
|
|
"""Regression: v0.40.0 must not fail when the live ``document_items``
|
|
table carries columns added by later migrations.
|
|
|
|
Mirrors what happens in practice: ``_init_tables`` always creates
|
|
``document_items`` with the latest schema (picture_data, heading_level,
|
|
tree_depth). v0.40.0 then runs against that table — its input must be
|
|
accepted even though it only writes the original 6 columns.
|
|
"""
|
|
docling_doc = _simple_docling_doc()
|
|
structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
|
|
|
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
|
# _init_tables already created document_items with the latest schema.
|
|
names = {f.name for f in await store.document_items_table.schema()}
|
|
assert {"picture_data", "heading_level", "tree_depth"} <= names
|
|
|
|
await store.documents_table.add(
|
|
[
|
|
DocumentRecord(
|
|
id="doc-1",
|
|
content="x",
|
|
docling_document=structure,
|
|
docling_pages=pages,
|
|
docling_version=docling_doc.version,
|
|
)
|
|
]
|
|
)
|
|
|
|
await _apply_populate_document_items(store)
|
|
|
|
rows = await (
|
|
store.document_items_table.query().where("document_id = 'doc-1'").to_list()
|
|
)
|
|
assert len(rows) >= 2
|
|
# Columns we didn't write should be null / default-typed.
|
|
for row in rows:
|
|
assert row.get("picture_data") is None
|