Merge pull request #381 from ggozad/fix/migration-schema-coupling

fix migration failures on pre-v0.48.0 document_items schemas

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
 

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 = [
+                live_schema = await store.document_items_table.schema()
-                    DocumentItemRecord(
+                records = _build_items_arrow_table(
-                        document_id=item.document_id,
+                    live_schema,
-                        position=item.position,
+                    [
-                        self_ref=item.self_ref,
+                        {
-                        label=item.label,
+                            "document_id": item.document_id,
-                        text=item.text,
+                            "position": item.position,
-                        page_numbers=json.dumps(item.page_numbers),
+                            "self_ref": item.self_ref,
-                    )
+                            "label": item.label,
-                    for item in items
+                            "text": item.text,
-                ]
+                            "page_numbers": json.dumps(item.page_numbers),
+                        }
+                        for item in items
+                    ],
+                )
                 await store.document_items_table.add(records)
 
             migrated += 1

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 = [
+                new_records = pa.Table.from_pylist(
-                    DocumentItemRecord(
+                    [
-                        document_id=doc_id,
+                        {
-                        position=existing_by_ref[ref]["position"],
+                            "document_id": doc_id,
-                        self_ref=ref,
+                            "position": existing_by_ref[ref]["position"],
-                        label=existing_by_ref[ref].get("label", ""),
+                            "self_ref": ref,
-                        text=existing_by_ref[ref].get("text", ""),
+                            "label": existing_by_ref[ref].get("label") or "",
-                        page_numbers=existing_by_ref[ref].get("page_numbers", "[]"),
+                            "text": existing_by_ref[ref].get("text") or "",
-                        picture_data=img_bytes,
+                            "page_numbers": existing_by_ref[ref].get("page_numbers")
-                    )
+                            or "[]",
-                    for ref, img_bytes in updates
+                            "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"])

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

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