danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #587 from ggozad/fix/v0-38-0-migration-zstd

Browse source 
Accept zstd docling blobs in the 0.38.0 migration
This commit is contained in:
Yiorgis Gozadinos 2026-08-31 04:29:51 -05:00 committed by GitHub
commit b87ae16910
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 737 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -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

2
haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py
View file

@ -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

6
haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py
View file

@ -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",
)

17
haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py
View file

@ -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")

84
tests/store/legacy_documents.py
View file

@ -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)

116
tests/store/test_v0_20_0_migration.py Normal file
View file

@ -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

234
tests/store/test_v0_25_0_migration.py Normal file
View file

@ -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

302
tests/store/test_v0_38_0_migration.py Normal file
View file

@ -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