diff --git a/CHANGELOG.md b/CHANGELOG.md index 482ec489..e0f4231a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Fixed + +- Migration to 0.38.0 no longer fails with `UnicodeDecodeError` on a + `docling_document` blob written as zstd. + ## [0.79.0] - 2026-08-28 ### Changed diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py index 5d23d8ad..c20cb563 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py @@ -7,7 +7,7 @@ from haiku.rag.store.engine import Store from haiku.rag.store.upgrades import Upgrade -async def _apply_add_docling_document_columns(store: Store) -> None: # pragma: no cover +async def _apply_add_docling_document_columns(store: Store) -> None: """Add 'docling_document_json' and 'docling_version' columns to documents table.""" # Read existing rows using Arrow for schema-agnostic access diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py index 04ed0c12..47558ae3 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) BATCH_SIZE = 10 -async def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover +async def _apply_compress_docling_document(store: Store) -> None: """Migrate docling_document_json (str) to docling_document (compressed bytes).""" class DocumentRecordV4(LanceModel): @@ -223,7 +223,7 @@ async def _apply_compress_docling_document(store: Store) -> None: # pragma: no for table in [store.documents_table, store.chunks_table, store.settings_table]: try: await table.optimize(cleanup_older_than=timedelta(seconds=0)) - except Exception: + except Exception: # pragma: no cover - vacuum failure must not fail migration pass logger.info("Migration complete") @@ -232,5 +232,5 @@ async def _apply_compress_docling_document(store: Store) -> None: # pragma: no upgrade_compress_docling_document = Upgrade( version="0.25.0", apply=_apply_compress_docling_document, - description="Compress docling_document with gzip and use large_binary type", + description="Compress docling_document and use large_binary type", ) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py index 4fe74222..2d9a3b5e 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py @@ -7,7 +7,7 @@ import pyarrow as pa from lancedb.pydantic import LanceModel from pydantic import Field -from haiku.rag.store.compression import compress_docling_split +from haiku.rag.store.compression import compress_docling_split, decompress_json from haiku.rag.store.engine import Store from haiku.rag.store.upgrades import Upgrade @@ -16,7 +16,7 @@ logger = logging.getLogger(__name__) BATCH_SIZE = 5 -async def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover +async def _apply_split_pages_zstd(store: Store) -> None: """Split docling_document into structure + pages and re-compress with zstd.""" class DocumentRecordV5(LanceModel): @@ -44,18 +44,21 @@ async def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover return pa.schema(fields) def migrate_row(row: dict) -> DocumentRecordV5: - """Migrate a single row: decompress gzip, split pages, re-compress with zstd.""" + """Migrate a single row: decompress, split pages, re-compress with zstd.""" docling_blob = row.get("docling_document") structure_bytes: bytes | None = None pages_bytes: bytes | None = None if docling_blob and isinstance(docling_blob, bytes): - # Decompress from gzip + # v0.25.0 blobs are gzip, zstd or uncompressed, depending on the + # version that wrote them. try: json_str = gzip.decompress(docling_blob).decode("utf-8") except Exception: - # May already be zstd or uncompressed — try as-is - json_str = docling_blob.decode("utf-8") + try: + json_str = decompress_json(docling_blob) + except Exception: + json_str = docling_blob.decode("utf-8") # Split structure and pages, re-compress with zstd structure_bytes, pages_bytes = compress_docling_split(json.loads(json_str)) @@ -222,7 +225,7 @@ async def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover for table in [store.documents_table, store.chunks_table, store.settings_table]: try: await table.optimize(cleanup_older_than=timedelta(seconds=0)) - except Exception: + except Exception: # pragma: no cover - vacuum failure must not fail migration pass logger.info("Migration complete") diff --git a/tests/store/legacy_documents.py b/tests/store/legacy_documents.py index f6ba2781..49cf0001 100644 --- a/tests/store/legacy_documents.py +++ b/tests/store/legacy_documents.py @@ -1,10 +1,11 @@ -"""Helpers to seed a `documents` table in its pre-0.58 shape. +"""Helpers to seed a `documents` table in the shapes it had before 0.58. -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. +Each migration in the early chain rewrites the whole `documents` table, so its +tests need the table as the *previous* version left it. The record classes here +are those frozen shapes, named for the schema version the migration chain gives +them: V2 predates docling, V3 carries it as JSON text (v0.20.0), V4 as one +compressed blob (v0.25.0), and `LegacyDocumentRecord` is V5, the split-blob +shape that stood until the document_meta split (v0.58.0). """ from uuid import uuid4 @@ -16,6 +17,46 @@ from pydantic import Field from haiku.rag.store.engine import Store +class DocumentRecordV2(LanceModel): + """The pre-0.20 `documents` record, before any docling column.""" + + id: str = Field(default_factory=lambda: str(uuid4())) + content: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + +class DocumentRecordV3(LanceModel): + """The 0.20.0 `documents` record, with docling stored as JSON text.""" + + id: str = Field(default_factory=lambda: str(uuid4())) + content: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + docling_document_json: str | None = None + docling_version: str | None = None + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + +class DocumentRecordV4(LanceModel): + """The 0.25.0 `documents` record, with docling as one compressed blob.""" + + 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_version: str | None = None + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + class LegacyDocumentRecord(LanceModel): """The pre-0.58 `documents` record (mutable attributes still inline).""" @@ -31,26 +72,33 @@ class LegacyDocumentRecord(LanceModel): 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"} +def documents_schema(model: type[LanceModel]) -> pa.Schema: + """Arrow schema for a historical record, with large_binary docling columns.""" + blobs = {"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 + pa.field(f.name, pa.large_binary()) if f.name in blobs else f + for f in model.to_arrow_schema() ] ) +def legacy_documents_schema() -> pa.Schema: + return documents_schema(LegacyDocumentRecord) + + +async def seed_documents(store: Store, schema: pa.Schema, records: list) -> None: + """Recreate the `documents` table with a historical schema and add records.""" + 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=schema) + if records: + await store.documents_table.add(records) + + 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) + await seed_documents(store, legacy_documents_schema(), records) diff --git a/tests/store/test_v0_20_0_migration.py b/tests/store/test_v0_20_0_migration.py new file mode 100644 index 00000000..1bc8e09a --- /dev/null +++ b/tests/store/test_v0_20_0_migration.py @@ -0,0 +1,116 @@ +"""Tests for the v0.20.0 docling-column migration. + +The migration rebuilds `documents` with `docling_document_json` and +`docling_version` added, carrying the existing rows across as NULL in both. +""" + +import pyarrow as pa +import pytest + +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades.v0_20_0 import _apply_add_docling_document_columns +from tests.store.legacy_documents import ( + DocumentRecordV2, + documents_schema, + seed_documents, +) + + +async def _seed_v2(store: Store, records: list[DocumentRecordV2]) -> None: + await seed_documents(store, documents_schema(DocumentRecordV2), records) + + +@pytest.mark.asyncio +async def test_existing_rows_survive_with_empty_docling_columns(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v2( + store, + [ + DocumentRecordV2( + id="doc-1", + content="hello", + uri="test://doc-1", + title="Doc 1", + metadata='{"k": "v"}', + created_at="2026-01-01", + updated_at="2026-01-02", + ), + DocumentRecordV2(id="doc-2", content="second"), + ], + ) + + await _apply_add_docling_document_columns(store) + + rows = sorted( + await store.documents_table.query().to_list(), key=lambda r: r["id"] + ) + + assert [row["id"] for row in rows] == ["doc-1", "doc-2"] + assert rows[0]["content"] == "hello" + assert rows[0]["uri"] == "test://doc-1" + assert rows[0]["title"] == "Doc 1" + assert rows[0]["metadata"] == '{"k": "v"}' + assert rows[0]["created_at"] == "2026-01-01" + assert rows[0]["updated_at"] == "2026-01-02" + for row in rows: + assert row["docling_document_json"] is None + assert row["docling_version"] is None + + +@pytest.mark.asyncio +async def test_null_metadata_becomes_an_empty_json_object(temp_db_path): + """`metadata` is non-nullable from 0.20.0 on, so a NULL must be coerced.""" + nullable_metadata = pa.schema( + [ + pa.field(f.name, f.type, nullable=f.nullable or f.name == "metadata") + for f in documents_schema(DocumentRecordV2) + ] + ) + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_documents(store, nullable_metadata, []) + await store.documents_table.add( + [ + { + "id": "doc-1", + "content": "hello", + "uri": None, + "title": None, + "metadata": None, + "created_at": "", + "updated_at": "", + } + ] + ) + + await _apply_add_docling_document_columns(store) + + rows = await store.documents_table.query().to_list() + + assert rows[0]["metadata"] == "{}" + + +@pytest.mark.asyncio +async def test_empty_database_is_rebuilt_on_the_new_schema(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v2(store, []) + + await _apply_add_docling_document_columns(store) + + names = {field.name for field in await store.documents_table.schema()} + rows = await store.documents_table.query().to_list() + + assert {"docling_document_json", "docling_version"} <= names + assert rows == [] + + +@pytest.mark.asyncio +async def test_missing_documents_table_is_created(temp_db_path): + """Reruns after an interrupted migration find no documents table at all.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await store.db.drop_table("documents") + + await _apply_add_docling_document_columns(store) + + names = {field.name for field in await store.documents_table.schema()} + + assert {"docling_document_json", "docling_version"} <= names diff --git a/tests/store/test_v0_25_0_migration.py b/tests/store/test_v0_25_0_migration.py new file mode 100644 index 00000000..c6c467f2 --- /dev/null +++ b/tests/store/test_v0_25_0_migration.py @@ -0,0 +1,234 @@ +"""Tests for the v0.25.0 docling-compression migration. + +The migration replaces the `docling_document_json` text column with a +compressed `docling_document` blob. Rerunning it over an already-compressed +table must leave those blobs untouched. +""" + +import json + +import pytest + +from haiku.rag.store.compression import compress_json, decompress_json +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades.v0_25_0 import _apply_compress_docling_document +from tests.store.legacy_documents import ( + DocumentRecordV3, + DocumentRecordV4, + documents_schema, + seed_documents, +) + +STAGING = "documents_v4_staging" +DOC_JSON = json.dumps({"schema_name": "DoclingDocument", "name": "test"}) + + +async def _seed_v3(store: Store, records: list[DocumentRecordV3]) -> None: + await seed_documents(store, documents_schema(DocumentRecordV3), records) + + +async def _seed_v4(store: Store, records: list[DocumentRecordV4]) -> None: + await seed_documents(store, documents_schema(DocumentRecordV4), records) + + +async def _read_migrated(store: Store, doc_id: str) -> dict: + rows = await store.documents_table.query().where(f"id = '{doc_id}'").to_list() + assert len(rows) == 1 + return rows[0] + + +@pytest.mark.asyncio +async def test_json_text_column_becomes_a_compressed_blob(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3( + store, + [ + DocumentRecordV3( + id="doc-1", + content="hello", + uri="test://doc-1", + title="Doc 1", + metadata='{"k": "v"}', + docling_document_json=DOC_JSON, + docling_version="1.10.0", + created_at="2026-01-01", + updated_at="2026-01-02", + ), + DocumentRecordV3(id="doc-2", content="no docling"), + ], + ) + + await _apply_compress_docling_document(store) + + with_docling = await _read_migrated(store, "doc-1") + without_docling = await _read_migrated(store, "doc-2") + + assert decompress_json(with_docling["docling_document"]) == DOC_JSON + assert with_docling["uri"] == "test://doc-1" + assert with_docling["title"] == "Doc 1" + assert with_docling["metadata"] == '{"k": "v"}' + assert with_docling["docling_version"] == "1.10.0" + assert with_docling["created_at"] == "2026-01-01" + assert with_docling["updated_at"] == "2026-01-02" + assert without_docling["docling_document"] is None + + +@pytest.mark.asyncio +async def test_already_compressed_blobs_are_left_byte_identical(temp_db_path): + """Rerunning over a migrated table must not re-compress what it finds.""" + blob = compress_json(DOC_JSON) + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4( + store, + [DocumentRecordV4(id="doc-1", content="hello", docling_document=blob)], + ) + + await _apply_compress_docling_document(store) + + row = await _read_migrated(store, "doc-1") + + assert row["docling_document"] == blob + + +@pytest.mark.asyncio +async def test_uncompressed_blob_is_compressed(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4( + store, + [ + DocumentRecordV4( + id="doc-1", + content="hello", + docling_document=DOC_JSON.encode("utf-8"), + ) + ], + ) + + await _apply_compress_docling_document(store) + + row = await _read_migrated(store, "doc-1") + + assert decompress_json(row["docling_document"]) == DOC_JSON + + +@pytest.mark.asyncio +async def test_migrates_batches_larger_than_batch_size(temp_db_path): + """BATCH_SIZE is 10; the staging round-trip must carry every document.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3( + store, + [ + DocumentRecordV3( + id=f"doc-{n}", + content=f"body {n}", + docling_document_json=json.dumps({"name": f"doc-{n}"}), + ) + for n in range(23) + ], + ) + + await _apply_compress_docling_document(store) + + rows = await store.documents_table.query().to_list() + assert STAGING not in (await store.db.list_tables()).tables + + assert {row["id"] for row in rows} == {f"doc-{n}" for n in range(23)} + for row in rows: + assert json.loads(decompress_json(row["docling_document"]))["name"] == row["id"] + + +@pytest.mark.asyncio +async def test_stale_staging_table_is_replaced(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3( + store, + [ + DocumentRecordV3( + id="doc-1", content="hello", docling_document_json=DOC_JSON + ) + ], + ) + await store.db.create_table(STAGING, schema=documents_schema(DocumentRecordV4)) + + await _apply_compress_docling_document(store) + + rows = await store.documents_table.query().to_list() + assert STAGING not in (await store.db.list_tables()).tables + + assert [row["id"] for row in rows] == ["doc-1"] + + +@pytest.mark.asyncio +async def test_recovers_documents_from_staging_when_documents_table_is_empty( + temp_db_path, +): + """An interrupted run can leave the documents table emptied and every + migrated row in staging; the rerun must adopt staging rather than drop it.""" + blob = compress_json(DOC_JSON) + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3(store, []) + staging = await store.db.create_table( + STAGING, schema=documents_schema(DocumentRecordV4) + ) + await staging.add( + [DocumentRecordV4(id="doc-1", content="hello", docling_document=blob)] + ) + + await _apply_compress_docling_document(store) + + row = await _read_migrated(store, "doc-1") + assert STAGING not in (await store.db.list_tables()).tables + + assert row["docling_document"] == blob + + +@pytest.mark.asyncio +async def test_unreadable_documents_table_falls_back_to_staging(temp_db_path): + """A documents table without an `id` column cannot be enumerated.""" + blob = compress_json(DOC_JSON) + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await seed_documents( + store, + documents_schema(DocumentRecordV3).remove(0), + [], + ) + staging = await store.db.create_table( + STAGING, schema=documents_schema(DocumentRecordV4) + ) + await staging.add( + [DocumentRecordV4(id="doc-1", content="hello", docling_document=blob)] + ) + + await _apply_compress_docling_document(store) + + row = await _read_migrated(store, "doc-1") + assert STAGING not in (await store.db.list_tables()).tables + + assert row["docling_document"] == blob + + +@pytest.mark.asyncio +async def test_empty_database_is_rebuilt_on_the_new_schema(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3(store, []) + + await _apply_compress_docling_document(store) + + names = {field.name for field in await store.documents_table.schema()} + + assert "docling_document" in names + assert "docling_document_json" not in names + + +@pytest.mark.asyncio +async def test_empty_staging_table_is_not_mistaken_for_recovery(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v3(store, []) + await store.db.create_table(STAGING, schema=documents_schema(DocumentRecordV4)) + + await _apply_compress_docling_document(store) + + names = {field.name for field in await store.documents_table.schema()} + + assert "docling_document" in names + assert "docling_document_json" not in names diff --git a/tests/store/test_v0_38_0_migration.py b/tests/store/test_v0_38_0_migration.py new file mode 100644 index 00000000..b3d30f5c --- /dev/null +++ b/tests/store/test_v0_38_0_migration.py @@ -0,0 +1,302 @@ +"""Tests for the v0.38.0 page-splitting migration. + +The migration reads the single ``docling_document`` blob written by v0.25.0, +splits page images into ``docling_pages``, and re-compresses both with zstd. +That blob reaches it gzip-compressed, zstd-compressed or uncompressed. +""" + +import gzip +import json + +import lancedb +import pytest + +from haiku.rag.store.compression import compress_json, decompress_json +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades.v0_38_0 import _apply_split_pages_zstd +from tests.store.legacy_documents import ( + DocumentRecordV4, + LegacyDocumentRecord, + documents_schema, + seed_documents, +) + +STAGING = "documents_v5_staging" + + +def _docling_doc(name: str = "test", with_pages: bool = True) -> dict: + """A minimal DoclingDocument dict carrying one page image.""" + doc: dict = { + "schema_name": "DoclingDocument", + "version": "1.10.0", + "name": name, + "texts": [], + "tables": [], + "pictures": [], + "groups": [], + "body": {"self_ref": "#/body", "children": [], "label": "unspecified"}, + "furniture": { + "self_ref": "#/furniture", + "children": [], + "label": "unspecified", + }, + } + if with_pages: + doc["pages"] = {"1": {"page_no": 1, "size": {"width": 10.0, "height": 20.0}}} + return doc + + +async def _seed_v4(store: Store, records: list[DocumentRecordV4]) -> None: + await seed_documents(store, documents_schema(DocumentRecordV4), records) + + +async def _read_migrated(store: Store, doc_id: str) -> dict: + rows = await store.documents_table.query().where(f"id = '{doc_id}'").to_list() + assert len(rows) == 1 + return rows[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "encode", + [ + pytest.param( + lambda doc: gzip.compress(json.dumps(doc).encode("utf-8")), id="gzip" + ), + pytest.param(lambda doc: compress_json(json.dumps(doc)), id="zstd"), + pytest.param(lambda doc: json.dumps(doc).encode("utf-8"), id="uncompressed"), + ], +) +async def test_migrates_every_v0_25_0_blob_encoding(temp_db_path, encode): + """Every encoding a v0.25.0 database can carry migrates to zstd.""" + doc = _docling_doc() + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4( + store, + [ + DocumentRecordV4( + id="doc-1", + content="hello", + uri="test://doc-1", + title="Doc 1", + metadata='{"k": "v"}', + docling_document=encode(doc), + docling_version="1.10.0", + created_at="2026-01-01", + updated_at="2026-01-02", + ) + ], + ) + + await _apply_split_pages_zstd(store) + + row = await _read_migrated(store, "doc-1") + + structure = json.loads(decompress_json(row["docling_document"])) + assert structure["name"] == "test" + assert "pages" not in structure + assert json.loads(decompress_json(row["docling_pages"]))["1"]["page_no"] == 1 + assert row["uri"] == "test://doc-1" + assert row["title"] == "Doc 1" + assert row["metadata"] == '{"k": "v"}' + assert row["docling_version"] == "1.10.0" + assert row["created_at"] == "2026-01-01" + assert row["updated_at"] == "2026-01-02" + + +@pytest.mark.asyncio +async def test_document_without_pages_gets_null_pages_column(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4( + store, + [ + DocumentRecordV4( + id="doc-1", + content="hello", + docling_document=compress_json( + json.dumps(_docling_doc(with_pages=False)) + ), + ), + DocumentRecordV4(id="doc-2", content="no blob"), + ], + ) + + await _apply_split_pages_zstd(store) + + with_blob = await _read_migrated(store, "doc-1") + without_blob = await _read_migrated(store, "doc-2") + + assert with_blob["docling_pages"] is None + assert json.loads(decompress_json(with_blob["docling_document"]))["name"] == "test" + assert without_blob["docling_document"] is None + assert without_blob["docling_pages"] is None + + +@pytest.mark.asyncio +async def test_migrates_batches_larger_than_batch_size(temp_db_path): + """BATCH_SIZE is 5; the staging round-trip must carry every document.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4( + store, + [ + DocumentRecordV4( + id=f"doc-{n}", + content=f"body {n}", + docling_document=compress_json( + json.dumps(_docling_doc(name=f"doc-{n}")) + ), + ) + for n in range(12) + ], + ) + + await _apply_split_pages_zstd(store) + + rows = await store.documents_table.query().to_list() + assert STAGING not in (await store.db.list_tables()).tables + + assert {row["id"] for row in rows} == {f"doc-{n}" for n in range(12)} + for row in rows: + structure = json.loads(decompress_json(row["docling_document"])) + assert structure["name"] == row["id"] + + +@pytest.mark.asyncio +async def test_stale_staging_table_is_replaced(temp_db_path): + """A staging table left by an interrupted run is dropped, not appended to.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4(store, []) + await store.db.create_table( + STAGING, schema=documents_schema(LegacyDocumentRecord) + ) + await store.documents_table.add( + [ + DocumentRecordV4( + id="doc-1", + content="hello", + docling_document=compress_json(json.dumps(_docling_doc())), + ) + ] + ) + + await _apply_split_pages_zstd(store) + + rows = await store.documents_table.query().to_list() + assert STAGING not in (await store.db.list_tables()).tables + + assert [row["id"] for row in rows] == ["doc-1"] + + +@pytest.mark.asyncio +async def test_recovers_documents_from_staging_when_documents_table_is_empty( + temp_db_path, +): + """An interrupted run can leave the documents table emptied and every + migrated row in staging; the rerun must adopt staging rather than drop it.""" + doc = _docling_doc() + pages = compress_json(json.dumps(doc.pop("pages"))) + structure = compress_json(json.dumps(doc)) + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4(store, []) + staging = await store.db.create_table( + STAGING, schema=documents_schema(LegacyDocumentRecord) + ) + await staging.add( + [ + { + "id": "doc-1", + "content": "hello", + "uri": "test://doc-1", + "title": "Doc 1", + "metadata": "{}", + "docling_document": structure, + "docling_pages": pages, + "docling_version": "1.10.0", + "created_at": "", + "updated_at": "", + } + ] + ) + + await _apply_split_pages_zstd(store) + + row = await _read_migrated(store, "doc-1") + assert STAGING not in (await store.db.list_tables()).tables + + assert row["docling_document"] == structure + assert row["docling_pages"] == pages + + +@pytest.mark.asyncio +async def test_unreadable_documents_table_falls_back_to_staging( + temp_db_path, monkeypatch +): + """An unreadable documents table falls back to adopting staging.""" + structure = compress_json(json.dumps(_docling_doc(with_pages=False))) + reads: list[str] = [] + original = lancedb.AsyncTable.query + + def failing_query(self): + reads.append(self.name) + if self.name == "documents" and reads.count("documents") == 1: + raise OSError("simulated read failure") + return original(self) + + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4(store, []) + staging = await store.db.create_table( + STAGING, schema=documents_schema(LegacyDocumentRecord) + ) + await staging.add( + [ + { + "id": "doc-1", + "content": "hello", + "uri": None, + "title": None, + "metadata": "{}", + "docling_document": structure, + "docling_pages": None, + "docling_version": None, + "created_at": "", + "updated_at": "", + } + ] + ) + + monkeypatch.setattr(lancedb.AsyncTable, "query", failing_query) + await _apply_split_pages_zstd(store) + monkeypatch.undo() + + row = await _read_migrated(store, "doc-1") + assert STAGING not in (await store.db.list_tables()).tables + + assert row["docling_document"] == structure + + +@pytest.mark.asyncio +async def test_empty_database_is_rebuilt_on_the_new_schema(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4(store, []) + + await _apply_split_pages_zstd(store) + + names = {field.name for field in await store.documents_table.schema()} + + assert "docling_pages" in names + + +@pytest.mark.asyncio +async def test_empty_staging_table_is_not_mistaken_for_recovery(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _seed_v4(store, []) + await store.db.create_table( + STAGING, schema=documents_schema(LegacyDocumentRecord) + ) + + await _apply_split_pages_zstd(store) + + names = {field.name for field in await store.documents_table.schema()} + + assert "docling_pages" in names