From f41ec379dfca6aa41aabb6c96109f68faed6896b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 15 Jul 2026 16:30:46 +0300 Subject: [PATCH] Add Store.restore_tag --- haiku_rag_slim/haiku/rag/store/engine.py | 194 ++++++++++++--- tests/store/test_restore.py | 304 +++++++++++++++++++++++ 2 files changed, 464 insertions(+), 34 deletions(-) create mode 100644 tests/store/test_restore.py diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 117f73dc..a22aa773 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -2,7 +2,7 @@ import asyncio import json import logging from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from enum import Enum from importlib import metadata from pathlib import Path @@ -192,6 +192,24 @@ REQUIRED_TABLES: tuple[str, ...] = ( # version; guards against timestamp precision at the boundary. TAG_RETENTION_MARGIN = timedelta(seconds=1) +# Restore order for multi-table restore and its rollback. documents restores +# last: writes land in it last on the ingest path, making it the closest +# available database commit point. +RESTORE_TABLE_ORDER: tuple[str, ...] = tuple( + name for name in REQUIRED_TABLES if name != "documents" +) + ("documents",) + + +def _safety_tag_name(existing: set[str]) -> str: + """Collision-resistant name for the pre-restore safety tag.""" + base = f"before-restore-{datetime.now(UTC).strftime('%Y%m%dT%H%M%SZ')}" + if base not in existing: + return base + n = 2 + while f"{base}-{n}" in existing: + n += 1 + return f"{base}-{n}" + @dataclass class TagInfo: @@ -883,42 +901,47 @@ class Store: """ self._assert_writable() self._assert_not_rebuilding() - tables = self._tables() async with self._rebuild_lock, self._write_lock: - existing = [ - table_name - for table_name, table in tables.items() - if name in await table.tags.list() - ] - if len(existing) == len(tables): - raise ValueError(f"Tag '{name}' already exists") - if existing: - raise ValueError( - f"Tag '{name}' already exists on some tables " - f"({', '.join(existing)}); delete it first with delete_tag" - ) + await self._create_tag_locked(name) - versions = await self.current_table_versions() - created: list[str] = [] - try: - for table_name, table in tables.items(): - await table.tags.create(name, versions[table_name]) - created.append(table_name) - except Exception as exc: - failed_cleanup: list[str] = [] - for table_name in created: - try: - await tables[table_name].tags.delete(name) - except Exception: - failed_cleanup.append(table_name) - if failed_cleanup: - raise RuntimeError( - f"Tag '{name}' creation failed ({exc}) and cleanup " - f"failed on: {', '.join(failed_cleanup)}. A partial " - "tag may remain; delete it with delete_tag." - ) from exc - raise + async def _create_tag_locked(self, name: str) -> None: + """Create a tag on every table; the caller must hold the write lock.""" + tables = self._tables() + + existing = [ + table_name + for table_name, table in tables.items() + if name in await table.tags.list() + ] + if len(existing) == len(tables): + raise ValueError(f"Tag '{name}' already exists") + if existing: + raise ValueError( + f"Tag '{name}' already exists on some tables " + f"({', '.join(existing)}); delete it first with delete_tag" + ) + + versions = await self.current_table_versions() + created: list[str] = [] + try: + for table_name, table in tables.items(): + await table.tags.create(name, versions[table_name]) + created.append(table_name) + except Exception as exc: + failed_cleanup: list[str] = [] + for table_name in created: + try: + await tables[table_name].tags.delete(name) + except Exception: + failed_cleanup.append(table_name) + if failed_cleanup: + raise RuntimeError( + f"Tag '{name}' creation failed ({exc}) and cleanup " + f"failed on: {', '.join(failed_cleanup)}. A partial " + "tag may remain; delete it with delete_tag." + ) from exc + raise async def list_tags(self) -> dict[str, TagInfo]: """Aggregate per-table tags into database-level tags. @@ -968,6 +991,109 @@ class Store: "Remnants remain; retry delete_tag." ) + async def _restore_tables( + self, versions: dict[str, int], *, best_effort: bool = False + ) -> list[tuple[str, Exception]]: + """Restore every table to the given versions, documents last. + + Stops at the first failure by default; with best_effort, continues + through all tables. Returns the failures either way. + """ + tables = self._tables() + failures: list[tuple[str, Exception]] = [] + for table_name in RESTORE_TABLE_ORDER: + try: + await tables[table_name].restore(int(versions[table_name])) + except Exception as exc: + failures.append((table_name, exc)) + if not best_effort: + break + return failures + + async def restore_tag(self, name: str) -> str: + """Restore every table to the versions of a complete tag. + + Creates a complete safety tag for the pre-restore state before + changing any table and returns its name. Each table restore writes a + new latest version; nothing is left checked out read-only. + + In-process coordination only: all other writers must be stopped for + the duration of the operation. + + Raises: + ReadOnlyError: If the store is in read-only mode. + ValueError: If a rebuild is in progress, the tag does not exist, + or the tag is partial. + RuntimeError: If the safety tag could not be created (no table + changed), or a table restore failed (the error states whether + rollback succeeded). + """ + self._assert_writable() + self._assert_not_rebuilding() + + async with self._rebuild_lock, self._write_lock: + tags = await self.list_tags() + info = tags.get(name) + if info is None: + raise ValueError(f"Tag '{name}' does not exist") + if not info.complete: + raise ValueError( + f"Tag '{name}' is partial (missing tables: " + f"{', '.join(info.missing_tables)}) and cannot be " + "restored; delete it with delete_tag" + ) + + snapshot = await self.current_table_versions() + safety_tag = _safety_tag_name(set(tags)) + try: + await self._create_tag_locked(safety_tag) + except Exception as exc: + raise RuntimeError( + f"Restore of tag '{name}' did not begin: safety tag " + f"creation failed ({exc}). No table was changed." + ) from exc + + try: + failures = await self._restore_tables(info.tables) + except asyncio.CancelledError: + # CancelledError is a BaseException and escapes the + # per-table handler; roll back before re-raising, shielded + # from further cancellation. + rollback_failures = await asyncio.shield( + self._restore_tables(snapshot, best_effort=True) + ) + if rollback_failures: + failed_names = ", ".join(t for t, _ in rollback_failures) + raise RuntimeError( + f"Restore of tag '{name}' was cancelled and rollback " + f"failed on: {failed_names}. The database may be " + f"cross-table inconsistent; manual recovery is " + f"required using safety tag '{safety_tag}'." + ) from None + raise + if failures: + failed_table, cause = failures[0] + rollback_failures = await self._restore_tables( + snapshot, best_effort=True + ) + if rollback_failures: + failed_names = ", ".join(t for t, _ in rollback_failures) + raise RuntimeError( + f"Restore of tag '{name}' failed on table " + f"'{failed_table}' and rollback failed on: " + f"{failed_names}. The database may be cross-table " + f"inconsistent; manual recovery is required using " + f"safety tag '{safety_tag}'." + ) from cause + raise RuntimeError( + f"Restore of tag '{name}' failed on table " + f"'{failed_table}'; all tables were rolled back to the " + f"pre-restore state. Safety tag '{safety_tag}' is " + "preserved." + ) from cause + + return safety_tag + async def list_table_versions(self, table_name: str) -> list[dict[str, Any]]: """List version history for a table. diff --git a/tests/store/test_restore.py b/tests/store/test_restore.py new file mode 100644 index 00000000..337b133e --- /dev/null +++ b/tests/store/test_restore.py @@ -0,0 +1,304 @@ +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")