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.
229 lines
7.8 KiB
Python
229 lines
7.8 KiB
Python
"""Tests for the v0.45.0 picture-data backfill migration.
|
|
|
|
The migration walks every document, decodes picture image URIs out of the
|
|
stored docling blob, writes the bytes to ``document_items.picture_data``,
|
|
and then strips the URIs from the blob (which lives separately on the
|
|
documents table).
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
from haiku.rag.store.compression import compress_json, decompress_json
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.schema import DocumentItemRecord, DocumentRecord
|
|
from haiku.rag.store.upgrades.v0_45_0 import _apply_extract_picture_bytes
|
|
|
|
PNG_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes"
|
|
PNG_DATA_URI = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode('ascii')}"
|
|
|
|
|
|
def _docling_blob_with_picture(self_ref: str = "#/pictures/0") -> bytes:
|
|
"""Build a compressed docling-document blob carrying one picture with
|
|
an inline data URI — the shape v0.45.0 was written to backfill from."""
|
|
doc = {
|
|
"schema_name": "DoclingDocument",
|
|
"version": "1.10.0",
|
|
"name": "test",
|
|
"pictures": [
|
|
{
|
|
"self_ref": self_ref,
|
|
"image": {"mimetype": "image/png", "uri": PNG_DATA_URI},
|
|
}
|
|
],
|
|
"tables": [],
|
|
"texts": [],
|
|
"groups": [],
|
|
"body": {"self_ref": "#/body", "children": [], "label": "unspecified"},
|
|
"furniture": {
|
|
"self_ref": "#/furniture",
|
|
"children": [],
|
|
"label": "unspecified",
|
|
},
|
|
}
|
|
return compress_json(json.dumps(doc))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migration_backfills_picture_bytes_and_strips_blob(temp_db_path):
|
|
"""Happy path: doc has a picture with bytes inline in the blob plus a
|
|
matching items row (with picture_data NULL). The migration writes the
|
|
bytes onto the items row and clears the URI from the blob."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_id = "doc-1"
|
|
# Insert a doc whose blob carries a picture data URI.
|
|
await store.documents_table.add(
|
|
[
|
|
DocumentRecord(
|
|
id=doc_id,
|
|
content="x",
|
|
docling_document=_docling_blob_with_picture(),
|
|
docling_version="1.10.0",
|
|
)
|
|
]
|
|
)
|
|
# Insert a matching items row, picture_data deliberately empty.
|
|
await store.document_items_table.add(
|
|
[
|
|
DocumentItemRecord(
|
|
document_id=doc_id,
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text="",
|
|
page_numbers="[1]",
|
|
picture_data=None,
|
|
)
|
|
]
|
|
)
|
|
|
|
await _apply_extract_picture_bytes(store)
|
|
|
|
# picture_data is populated on the items row.
|
|
rows = await (
|
|
store.document_items_table.query()
|
|
.where(f"document_id = '{doc_id}' AND self_ref = '#/pictures/0'")
|
|
.to_list()
|
|
)
|
|
assert len(rows) == 1
|
|
assert rows[0]["picture_data"] == PNG_BYTES
|
|
|
|
# The docling blob no longer carries the data URI on the picture.
|
|
doc_rows = await (
|
|
store.documents_table.query()
|
|
.where(f"id = '{doc_id}'")
|
|
.select(["docling_document"])
|
|
.to_list()
|
|
)
|
|
blob = json.loads(decompress_json(doc_rows[0]["docling_document"]))
|
|
assert blob["pictures"][0]["image"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migration_runs_against_pre_v0_48_0_schema(temp_db_path):
|
|
"""Regression: v0.45.0 must not depend on columns introduced later.
|
|
|
|
Reproduces the client report where a DB on which v0.40.0 had already
|
|
run (and thus had a ``document_items`` table without ``picture_data``,
|
|
``heading_level`` or ``tree_depth``) failed v0.45.0 with
|
|
``Field 'heading_level' not found in target schema`` because the
|
|
migration was building rows from the current ``DocumentItemRecord``
|
|
Pydantic model — which now carries fields added in v0.48.0.
|
|
"""
|
|
legacy_items_schema = pa.schema(
|
|
[
|
|
pa.field("document_id", pa.string()),
|
|
pa.field("position", pa.int64()),
|
|
pa.field("self_ref", pa.string()),
|
|
pa.field("label", pa.string()),
|
|
pa.field("text", pa.string()),
|
|
pa.field("page_numbers", pa.string()),
|
|
]
|
|
)
|
|
|
|
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
|
# Replace the freshly-created (latest-schema) document_items table
|
|
# with one that only carries the v0.40.0 columns.
|
|
await store.db.drop_table("document_items")
|
|
legacy_items = await store.db.create_table(
|
|
"document_items", schema=legacy_items_schema
|
|
)
|
|
|
|
doc_id = "doc-1"
|
|
await store.documents_table.add(
|
|
[
|
|
DocumentRecord(
|
|
id=doc_id,
|
|
content="x",
|
|
docling_document=_docling_blob_with_picture(),
|
|
docling_version="1.10.0",
|
|
)
|
|
]
|
|
)
|
|
await legacy_items.add(
|
|
[
|
|
{
|
|
"document_id": doc_id,
|
|
"position": 0,
|
|
"self_ref": "#/pictures/0",
|
|
"label": "picture",
|
|
"text": "",
|
|
"page_numbers": "[1]",
|
|
}
|
|
]
|
|
)
|
|
await store.set_haiku_version("0.40.0")
|
|
|
|
# Re-open and reach for the table fresh so we use the legacy-schema handle.
|
|
async with Store(temp_db_path, skip_migration_check=True) as store:
|
|
await _apply_extract_picture_bytes(store)
|
|
|
|
rows = await (
|
|
store.document_items_table.query()
|
|
.where(f"document_id = '{doc_id}' AND self_ref = '#/pictures/0'")
|
|
.to_list()
|
|
)
|
|
assert len(rows) == 1
|
|
assert rows[0]["picture_data"] == PNG_BYTES
|
|
|
|
# And the blob is stripped.
|
|
doc_rows = await (
|
|
store.documents_table.query()
|
|
.where(f"id = '{doc_id}'")
|
|
.select(["docling_document"])
|
|
.to_list()
|
|
)
|
|
blob = json.loads(decompress_json(doc_rows[0]["docling_document"]))
|
|
assert blob["pictures"][0]["image"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migration_is_idempotent(temp_db_path):
|
|
"""Running the migration twice on the same DB is a no-op the second
|
|
time — the blob is already stripped."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_id = "doc-1"
|
|
await store.documents_table.add(
|
|
[
|
|
DocumentRecord(
|
|
id=doc_id,
|
|
content="x",
|
|
docling_document=_docling_blob_with_picture(),
|
|
docling_version="1.10.0",
|
|
)
|
|
]
|
|
)
|
|
await store.document_items_table.add(
|
|
[
|
|
DocumentItemRecord(
|
|
document_id=doc_id,
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text="",
|
|
page_numbers="[1]",
|
|
)
|
|
]
|
|
)
|
|
|
|
await _apply_extract_picture_bytes(store)
|
|
# Capture state after first run.
|
|
rows_first = await (
|
|
store.document_items_table.query()
|
|
.where(f"document_id = '{doc_id}'")
|
|
.to_list()
|
|
)
|
|
first_bytes = rows_first[0]["picture_data"]
|
|
|
|
# Re-run.
|
|
await _apply_extract_picture_bytes(store)
|
|
|
|
rows_second = await (
|
|
store.document_items_table.query()
|
|
.where(f"document_id = '{doc_id}'")
|
|
.to_list()
|
|
)
|
|
assert rows_second[0]["picture_data"] == first_bytes
|