Add Store.restore_tag
This commit is contained in:
Yiorgis Gozadinos
parent
a1f3435df3
commit
f41ec379df
2 changed files with 464 additions and 34 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
304
tests/store/test_restore.py
Normal file
304
tests/store/test_restore.py
Normal file
|
|
@ -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")
|
||||
Loading…
Reference in a new issue