danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Protect restore rollback from cancellation; report delete_tag listing failures

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-07-16 10:15:26 +03:00
parent ce1b9ac88e
commit 1f88944ada
No known key found for this signature in database
6 changed files with 132 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -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

2
docs/cli.md
View file

@ -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

12
docs/python.md
View file

@ -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

60
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -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 "

45
tests/store/test_restore.py
View file

@ -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"}

33
tests/store/test_tags.py
View file

@ -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() == {}