diff --git a/CHANGELOG.md b/CHANGELOG.md index df8b76f1..e2edc3bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed - mxbai reranker crashing inside the chat TUI with `ValueError: bad value(s) in fds_to_keep`. tqdm constructs a `multiprocessing.RLock` on first use, whose `resource_tracker` spawn picks up `sys.stderr.fileno()`; Textual's redirected stderr returns `-1`, failing the `fork_exec` validation. The reranker now pins tqdm's class lock to a `threading.RLock`. +- `migrate` failing on DBs upgrading through v0.45.0 with `Field 'heading_level' not found in target schema`. The v0.45.0 picture-data backfill was building rows from the current `DocumentItemRecord` Pydantic model — which now carries `heading_level` / `tree_depth` added in v0.48.0 — and feeding them to `merge_insert` against a pre-v0.48.0 schema. Both v0.40.0 and v0.45.0 now build their PyArrow inputs from explicit per-migration column sets so future model additions can't retroactively break them. ## [0.48.0] - 2026-05-20 diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py index 3e61aaad..a0ebcaa3 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py @@ -1,12 +1,72 @@ import json import logging -from haiku.rag.store.engine import DocumentItemRecord, Store +import pyarrow as pa + +from haiku.rag.store.engine import Store from haiku.rag.store.upgrades import Upgrade from haiku.rag.utils import escape_sql_string logger = logging.getLogger(__name__) +# Columns v0.40.0 owns. Later migrations add more (``picture_data`` in v0.45.0, +# ``heading_level`` / ``tree_depth`` in v0.48.0); those are filled with +# type-appropriate defaults at insert time so the migration stays independent +# of future model changes. +_V0_40_0_OWNED_COLUMNS: frozenset[str] = frozenset( + {"document_id", "position", "self_ref", "label", "text", "page_numbers"} +) + + +def _default_for(field: pa.Field): + """Return a sensible placeholder for a column v0.40.0 doesn't write. + + Nullable columns get ``None``; non-nullable columns get the Arrow-typed + zero (``0`` for ints, ``""`` for strings, ``b""`` for binary). This + matches the Pydantic ``Field(default=...)`` values on the live model. + """ + if field.nullable: + return None + if pa.types.is_integer(field.type) or pa.types.is_floating(field.type): + return 0 + # Defensive branches for column types not yet introduced after v0.40.0. + # Kept so a future non-nullable string / binary / bool addition doesn't + # turn the migration into a TypeError on real DBs. + if pa.types.is_string(field.type) or pa.types.is_large_string( + field.type + ): # pragma: no cover + return "" + if pa.types.is_binary(field.type) or pa.types.is_large_binary( + field.type + ): # pragma: no cover + return b"" + if pa.types.is_boolean(field.type): # pragma: no cover + return False + raise TypeError( # pragma: no cover + f"No default for non-nullable arrow type {field.type}" + ) + + +def _build_items_arrow_table(live_schema: pa.Schema, rows: list[dict]) -> pa.Table: + """Materialise ``rows`` as a PyArrow table matching ``live_schema``. + + ``rows`` provides values for the v0.40.0-owned columns; any additional + columns the live schema carries (added by later migrations) are filled + with type-appropriate defaults so ``Table.add`` accepts the input even + when those columns are non-nullable. + """ + columns: dict[str, pa.Array] = {} + for field in live_schema: + if field.name in _V0_40_0_OWNED_COLUMNS: + columns[field.name] = pa.array( + [row[field.name] for row in rows], type=field.type + ) + else: + columns[field.name] = pa.array( + [_default_for(field)] * len(rows), type=field.type + ) + return pa.table(columns, schema=live_schema) + async def _apply_populate_document_items(store: Store) -> None: # pragma: no cover """Populate document_items table from existing docling documents.""" @@ -55,17 +115,21 @@ async def _apply_populate_document_items(store: Store) -> None: # pragma: no co items = extract_items(doc_id, docling_doc) if items: - records = [ - DocumentItemRecord( - document_id=item.document_id, - position=item.position, - self_ref=item.self_ref, - label=item.label, - text=item.text, - page_numbers=json.dumps(item.page_numbers), - ) - for item in items - ] + live_schema = await store.document_items_table.schema() + records = _build_items_arrow_table( + live_schema, + [ + { + "document_id": item.document_id, + "position": item.position, + "self_ref": item.self_ref, + "label": item.label, + "text": item.text, + "page_numbers": json.dumps(item.page_numbers), + } + for item in items + ], + ) await store.document_items_table.add(records) migrated += 1 diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py index e97ef0c6..4cd6e781 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py @@ -5,7 +5,7 @@ import logging import pyarrow as pa -from haiku.rag.store.engine import DocumentItemRecord, Store +from haiku.rag.store.engine import Store from haiku.rag.store.upgrades import Upgrade from haiku.rag.utils import escape_sql_string @@ -13,6 +13,21 @@ logger = logging.getLogger(__name__) PROGRESS_INTERVAL = 5 +# Schema for the merge_insert input. Pinned to the columns that exist at +# v0.45.0 time so the migration stays independent of future model changes +# (e.g. heading_level / tree_depth added in v0.48.0). +_V0_45_0_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()), + pa.field("picture_data", pa.large_binary()), + ] +) + async def _ensure_picture_data_column(store: Store) -> None: """Add ``document_items.picture_data`` (large_binary) if it doesn't exist. @@ -125,18 +140,22 @@ async def _apply_extract_picture_bytes(store: Store) -> None: ) existing_by_ref = {r["self_ref"]: r for r in existing_items} - new_records = [ - DocumentItemRecord( - document_id=doc_id, - position=existing_by_ref[ref]["position"], - self_ref=ref, - label=existing_by_ref[ref].get("label", ""), - text=existing_by_ref[ref].get("text", ""), - page_numbers=existing_by_ref[ref].get("page_numbers", "[]"), - picture_data=img_bytes, - ) - for ref, img_bytes in updates - ] + new_records = pa.Table.from_pylist( + [ + { + "document_id": doc_id, + "position": existing_by_ref[ref]["position"], + "self_ref": ref, + "label": existing_by_ref[ref].get("label") or "", + "text": existing_by_ref[ref].get("text") or "", + "page_numbers": existing_by_ref[ref].get("page_numbers") + or "[]", + "picture_data": img_bytes, + } + for ref, img_bytes in updates + ], + schema=_V0_45_0_ITEMS_SCHEMA, + ) await ( store.document_items_table.merge_insert(["document_id", "self_ref"]) diff --git a/tests/store/test_v0_40_0_migration.py b/tests/store/test_v0_40_0_migration.py new file mode 100644 index 00000000..e11f630e --- /dev/null +++ b/tests/store/test_v0_40_0_migration.py @@ -0,0 +1,65 @@ +"""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 import Store +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 + + +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_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 diff --git a/tests/store/test_v0_45_0_migration.py b/tests/store/test_v0_45_0_migration.py index efdd12b3..6823a2e0 100644 --- a/tests/store/test_v0_45_0_migration.py +++ b/tests/store/test_v0_45_0_migration.py @@ -9,6 +9,7 @@ documents table). import base64 import json +import pyarrow as pa import pytest from haiku.rag.store import Store @@ -101,6 +102,84 @@ async def test_migration_backfills_picture_bytes_and_strips_blob(temp_db_path): 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