304 lines
11 KiB
Python
304 lines
11 KiB
Python
import re
|
|
|
|
import pytest
|
|
from lancedb.table import AsyncTable, AsyncTags
|
|
|
|
from haiku.rag.store import ReadOnlyError, Store
|
|
from haiku.rag.store.engine import RESTORE_TABLE_ORDER
|
|
from haiku.rag.store.models import Document
|
|
from haiku.rag.store.repositories.document import DocumentRepository
|
|
|
|
SAFETY_TAG_PATTERN = r"before-restore-\d{8}T\d{6}Z"
|
|
|
|
|
|
async def _doc_contents(store: Store) -> set[str]:
|
|
docs = await DocumentRepository(store).list_all(include_content=True)
|
|
return {d.content for d in docs}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_tag_restores_all_tables(temp_db_path):
|
|
"""A complete tag restores every table; rows added after the tag are
|
|
absent from the restored latest state, which stays writable."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
pre_restore_docs_version = await store.documents_table.version()
|
|
|
|
safety_tag = await store.restore_tag("release-1")
|
|
|
|
assert re.fullmatch(SAFETY_TAG_PATTERN, safety_tag)
|
|
assert await _doc_contents(store) == {"First document"}
|
|
|
|
# restore writes a NEW latest version; the table is not a read-only
|
|
# checkout and stays writable.
|
|
assert await store.documents_table.version() > pre_restore_docs_version
|
|
await repo.create(Document(content="Third document"))
|
|
assert await _doc_contents(store) == {"First document", "Third document"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_safety_tag_matches_pre_restore_state(temp_db_path):
|
|
"""The safety tag records the exact pre-restore version map, and
|
|
restoring it returns the database to its prior logical state."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
snapshot = await store.current_table_versions()
|
|
|
|
safety_tag = await store.restore_tag("release-1")
|
|
|
|
tags = await store.list_tags()
|
|
assert tags[safety_tag].complete is True
|
|
assert tags[safety_tag].tables == snapshot
|
|
|
|
await store.restore_tag(safety_tag)
|
|
assert await _doc_contents(store) == {"First document", "Second document"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_missing_tag_makes_no_changes(temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await DocumentRepository(store).create(Document(content="First document"))
|
|
versions = await store.current_table_versions()
|
|
|
|
with pytest.raises(ValueError, match="does not exist"):
|
|
await store.restore_tag("nope")
|
|
|
|
assert await store.current_table_versions() == versions
|
|
assert await store.list_tags() == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_partial_tag_makes_no_changes(temp_db_path):
|
|
"""A partial tag can never be restored; the error lists every missing
|
|
table and no safety tag is created."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
version = await store.chunks_table.version()
|
|
await store.chunks_table.tags.create("stale", version)
|
|
versions = await store.current_table_versions()
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
await store.restore_tag("stale")
|
|
|
|
msg = str(exc_info.value)
|
|
for table_name in ("documents", "document_meta", "document_items", "settings"):
|
|
assert table_name in msg
|
|
|
|
assert await store.current_table_versions() == versions
|
|
assert set(await store.list_tags()) == {"stale"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_safety_tag_name_collision(temp_db_path, monkeypatch):
|
|
"""A colliding safety-tag name gets a numeric suffix."""
|
|
import haiku.rag.store.engine as engine_mod
|
|
|
|
class FixedDatetime:
|
|
@staticmethod
|
|
def now(tz=None):
|
|
from datetime import UTC, datetime
|
|
|
|
return datetime(2026, 7, 15, 14, 30, 12, tzinfo=UTC)
|
|
|
|
monkeypatch.setattr(engine_mod, "datetime", FixedDatetime)
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await DocumentRepository(store).create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await store.create_tag("before-restore-20260715T143012Z")
|
|
|
|
safety_tag = await store.restore_tag("release-1")
|
|
assert safety_tag == "before-restore-20260715T143012Z-2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_safety_tag_failure_leaves_state_untouched(
|
|
temp_db_path, monkeypatch
|
|
):
|
|
"""If the safety tag cannot be created, restore never begins."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
versions = await store.current_table_versions()
|
|
|
|
async def failing_create(self, name: str, version: int) -> None:
|
|
raise RuntimeError("tag boom")
|
|
|
|
monkeypatch.setattr(AsyncTags, "create", failing_create)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await store.restore_tag("release-1")
|
|
|
|
msg = str(exc_info.value)
|
|
assert "did not begin" in msg
|
|
assert "No table was changed" in msg
|
|
assert "tag boom" in msg
|
|
assert exc_info.value.__cause__ is not None
|
|
|
|
monkeypatch.undo()
|
|
assert await store.current_table_versions() == versions
|
|
assert await _doc_contents(store) == {"First document", "Second document"}
|
|
assert set(await store.list_tags()) == {"release-1"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_midway_failure_rolls_back(temp_db_path, monkeypatch):
|
|
"""A restore failure after some tables were restored rolls every table
|
|
back to the pre-restore snapshot; the error names the failed table and
|
|
the safety tag."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
|
|
real_restore = AsyncTable.restore
|
|
calls = {"n": 0}
|
|
|
|
async def flaky_restore(self, version=None):
|
|
calls["n"] += 1
|
|
if calls["n"] == 3:
|
|
raise RuntimeError("restore boom")
|
|
return await real_restore(self, version)
|
|
|
|
monkeypatch.setattr(AsyncTable, "restore", flaky_restore)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await store.restore_tag("release-1")
|
|
|
|
msg = str(exc_info.value)
|
|
assert RESTORE_TABLE_ORDER[2] in msg
|
|
assert "rolled back" in msg
|
|
assert "before-restore-" in msg
|
|
|
|
monkeypatch.undo()
|
|
assert await _doc_contents(store) == {"First document", "Second document"}
|
|
assert any(t.startswith("before-restore-") for t in await store.list_tags())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_rollback_failure_reports_inconsistency(
|
|
temp_db_path, monkeypatch
|
|
):
|
|
"""When rollback also fails, the error lists the failed tables, names
|
|
the safety tag, and states manual recovery is required."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
|
|
real_restore = AsyncTable.restore
|
|
calls = {"n": 0}
|
|
|
|
async def flaky_restore(self, version=None):
|
|
calls["n"] += 1
|
|
if calls["n"] >= 3:
|
|
raise RuntimeError("restore boom")
|
|
return await real_restore(self, version)
|
|
|
|
monkeypatch.setattr(AsyncTable, "restore", flaky_restore)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await store.restore_tag("release-1")
|
|
|
|
msg = str(exc_info.value)
|
|
assert "inconsistent" in msg
|
|
assert "manual recovery" in msg
|
|
assert "before-restore-" in msg
|
|
for table_name in RESTORE_TABLE_ORDER:
|
|
assert table_name in msg
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_cancellation_rolls_back(temp_db_path, monkeypatch):
|
|
"""Cancellation mid-restore must not bypass rollback: the tables return
|
|
to the pre-restore snapshot and the cancellation re-raises."""
|
|
import asyncio
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
|
|
real_restore = AsyncTable.restore
|
|
calls = {"n": 0}
|
|
|
|
async def cancelled_restore(self, version=None):
|
|
calls["n"] += 1
|
|
if calls["n"] == 3:
|
|
raise asyncio.CancelledError()
|
|
return await real_restore(self, version)
|
|
|
|
monkeypatch.setattr(AsyncTable, "restore", cancelled_restore)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await store.restore_tag("release-1")
|
|
|
|
monkeypatch.undo()
|
|
assert await _doc_contents(store) == {"First document", "Second document"}
|
|
assert any(t.startswith("before-restore-") for t in await store.list_tags())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_cancellation_with_failed_rollback_reports(
|
|
temp_db_path, monkeypatch
|
|
):
|
|
"""If rollback after a cancellation also fails, the manual-recovery
|
|
error is raised instead of the bare cancellation."""
|
|
import asyncio
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentRepository(store)
|
|
await repo.create(Document(content="First document"))
|
|
await store.create_tag("release-1")
|
|
await repo.create(Document(content="Second document"))
|
|
|
|
real_restore = AsyncTable.restore
|
|
calls = {"n": 0}
|
|
|
|
async def broken_restore(self, version=None):
|
|
calls["n"] += 1
|
|
if calls["n"] < 3:
|
|
return await real_restore(self, version)
|
|
if calls["n"] == 3:
|
|
raise asyncio.CancelledError()
|
|
raise RuntimeError("restore boom")
|
|
|
|
monkeypatch.setattr(AsyncTable, "restore", broken_restore)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await store.restore_tag("release-1")
|
|
|
|
msg = str(exc_info.value)
|
|
assert "cancel" in msg.lower()
|
|
assert "manual recovery" in msg
|
|
assert "before-restore-" in msg
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_read_only_raises(temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await store.create_tag("release-1")
|
|
|
|
async with Store(temp_db_path, read_only=True) as store:
|
|
with pytest.raises(ReadOnlyError):
|
|
await store.restore_tag("release-1")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_rejected_during_rebuild(temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await store.create_tag("release-1")
|
|
|
|
async with store._rebuild_lock:
|
|
with pytest.raises(ValueError, match="[Rr]ebuild in progress"):
|
|
await store.restore_tag("release-1")
|