fix(repair): require clean SQLite recovery finalization

Co-authored-by: Aaron Zendejas <phileokairos@gmail.com>
This commit is contained in:
Igor Lins e Silva 2026-07-14 19:12:56 -03:00
parent 32e545eaaa
commit acd62442be
4 changed files with 175 additions and 11 deletions

View File

@ -1132,7 +1132,7 @@ def cmd_repair(args):
if getattr(args, "mode", "legacy") == "from-sqlite":
from .migrate import confirm_destructive_action
from .repair import RebuildPartialError, rebuild_from_sqlite
from .repair import RebuildCleanupError, RebuildPartialError, rebuild_from_sqlite
source_path = getattr(args, "source", None)
source_path = (
@ -1170,6 +1170,13 @@ def cmd_repair(args):
f"Failed in collection: {exc.failed_collection}"
)
sys.exit(1)
except RebuildCleanupError:
# All rows may have landed, but rebuild_from_sqlite deliberately
# withholds success until FTS5 rebuild, VACUUM, and quick_check are
# clean. Its exception already includes the retained destination
# and archive/source recovery paths.
print("\n Rebuild cleanup failed — see recovery details above.")
sys.exit(1)
# An empty counts dict is rebuild_from_sqlite's documented signal
# for a validation refusal (missing source, existing dest,
# in-place without --archive-existing). The library already

View File

