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

Test the 0.20.0 and 0.25.0 migrations

Browse source 
Both were exempt from coverage, which is how the 0.38.0 blob-encoding bug
reached a release: nothing exercised the chain that produced it.

Move the historical `documents` shapes into tests/store/legacy_documents.py
so each migration's tests can seed the table as its predecessor left it.
This commit is contained in:
Yiorgis Gozadinos 2026-08-31 11:35:04 +03:00
parent 6133e5ff4a
commit e0bde9d59a
No known key found for this signature in database
6 changed files with 453 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

96
tests/store/test_v0_38_0_migration.py
View file

@ -9,14 +9,17 @@ import gzip
import json
import lancedb
import pyarrow as pa
import pytest
from lancedb.pydantic import LanceModel
from pydantic import Field
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"
@ -43,47 +46,8 @@ def _docling_doc(name: str = "test", with_pages: bool = True) -> dict:
return doc
class DocumentRecordV4(LanceModel):
id: str
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 DocumentRecordV5(DocumentRecordV4):
docling_pages: bytes | None = None
def _large_binary_schema(model: type[LanceModel]) -> pa.Schema:
blobs = {"docling_document", "docling_pages"}
return pa.schema(
[
pa.field(field.name, pa.large_binary()) if field.name in blobs else field
for field in model.to_arrow_schema()
]
)
def _v4_schema() -> pa.Schema:
return _large_binary_schema(DocumentRecordV4)
def _v5_schema() -> pa.Schema:
return _large_binary_schema(DocumentRecordV5)
async def _make_v4_documents_table(store: Store) -> None:
"""Replace the live documents table with the pre-0.38.0 schema."""
del store.documents_table
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=_v4_schema()
)
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:
@ -107,8 +71,8 @@ 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 _make_v4_documents_table(store)
await store.documents_table.add(
await _seed_v4(
store,
[
DocumentRecordV4(
id="doc-1",
@ -121,7 +85,7 @@ async def test_migrates_every_v0_25_0_blob_encoding(temp_db_path, encode):
created_at="2026-01-01",
updated_at="2026-01-02",
)
]
],
)
await _apply_split_pages_zstd(store)
@ -143,8 +107,8 @@ async def test_migrates_every_v0_25_0_blob_encoding(temp_db_path, encode):
@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 _make_v4_documents_table(store)
await store.documents_table.add(
await _seed_v4(
store,
[
DocumentRecordV4(
id="doc-1",
@ -154,7 +118,7 @@ async def test_document_without_pages_gets_null_pages_column(temp_db_path):
),
),
DocumentRecordV4(id="doc-2", content="no blob"),
]
],
)
await _apply_split_pages_zstd(store)
@ -172,8 +136,8 @@ async def test_document_without_pages_gets_null_pages_column(temp_db_path):
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 _make_v4_documents_table(store)
await store.documents_table.add(
await _seed_v4(
store,
[
DocumentRecordV4(
id=f"doc-{n}",
@ -183,7 +147,7 @@ async def test_migrates_batches_larger_than_batch_size(temp_db_path):
),
)
for n in range(12)
]
],
)
await _apply_split_pages_zstd(store)
@ -201,8 +165,10 @@ async def test_migrates_batches_larger_than_batch_size(temp_db_path):
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 _make_v4_documents_table(store)
await store.db.create_table(STAGING, schema=_v4_schema())
await _seed_v4(store, [])
await store.db.create_table(
STAGING, schema=documents_schema(LegacyDocumentRecord)
)
await store.documents_table.add(
[
DocumentRecordV4(
@ -232,8 +198,10 @@ async def test_recovers_documents_from_staging_when_documents_table_is_empty(
structure = compress_json(json.dumps(doc))
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await _make_v4_documents_table(store)
staging = await store.db.create_table(STAGING, schema=_v5_schema())
await _seed_v4(store, [])
staging = await store.db.create_table(
STAGING, schema=documents_schema(LegacyDocumentRecord)
)
await staging.add(
[
{
@ -276,8 +244,10 @@ async def test_unreadable_documents_table_falls_back_to_staging(
return original(self)
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await _make_v4_documents_table(store)
staging = await store.db.create_table(STAGING, schema=_v5_schema())
await _seed_v4(store, [])
staging = await store.db.create_table(
STAGING, schema=documents_schema(LegacyDocumentRecord)
)
await staging.add(
[
{
@ -308,7 +278,7 @@ async def test_unreadable_documents_table_falls_back_to_staging(
@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 _make_v4_documents_table(store)
await _seed_v4(store, [])
await _apply_split_pages_zstd(store)
@ -320,8 +290,10 @@ async def test_empty_database_is_rebuilt_on_the_new_schema(temp_db_path):
@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 _make_v4_documents_table(store)
await store.db.create_table(STAGING, schema=_v5_schema())
await _seed_v4(store, [])
await store.db.create_table(
STAGING, schema=documents_schema(LegacyDocumentRecord)
)
await _apply_split_pages_zstd(store)