haiku.rag/tests/store/test_v0_75_0_migration.py
Yiorgis Gozadinos 11644c7f43
Trim comments, docstrings and docs to what they need to say
Also drops two things that were stale rather than merely verbose: the CLI docs
note for 0.75.0, which was the only release-tagged note in the docs tree while
the CHANGELOG already records that existing databases need `haiku-rag migrate`;
and "(created or corrected)" from the migration log line, left over from the
earlier behaviour that replaced wrong-typed indexes.
2026-08-18 12:38:02 +03:00

105 lines
3.5 KiB
Python

import pytest
from lancedb.index import BTree
from haiku.rag.store.engine import Store
from haiku.rag.store.models import Document
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.upgrades.v0_75_0 import _apply_index_hot_lookup_keys
# The indexes a pre-0.75.0 database lacked.
LEGACY_DROPPED = {
"documents": ["id_idx"],
"chunks": ["id_idx", "document_id_idx"],
"document_items": ["label_idx"],
}
async def _indexed(table) -> dict[str, str]:
return {
column: index.index_type
for index in await table.list_indices()
for column in index.columns
}
async def _make_legacy(store: Store) -> None:
for table_name, indexes in LEGACY_DROPPED.items():
table = store._tables()[table_name]
for index in indexes:
await table.drop_index(index)
@pytest.mark.asyncio
async def test_adds_the_missing_indexes(temp_db_path):
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await _make_legacy(store)
await _apply_index_hot_lookup_keys(store)
assert await _indexed(store.documents_table) == {"id": "BTree"}
assert await _indexed(store.chunks_table) == {
"content_fts": "FTS",
"id": "BTree",
"document_id": "BTree",
}
assert await _indexed(store.document_items_table) == {
"document_id": "BTree",
"position": "BTree",
"self_ref": "BTree",
"label": "Bitmap",
}
@pytest.mark.asyncio
async def test_keeps_every_row(temp_db_path):
"""Indexing must not touch data."""
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
repo = DocumentRepository(store)
await repo.create(Document(content="Kept", uri="test://kept"))
await _make_legacy(store)
await _apply_index_hot_lookup_keys(store)
docs = await repo.list_all(include_content=True)
assert [d.content for d in docs] == ["Kept"]
@pytest.mark.asyncio
async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path):
"""A second run must not re-index: replace=True rebuilds."""
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await _apply_index_hot_lookup_keys(store)
versions = {name: await t.version() for name, t in store._tables().items()}
await _apply_index_hot_lookup_keys(store)
assert {name: await t.version() for name, t in store._tables().items()} == (
versions
)
@pytest.mark.asyncio
async def test_replaces_a_wrong_typed_legacy_index(temp_db_path):
"""A BTree on `label` does not satisfy the declared Bitmap."""
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await store.document_items_table.create_index(
"label", config=BTree(), replace=True
)
await _apply_index_hot_lookup_keys(store)
indexed = await _indexed(store.document_items_table)
assert indexed["label"] == "Bitmap"
@pytest.mark.asyncio
async def test_leaves_undeclared_indexes_alone(temp_db_path):
"""Indexes haiku.rag never declared are not dropped."""
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await store.document_meta_table.create_index(
"title", config=BTree(), replace=True
)
await _apply_index_hot_lookup_keys(store)
assert "title" in await _indexed(store.document_meta_table)