@ -877,7 +877,12 @@ class _DefaultProgress:
return f" (elapsed {_format_eta(elapsed)}, rate {rate:.1f}/s, ETA {_format_eta(eta)})"
def _vacuum_and_rebuild_fts5(palace_path: str, progress=print) -> None:
def _vacuum_and_rebuild_fts5(
palace_path: str,
progress=print,
*,
strict: bool = False,
) -> None:
"""VACUUM the palace SQLite file and rebuild the FTS5 index if present.
Repeated ``repair --yes`` runs delete and recreate the drawers collection,
@ -886,12 +891,16 @@ def _vacuum_and_rebuild_fts5(palace_path: str, progress=print) -> None:
internally inconsistent after multiple collection deletes; the rebuild
command fixes it atomically without touching any row data.
Failures are non-fatal: a warning is printed and the caller continues.
The repair itself succeeded at this point VACUUM/FTS5 are best-effort
cleanup, not correctness requirements.
Existing repair paths use the default best-effort behavior: failures log a
warning and return because their primary rebuild has already succeeded.
SQLite recovery passes ``strict=True`` because its bulk upserts can leave
this derived index malformed; that path must not report success until the
rebuild, VACUUM, and a final quick_check all complete.
"""
sqlite_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.exists(sqlite_path):
if strict:
raise FileNotFoundError(f"recovered palace has no SQLite database: {sqlite_path}")
return
try:
with closing(sqlite3.connect(sqlite_path, isolation_level=None)) as conn:
@ -907,7 +916,18 @@ def _vacuum_and_rebuild_fts5(palace_path: str, progress=print) -> None:
progress(" FTS5 index rebuilt.")
conn.execute("VACUUM")
progress(" SQLite VACUUM complete.")
if strict:
rows = conn.execute("PRAGMA quick_check").fetchall()
errors = [str(row[0]) for row in rows if row and str(row[0]).lower() != "ok"]
if errors:
raise sqlite3.DatabaseError(
"post-recovery quick_check failed: " + "; ".join(errors[:3])
)
progress(" SQLite quick_check clean.")
except Exception as exc:
if strict:
progress(f" ERROR: required post-recovery cleanup failed: {exc}")
raise RuntimeError(f"required post-recovery cleanup failed: {exc}") from exc
progress(f" Warning: post-repair cleanup failed (non-fatal): {exc}")
@ -1093,6 +1113,29 @@ class RebuildPartialError(Exception):
self.archive_path = archive_path
class RebuildCleanupError(Exception):
"""Raised when all recoverable rows landed but final cleanup failed.
The destination is intentionally retained for inspection, and an in-place
rebuild's original archive remains untouched. Callers must not treat this
as success because the derived FTS5 index has not been verified clean.
"""
def __init__(
self,
message: str,
*,
counts: dict[str, int],
dest_palace: str,
archive_path: Optional[str],
):
super().__init__(message)
self.message = message
self.counts = counts
self.dest_palace = dest_palace
self.archive_path = archive_path
def _rebuild_one_collection(
*,
backend: ChromaBackend,
@ -1342,7 +1385,9 @@ def rebuild_from_sqlite(
chromadb upsert fails partway through; the dest palace is left in
place so the user can inspect what landed, and the in-place archive
(when applicable) is reported in the error so the user can re-run
against it.
against it. Raises :class:`RebuildCleanupError` if all rows land but the
required FTS5 rebuild, VACUUM, or final quick_check fails; this prevents a
structurally unverified recovery from being reported as complete.
.. warning::
@ -1491,14 +1536,39 @@ def rebuild_from_sqlite(
else:
print(f" done: {upserted} rows in {cname}")
print(f"\n Rebuild complete. {sum(counts.values())} total rows.")
if archive_path is not None:
print(f" Original palace archived at: {archive_path}")
print(f"{'=' * 55}\n")
return counts
finally:
backend.close()
# Bulk Chroma upserts can leave the derived FTS5 index internally
# inconsistent even when all source rows landed. Rebuild it only after
# the backend releases its SQLite handle; otherwise VACUUM cannot obtain
# the exclusive lock it needs on Windows.
try:
_vacuum_and_rebuild_fts5(dest_palace, strict=True)
except Exception as exc:
message_parts = [
f"Post-recovery cleanup failed after {sum(counts.values())} rows were rebuilt: {exc}",
f"Recovered palace retained at: {dest_palace}",
]
if archive_path is not None:
message_parts.append(f"Original palace remains archived at: {archive_path}")
else:
message_parts.append(f"Source palace is unchanged at: {source_palace}")
message = "\n ".join(message_parts)
print(f"\n ERROR: {message}")
raise RebuildCleanupError(
message,
counts=dict(counts),
dest_palace=dest_palace,
archive_path=archive_path,
) from exc
print(f"\n Rebuild complete. {sum(counts.values())} total rows.")
if archive_path is not None:
print(f" Original palace archived at: {archive_path}")
print(f"{'=' * 55}\n")
return counts
def status(palace_path=None, collection_name: Optional[str] = None) -> dict:
"""Read-only health check: compare sqlite vs HNSW element counts.

View File

@ -1779,6 +1779,34 @@ def test_cmd_repair_from_sqlite_success_does_not_exit(mock_config_cls, tmp_path)
cmd_repair(args)
@patch("mempalace.cli.MempalaceConfig")
def test_cmd_repair_from_sqlite_cleanup_failure_exits_nonzero(mock_config_cls, tmp_path, capsys):
from mempalace.repair import RebuildCleanupError
palace_dir = tmp_path / "palace"
source_dir = tmp_path / "source"
mock_config_cls.return_value.palace_path = str(palace_dir)
args = argparse.Namespace(
palace=str(palace_dir),
mode="from-sqlite",
source=str(source_dir),
archive_existing=False,
yes=True,
)
failure = RebuildCleanupError(
"cleanup failed",
counts={"mempalace_drawers": 1},
dest_palace=str(palace_dir),
archive_path=None,
)
with patch("mempalace.repair.rebuild_from_sqlite", side_effect=failure):
with pytest.raises(SystemExit) as excinfo:
cmd_repair(args)
assert excinfo.value.code == 1
assert "Rebuild cleanup failed" in capsys.readouterr().out
@patch("mempalace.cli.MempalaceConfig")
def test_cmd_repair_rebuild_index_alias_uses_sqlite_archive(mock_config_cls, tmp_path):
"""``repair rebuild-index`` must bypass Chroma reads and rebuild from SQLite."""

View File

@ -1752,6 +1752,60 @@ def test_rebuild_from_sqlite_roundtrips_via_real_chromadb(tmp_path):
assert closet_row["metadatas"][0] == {"wing": "alpha"}
def test_rebuild_from_sqlite_rebuilds_fts5_after_chroma_closes(tmp_path, monkeypatch):
"""The SQLite recovery path must finish by rebuilding Chroma's FTS5 index.
Large bulk upserts can leave the derived full-text index malformed even
when every source drawer survived. The repair is not complete until the
Chroma client releases its SQLite handle and FTS5 is rebuilt.
"""
source = tmp_path / "source"
dest = tmp_path / "dest"
_seed_palace(source, "mempalace_drawers", [("d1", "doc", {"wing": "w"})])
calls = []
real_rebuild = repair._vacuum_and_rebuild_fts5
def _spy(path, progress=print, *, strict=False):
calls.append((path, strict))
return real_rebuild(path, progress=progress, strict=strict)
monkeypatch.setattr(repair, "_vacuum_and_rebuild_fts5", _spy)
counts = repair.rebuild_from_sqlite(str(source), str(dest))
assert counts["mempalace_drawers"] == 1
assert calls == [(str(dest), True)]
def test_rebuild_from_sqlite_cleanup_failure_is_not_reported_as_success(
tmp_path, monkeypatch, capsys
):
source = tmp_path / "source"
dest = tmp_path / "dest"
_seed_palace(source, "mempalace_drawers", [("d1", "verbatim", {"wing": "w"})])
def _fail_cleanup(path, progress=print, *, strict=False):
assert path == str(dest)
assert strict is True
raise RuntimeError("simulated FTS5 rebuild failure")
monkeypatch.setattr(repair, "_vacuum_and_rebuild_fts5", _fail_cleanup)
with pytest.raises(repair.RebuildCleanupError) as excinfo:
repair.rebuild_from_sqlite(str(source), str(dest))
exc = excinfo.value
assert exc.counts["mempalace_drawers"] == 1
assert exc.dest_palace == str(dest)
assert exc.archive_path is None
assert dest.exists()
assert (source / "chroma.sqlite3").exists()
output = capsys.readouterr().out
assert "Rebuild complete" not in output
assert "Post-recovery cleanup failed" in output
def test_rebuild_from_sqlite_refuses_existing_dest(tmp_path):
"""Refuse to write into a directory that already exists when source
and dest differ. Without this, an unattended re-run would silently
@ -2017,6 +2071,11 @@ def test_vacuum_and_rebuild_fts5_missing_sqlite(tmp_path):
repair._vacuum_and_rebuild_fts5(str(tmp_path)) # no file — must not raise
def test_vacuum_and_rebuild_fts5_strict_requires_sqlite(tmp_path):
with pytest.raises(FileNotFoundError, match="has no SQLite database"):
repair._vacuum_and_rebuild_fts5(str(tmp_path), strict=True)
# ── FTS5 inverted-index auto-heal (#1596) ─────────────────────────────