From 1f88944ada5d0ac056eee40a9f6b8ba09f20b1ac Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 16 Jul 2026 10:15:26 +0300 Subject: [PATCH] Protect restore rollback from cancellation; report delete_tag listing failures --- CHANGELOG.md | 2 +- docs/cli.md | 2 + docs/python.md | 12 +++-- haiku_rag_slim/haiku/rag/store/engine.py | 60 +++++++++++++++++------- tests/store/test_restore.py | 45 ++++++++++++++++++ tests/store/test_tags.py | 33 +++++++++++++ 6 files changed, 132 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c575839..689fa53b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Removed -- `--before` global flag. There is no read-only replacement; create tags prospectively before important changes and use `tag restore` during a maintenance window. +- `--before` global flag and the `before` constructor arguments on `HaikuRAG`, `Store`, `HaikuRAGApp`, `ChatApp`/`run_chat`, and `InspectorApp`/`run_inspector`. There is no read-only replacement; create tags prospectively before important changes and use `tag restore` during a maintenance window. ### Fixed diff --git a/docs/cli.md b/docs/cli.md index eb404168..ad7252a4 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -557,6 +557,8 @@ haiku-rag tag delete release-1 A tag present on every table is complete. A tag missing from some tables (created outside haiku.rag, or left behind by a failure) is partial. `tag list` marks partial tags. Partial tags can be listed and deleted but never restored. +Create tags with other writers stopped. Tag creation coordinates writers within one process only; a writer in another process can commit between the per-table snapshot reads, and the tag then captures a mixed state. + Tagged versions survive `vacuum`. Vacuum retains the oldest tagged version and every newer version; versions older than the oldest tag remain eligible for cleanup. Delete tags you no longer need so cleanup can advance. ### Restore diff --git a/docs/python.md b/docs/python.md index d4c2fcb2..0e377024 100644 --- a/docs/python.md +++ b/docs/python.md @@ -436,7 +436,7 @@ This compacts tables and removes historical versions to keep disk usage in check ### Tags -Tag the current database state and restore it later, for example after an ingestion run. A tag covers all five tables and is created from a single version snapshot: +Tag the current database state and restore it later, for example after an ingestion run. A tag covers all five tables and is created from a single version snapshot. Create tags with other writers stopped: the snapshot is coordinated within one process only, and a writer in another process can commit between the per-table reads. ```python await client.store.create_tag("release-1") @@ -444,8 +444,6 @@ await client.store.create_tag("release-1") tags = await client.store.list_tags() for name, info in tags.items(): print(name, info.tables, info.complete) - -await client.store.delete_tag("release-1") ``` `restore_tag` brings the live database back to a tagged state. It creates a complete safety tag for the current state before changing any table and returns its name: @@ -454,7 +452,13 @@ await client.store.delete_tag("release-1") safety_tag = await client.store.restore_tag("release-1") ``` -Restore is a maintenance operation: stop all other writers first. A tag present on only some tables is partial; `list_tags` reports it via `missing_tables`, and partial tags can be deleted but never restored. Vacuum retains the oldest tagged version and everything newer, so delete tags you no longer need. +Restore is a maintenance operation: stop all other writers first. A tag present on only some tables is partial; `list_tags` reports it via `missing_tables`, and partial tags can be deleted but never restored. + +Delete tags you no longer need. Vacuum retains the oldest tagged version and everything newer: + +```python +await client.store.delete_tag("release-1") +``` ### Rebuilding the Database diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index a22aa773..0b53fc1f 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -569,8 +569,8 @@ class Store: Lance hard-errors when a tagged version falls inside the cleanup window and the Python API exposes no way to skip tagged versions, so - everything older than the oldest tag is retained until that tag is - deleted. + the oldest tagged version and everything newer are retained; versions + older than the oldest tag remain eligible for cleanup. """ tags = await table.tags.list() if not tags: @@ -892,6 +892,9 @@ class Store: Serializes with client writes via the write lock so a write cannot land between the version snapshot and the per-table tag creation. + This is in-process coordination only: a writer in another process + can commit between the per-table version reads, so create tags with + all other writers stopped when a consistent snapshot matters. Raises: ReadOnlyError: If the store is in read-only mode. @@ -977,19 +980,21 @@ class Store: found = False failed: list[str] = [] for table_name, table in self._tables().items(): - if name in await table.tags.list(): - found = True - try: + try: + if name in await table.tags.list(): + found = True await table.tags.delete(name) - except Exception: - failed.append(table_name) - if not found: - raise ValueError(f"Tag '{name}' does not exist") + except Exception: + failed.append(table_name) if failed: + # A listing failure obscures whether the tag exists on that + # table, so failures take precedence over not-found. raise RuntimeError( f"Tag '{name}' deletion failed on: {', '.join(failed)}. " - "Remnants remain; retry delete_tag." + "Remnants may remain; retry delete_tag." ) + if not found: + raise ValueError(f"Tag '{name}' does not exist") async def _restore_tables( self, versions: dict[str, int], *, best_effort: bool = False @@ -1010,6 +1015,28 @@ class Store: break return failures + async def _rollback_to_snapshot( + self, snapshot: dict[str, int] + ) -> tuple[list[tuple[str, Exception]], bool]: + """Best-effort rollback that a cancellation cannot interrupt. + + Runs the rollback as a task and keeps waiting for it even if this + coroutine is cancelled, so a Ctrl-C cannot leave the rollback half + applied. Returns the rollback failures and whether a cancellation + was absorbed; the caller must re-deliver an absorbed cancellation. + """ + task = asyncio.ensure_future(self._restore_tables(snapshot, best_effort=True)) + cancelled = False + while True: + try: + return await asyncio.shield(task), cancelled + except asyncio.CancelledError: + if task.done(): + # The rollback coroutine itself ended cancelled; there is + # nothing left to wait for. + raise + cancelled = True + async def restore_tag(self, name: str) -> str: """Restore every table to the versions of a complete tag. @@ -1057,11 +1084,8 @@ class Store: 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) - ) + # per-table handler; roll back before re-raising. + rollback_failures, _ = await self._rollback_to_snapshot(snapshot) if rollback_failures: failed_names = ", ".join(t for t, _ in rollback_failures) raise RuntimeError( @@ -1073,8 +1097,8 @@ class Store: raise if failures: failed_table, cause = failures[0] - rollback_failures = await self._restore_tables( - snapshot, best_effort=True + rollback_failures, cancelled = await self._rollback_to_snapshot( + snapshot ) if rollback_failures: failed_names = ", ".join(t for t, _ in rollback_failures) @@ -1085,6 +1109,8 @@ class Store: f"inconsistent; manual recovery is required using " f"safety tag '{safety_tag}'." ) from cause + if cancelled: + raise asyncio.CancelledError() raise RuntimeError( f"Restore of tag '{name}' failed on table " f"'{failed_table}'; all tables were rolled back to the " diff --git a/tests/store/test_restore.py b/tests/store/test_restore.py index 3c7731ba..9f71e6b5 100644 --- a/tests/store/test_restore.py +++ b/tests/store/test_restore.py @@ -336,3 +336,48 @@ async def test_restore_old_version_marker_requires_explicit_migration(temp_db_pa assert await _doc_contents(store) == {"First document"} await store.restore_tag(safety_tag) assert await _doc_contents(store) == {"First document", "Second document"} + + +@pytest.mark.asyncio +async def test_restore_failure_rollback_survives_cancellation( + temp_db_path, monkeypatch +): + """Cancelling restore while it rolls back a failed restore must not + interrupt the rollback: all tables return to the snapshot before the + cancellation is delivered.""" + 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} + rollback_started = asyncio.Event() + release = asyncio.Event() + + async def flaky_restore(self, version=None): + calls["n"] += 1 + if calls["n"] == 3: + raise RuntimeError("restore boom") + if calls["n"] == 4: + rollback_started.set() + await release.wait() + return await real_restore(self, version) + + monkeypatch.setattr(AsyncTable, "restore", flaky_restore) + + task = asyncio.create_task(store.restore_tag("release-1")) + await rollback_started.wait() + task.cancel() + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + monkeypatch.undo() + # 3 forward calls (2 ok, 1 failed) + all 5 rollback calls ran. + assert calls["n"] == 8 + assert await _doc_contents(store) == {"First document", "Second document"} diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py index b007cac3..74288ada 100644 --- a/tests/store/test_tags.py +++ b/tests/store/test_tags.py @@ -412,3 +412,36 @@ async def test_list_table_versions_returns_history(temp_db_path): for v in versions: assert "version" in v assert "timestamp" in v + + +@pytest.mark.asyncio +async def test_delete_tag_reports_listing_failures(temp_db_path, monkeypatch): + """A tags.list() failure mid-delete is reported with the table named and + a recovery hint, instead of escaping raw after earlier deletions.""" + async with Store(temp_db_path, create=True) as store: + await store.create_tag("release-1") + + real_list = AsyncTags.list + calls = {"n": 0} + + async def flaky_list(self): + calls["n"] += 1 + if calls["n"] == 2: + raise RuntimeError("list boom") + return await real_list(self) + + monkeypatch.setattr(AsyncTags, "list", flaky_list) + + with pytest.raises(RuntimeError) as exc_info: + await store.delete_tag("release-1") + + msg = str(exc_info.value) + assert "document_meta" in msg + assert "retry delete_tag" in msg + + monkeypatch.undo() + tags = await store.list_tags() + assert set(tags["release-1"].tables) == {"document_meta"} + + await store.delete_tag("release-1") + assert await store.list_tags() == {}