c3749daf48
Tests / test-linux (3.13) (push) Failing after 0s
Tests / test-linux (3.11) (push) Failing after 1s
Tests / lint (push) Failing after 0s
Tests / test-linux (3.9) (push) Failing after 1s
Docker / build (push) Failing after 1s
Docker / build-gpu (push) Failing after 2s
Tests / test-windows (push) Has been cancelled
Tests / test-macos (push) Has been cancelled
2197 lines
87 KiB
Python
2197 lines
87 KiB
Python
"""Tests for mempalace.repair — scan, prune, and rebuild HNSW index."""
|
|
|
|
import os
|
|
import sqlite3
|
|
from contextlib import closing
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
|
|
from _chroma_palace_helper import make_minimal_chroma_sqlite
|
|
|
|
from mempalace import repair
|
|
|
|
|
|
# ── _get_palace_path ──────────────────────────────────────────────────
|
|
|
|
|
|
@patch("mempalace.repair.MempalaceConfig", create=True)
|
|
def test_get_palace_path_from_config(mock_config_cls):
|
|
mock_config_cls.return_value.palace_path = "/configured/palace"
|
|
with patch.dict("sys.modules", {}):
|
|
# Force reimport to pick up the mock
|
|
result = repair._get_palace_path()
|
|
assert isinstance(result, str)
|
|
|
|
|
|
def test_get_palace_path_fallback():
|
|
with patch("mempalace.repair._get_palace_path") as mock_get:
|
|
mock_get.return_value = os.path.join(os.path.expanduser("~"), ".mempalace", "palace")
|
|
result = mock_get()
|
|
assert ".mempalace" in result
|
|
|
|
|
|
def test_get_collection_name_from_config():
|
|
from mempalace.config import get_configured_collection_name
|
|
|
|
get_configured_collection_name.cache_clear()
|
|
with patch("mempalace.config.MempalaceConfig") as mock_config_cls:
|
|
mock_config_cls.return_value.collection_name = "custom_drawers"
|
|
assert repair._drawers_collection_name() == "custom_drawers"
|
|
get_configured_collection_name.cache_clear()
|
|
|
|
|
|
# ── _paginate_ids ─────────────────────────────────────────────────────
|
|
|
|
|
|
def test_paginate_ids_single_batch():
|
|
col = MagicMock()
|
|
col.get.return_value = {"ids": ["id1", "id2", "id3"]}
|
|
ids = repair._paginate_ids(col)
|
|
assert ids == ["id1", "id2", "id3"]
|
|
|
|
|
|
def test_paginate_ids_empty():
|
|
col = MagicMock()
|
|
col.get.return_value = {"ids": []}
|
|
ids = repair._paginate_ids(col)
|
|
assert ids == []
|
|
|
|
|
|
def test_paginate_ids_with_where():
|
|
col = MagicMock()
|
|
col.get.return_value = {"ids": ["id1"]}
|
|
repair._paginate_ids(col, where={"wing": "test"})
|
|
col.get.assert_called_with(where={"wing": "test"}, include=[], limit=1000, offset=0)
|
|
|
|
|
|
def test_paginate_ids_offset_exception_fallback():
|
|
col = MagicMock()
|
|
# First call raises, fallback returns ids, second fallback returns empty
|
|
col.get.side_effect = [
|
|
Exception("offset bug"),
|
|
{"ids": ["id1", "id2"]},
|
|
Exception("offset bug"),
|
|
{"ids": ["id1", "id2"]}, # same ids = no new = break
|
|
]
|
|
ids = repair._paginate_ids(col)
|
|
assert "id1" in ids
|
|
|
|
|
|
# ── _extract_drawers ──────────────────────────────────────────────────
|
|
|
|
|
|
def test_extract_drawers_preserves_valid_metadata():
|
|
"""Non-empty dict metadata passes through unchanged."""
|
|
col = MagicMock()
|
|
col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a", "room": "1"}, {"wing": "b", "room": "2"}],
|
|
}
|
|
all_ids, all_docs, all_metas = repair._extract_drawers(col, total=2, batch_size=2)
|
|
assert all_ids == ["id1", "id2"]
|
|
assert all_docs == ["doc1", "doc2"]
|
|
assert all_metas == [{"wing": "a", "room": "1"}, {"wing": "b", "room": "2"}]
|
|
|
|
|
|
def test_extract_drawers_sanitizes_none_metadata():
|
|
"""None entries in metadatas are coerced to the sentinel dict.
|
|
|
|
chromadb 1.5.x's `validate_metadata` raises `ValueError: Expected metadata
|
|
to be a non-empty dict, got 0 metadata attributes in add.` if it sees a
|
|
None entry; the sanitizer keeps the rebuild upsert from crashing.
|
|
"""
|
|
col = MagicMock()
|
|
col.get.return_value = {
|
|
"ids": ["id1", "id2", "id3"],
|
|
"documents": ["doc1", "doc2", "doc3"],
|
|
"metadatas": [{"wing": "a"}, None, {"wing": "c"}],
|
|
}
|
|
_, _, all_metas = repair._extract_drawers(col, total=3, batch_size=3)
|
|
assert all_metas[0] == {"wing": "a"}
|
|
assert all_metas[1] == {"_repaired_empty_meta": True}
|
|
assert all_metas[2] == {"wing": "c"}
|
|
|
|
|
|
def test_extract_drawers_sanitizes_empty_dict_metadata():
|
|
"""Empty dict {} entries are coerced to the sentinel dict.
|
|
|
|
chromadb 1.5.x rejects `{}` the same way it rejects `None`. The comment
|
|
in the previous code path mistakenly assumed otherwise.
|
|
"""
|
|
col = MagicMock()
|
|
col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{}, {"wing": "b"}],
|
|
}
|
|
_, _, all_metas = repair._extract_drawers(col, total=2, batch_size=2)
|
|
assert all_metas[0] == {"_repaired_empty_meta": True}
|
|
assert all_metas[1] == {"wing": "b"}
|
|
|
|
|
|
def test_extract_drawers_sanitization_preserves_alignment():
|
|
"""Sanitized output keeps the same length and ordering as input.
|
|
|
|
Critical invariant: ids[i] / documents[i] / metadatas[i] must stay in
|
|
lockstep through the sanitizer; otherwise the rebuild upsert mis-pairs
|
|
documents with metadata.
|
|
"""
|
|
col = MagicMock()
|
|
col.get.return_value = {
|
|
"ids": ["id1", "id2", "id3", "id4"],
|
|
"documents": ["d1", "d2", "d3", "d4"],
|
|
"metadatas": [None, {"k": "v"}, {}, None],
|
|
}
|
|
all_ids, all_docs, all_metas = repair._extract_drawers(col, total=4, batch_size=4)
|
|
assert len(all_ids) == len(all_docs) == len(all_metas) == 4
|
|
assert all_ids == ["id1", "id2", "id3", "id4"]
|
|
assert all_metas[0] == {"_repaired_empty_meta": True}
|
|
assert all_metas[1] == {"k": "v"}
|
|
assert all_metas[2] == {"_repaired_empty_meta": True}
|
|
assert all_metas[3] == {"_repaired_empty_meta": True}
|
|
|
|
|
|
def test_extract_drawers_multiple_batches():
|
|
"""Pagination handles batch boundaries without losing/duplicating rows."""
|
|
col = MagicMock()
|
|
col.get.side_effect = [
|
|
{"ids": ["id1", "id2"], "documents": ["d1", "d2"], "metadatas": [{"a": 1}, None]},
|
|
{"ids": ["id3"], "documents": ["d3"], "metadatas": [{}]},
|
|
{"ids": [], "documents": [], "metadatas": []},
|
|
]
|
|
all_ids, all_docs, all_metas = repair._extract_drawers(col, total=3, batch_size=2)
|
|
assert all_ids == ["id1", "id2", "id3"]
|
|
assert all_metas == [{"a": 1}, {"_repaired_empty_meta": True}, {"_repaired_empty_meta": True}]
|
|
|
|
|
|
# ── scan_palace ───────────────────────────────────────────────────────
|
|
|
|
|
|
def _install_mock_backend(mock_backend_cls, collection):
|
|
"""Wire mock_backend_cls so ChromaBackend().get_collection(...) returns *collection*."""
|
|
mock_backend = MagicMock()
|
|
mock_backend.get_collection.return_value = collection
|
|
mock_backend_cls.return_value = mock_backend
|
|
return mock_backend
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_scan_palace_no_ids(mock_backend_cls, tmp_path):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 0
|
|
mock_col.get.return_value = {"ids": []}
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
|
assert good == set()
|
|
assert bad == set()
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_scan_palace_all_good(mock_backend_cls, tmp_path):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
# _paginate_ids call
|
|
mock_col.get.side_effect = [
|
|
{"ids": ["id1", "id2"]}, # paginate
|
|
{"ids": ["id1", "id2"]}, # probe batch — both returned
|
|
]
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
|
assert "id1" in good
|
|
assert "id2" in good
|
|
assert len(bad) == 0
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_scan_palace_with_bad_ids(mock_backend_cls, tmp_path):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
|
|
def get_side_effect(**kwargs):
|
|
ids = kwargs.get("ids", None)
|
|
if ids is None:
|
|
# paginate call
|
|
return {"ids": ["good1", "bad1"]}
|
|
if "bad1" in ids and len(ids) == 1:
|
|
raise Exception("corrupt")
|
|
if "good1" in ids and len(ids) == 1:
|
|
return {"ids": ["good1"]}
|
|
# batch probe — raise to force per-id
|
|
raise Exception("batch fail")
|
|
|
|
mock_col.get.side_effect = get_side_effect
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
|
assert "good1" in good
|
|
assert "bad1" in bad
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_scan_palace_with_wing_filter(mock_backend_cls, tmp_path):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 1
|
|
mock_col.get.side_effect = [
|
|
{"ids": ["id1"]}, # paginate
|
|
{"ids": ["id1"]}, # probe
|
|
]
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
repair.scan_palace(palace_path=str(tmp_path), only_wing="test_wing")
|
|
# Verify where filter was passed
|
|
first_call = mock_col.get.call_args_list[0]
|
|
assert first_call.kwargs.get("where") == {"wing": "test_wing"}
|
|
|
|
|
|
# ── prune_corrupt ─────────────────────────────────────────────────────
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_prune_corrupt_no_file(mock_backend_cls, tmp_path):
|
|
# Should print message and return without error
|
|
repair.prune_corrupt(palace_path=str(tmp_path))
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_prune_corrupt_dry_run(mock_backend_cls, tmp_path):
|
|
bad_file = tmp_path / "corrupt_ids.txt"
|
|
bad_file.write_text("bad1\nbad2\n")
|
|
repair.prune_corrupt(palace_path=str(tmp_path), confirm=False)
|
|
# No backend calls in dry run
|
|
mock_backend_cls.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_prune_corrupt_confirmed(mock_backend_cls, tmp_path):
|
|
bad_file = tmp_path / "corrupt_ids.txt"
|
|
bad_file.write_text("bad1\nbad2\n")
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.side_effect = [10, 8]
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
repair.prune_corrupt(palace_path=str(tmp_path), confirm=True)
|
|
mock_col.delete.assert_called_once()
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_prune_corrupt_delete_failure_fallback(mock_backend_cls, tmp_path):
|
|
bad_file = tmp_path / "corrupt_ids.txt"
|
|
bad_file.write_text("bad1\nbad2\n")
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.side_effect = [10, 8]
|
|
# Batch delete fails, per-id succeeds
|
|
mock_col.delete.side_effect = [Exception("batch fail"), None, None]
|
|
_install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
repair.prune_corrupt(palace_path=str(tmp_path), confirm=True)
|
|
assert mock_col.delete.call_count == 3 # 1 batch + 2 individual
|
|
|
|
|
|
# ── rebuild_index ─────────────────────────────────────────────────────
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_no_palace(mock_backend_cls, tmp_path):
|
|
nonexistent = str(tmp_path / "nope")
|
|
repair.rebuild_index(palace_path=nonexistent)
|
|
mock_backend_cls.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair.shutil")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_empty_palace(mock_backend_cls, mock_shutil, tmp_path):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 0
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
mock_backend.delete_collection.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_read_failure_points_to_from_sqlite(mock_backend_cls, tmp_path):
|
|
"""A chromadb HNSW compactor failure makes the first ``count()`` read
|
|
raise; rebuild_index cannot recover it, so it must direct the user to
|
|
``repair --mode from-sqlite`` (rows are intact in chroma.sqlite3) rather
|
|
than re-mining from source files, which drops MCP-added drawers (#1843)."""
|
|
sqlite3.connect(str(tmp_path / "chroma.sqlite3")).close()
|
|
mock_col = MagicMock()
|
|
mock_col.count.side_effect = Exception("Failed to apply logs to the hnsw segment writer")
|
|
mock_backend_cls.return_value.get_collection.return_value = mock_col
|
|
msgs: list[str] = []
|
|
repair.rebuild_index(palace_path=str(tmp_path), progress=msgs.append)
|
|
out = "\n".join(msgs)
|
|
assert "mempalace repair --mode from-sqlite --archive-existing" in out
|
|
assert "may need to be re-mined" not in out
|
|
|
|
|
|
def test_index_read_recovery_guidance_recommends_from_sqlite():
|
|
"""The shared guidance helper names the from-sqlite recovery command in
|
|
full and never tells the user the palace ``may need to be re-mined`` —
|
|
the harmful pre-#1843 advice that silently drops MCP-added drawers."""
|
|
msg = repair.index_read_recovery_guidance()
|
|
assert "mempalace repair --mode from-sqlite --archive-existing" in msg
|
|
assert "may need to be re-mined" not in msg
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_success(mock_backend_cls, mock_copy, tmp_path):
|
|
# Create a valid sqlite file so the repair preflight can run quick_check.
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with sqlite3.connect(sqlite_path) as conn:
|
|
conn.execute("CREATE TABLE dummy(id INTEGER PRIMARY KEY)")
|
|
conn.commit()
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 2
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
# Verify: backed up sqlite only, not copytree.
|
|
mock_copy.assert_called_once()
|
|
assert "chroma.sqlite3" in str(mock_copy.call_args)
|
|
|
|
# Verify: deleted and recreated (cosine is the backend default)
|
|
assert mock_backend.create_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
]
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
]
|
|
|
|
# Verify: used upsert not add
|
|
mock_temp_col.upsert.assert_called_once()
|
|
mock_new_col.upsert.assert_called_once()
|
|
mock_new_col.add.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_ignores_missing_temp_collection_at_start(
|
|
mock_backend_cls, mock_copy, tmp_path
|
|
):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
def _fake_copy2(src, dst, **_):
|
|
with open(dst, "w") as handle:
|
|
handle.write("backup")
|
|
|
|
mock_copy.side_effect = _fake_copy2
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 2
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
mock_backend.delete_collection.side_effect = [
|
|
ValueError("Collection [mempalace_drawers__repair_tmp] does not exist"),
|
|
None,
|
|
None,
|
|
]
|
|
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
assert mock_copy.call_count == 1
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
]
|
|
|
|
|
|
def test_delete_collection_if_exists_reraises_unexpected_value_error():
|
|
mock_backend = MagicMock()
|
|
mock_backend.delete_collection.side_effect = ValueError("invalid collection name")
|
|
|
|
with pytest.raises(ValueError, match="invalid collection name"):
|
|
repair._delete_collection_if_exists(mock_backend, "/palace", "bad/name")
|
|
|
|
|
|
@patch("mempalace.repair.shutil")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_error_reading(mock_backend_cls, mock_shutil, tmp_path):
|
|
mock_backend = MagicMock()
|
|
mock_backend.get_collection.side_effect = Exception("corrupt")
|
|
mock_backend_cls.return_value = mock_backend
|
|
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
mock_backend.delete_collection.assert_not_called()
|
|
|
|
|
|
# ── #1208 truncation safety ───────────────────────────────────────────
|
|
|
|
|
|
def test_check_extraction_safety_passes_when_counts_match(tmp_path):
|
|
"""SQLite reports same count as extracted → no exception."""
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=500):
|
|
repair.check_extraction_safety(str(tmp_path), 500)
|
|
|
|
|
|
def test_check_extraction_safety_uses_configured_collection(tmp_path):
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=500) as count:
|
|
repair.check_extraction_safety(str(tmp_path), 500, collection_name="custom_drawers")
|
|
count.assert_called_once_with(str(tmp_path), "custom_drawers")
|
|
|
|
|
|
def test_check_extraction_safety_default_uses_configured_collection(tmp_path):
|
|
with (
|
|
patch("mempalace.repair._drawers_collection_name", return_value="custom_drawers"),
|
|
patch("mempalace.repair.sqlite_drawer_count", return_value=500) as count,
|
|
):
|
|
repair.check_extraction_safety(str(tmp_path), 500)
|
|
count.assert_called_once_with(str(tmp_path), "custom_drawers")
|
|
|
|
|
|
def test_check_extraction_safety_passes_when_sqlite_unreadable_and_under_cap(tmp_path):
|
|
"""SQLite check fails (None) but extraction is well under the cap → safe."""
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=None):
|
|
repair.check_extraction_safety(str(tmp_path), 5_000)
|
|
|
|
|
|
def test_check_extraction_safety_aborts_when_sqlite_higher(tmp_path):
|
|
"""SQLite reports more than extracted — the user-reported #1208 case."""
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=67_580):
|
|
try:
|
|
repair.check_extraction_safety(str(tmp_path), 10_000)
|
|
except repair.TruncationDetected as e:
|
|
assert e.sqlite_count == 67_580
|
|
assert e.extracted == 10_000
|
|
assert "67,580" in e.message
|
|
assert "10,000" in e.message
|
|
assert "57,580" in e.message # the loss number
|
|
else:
|
|
raise AssertionError("expected TruncationDetected")
|
|
|
|
|
|
def test_check_extraction_safety_aborts_when_unreadable_and_at_cap(tmp_path):
|
|
"""SQLite unreadable but extraction == default get() cap → suspicious."""
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=None):
|
|
try:
|
|
repair.check_extraction_safety(str(tmp_path), repair.CHROMADB_DEFAULT_GET_LIMIT)
|
|
except repair.TruncationDetected as e:
|
|
assert e.sqlite_count is None
|
|
assert e.extracted == repair.CHROMADB_DEFAULT_GET_LIMIT
|
|
assert "10,000" in e.message
|
|
else:
|
|
raise AssertionError("expected TruncationDetected")
|
|
|
|
|
|
def test_check_extraction_safety_override_skips_check(tmp_path):
|
|
"""``confirm_truncation_ok=True`` short-circuits both signals."""
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=99_999):
|
|
# Would normally abort — override allows through
|
|
repair.check_extraction_safety(str(tmp_path), 10_000, confirm_truncation_ok=True)
|
|
|
|
|
|
def test_sqlite_drawer_count_returns_none_on_missing_file(tmp_path):
|
|
"""Palace dir exists but no chroma.sqlite3 → None, not crash."""
|
|
assert repair.sqlite_drawer_count(str(tmp_path)) is None
|
|
|
|
|
|
def test_sqlite_drawer_count_returns_none_on_unreadable_schema(tmp_path):
|
|
"""File exists but isn't a chromadb sqlite → None, not crash."""
|
|
sqlite_path = os.path.join(str(tmp_path), "chroma.sqlite3")
|
|
with open(sqlite_path, "wb") as f:
|
|
f.write(b"not a sqlite file at all")
|
|
assert repair.sqlite_drawer_count(str(tmp_path)) is None
|
|
|
|
|
|
@patch("mempalace.repair.shutil")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_default_uses_configured_collection(mock_backend_cls, mock_shutil, tmp_path):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
|
|
with (
|
|
patch("mempalace.repair._drawers_collection_name", return_value="custom_drawers"),
|
|
patch("mempalace.repair.sqlite_drawer_count", return_value=2) as count,
|
|
):
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
mock_backend.get_collection.assert_called_once_with(str(tmp_path), "custom_drawers")
|
|
count.assert_called_once_with(str(tmp_path), "custom_drawers")
|
|
assert mock_backend.create_collection.call_args_list == [
|
|
call(str(tmp_path), "custom_drawers__repair_tmp"),
|
|
call(str(tmp_path), "custom_drawers"),
|
|
]
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "custom_drawers__repair_tmp"),
|
|
call(str(tmp_path), "custom_drawers"),
|
|
call(str(tmp_path), "custom_drawers__repair_tmp"),
|
|
]
|
|
|
|
|
|
def test_status_returns_uninitialized_when_db_missing(tmp_path, capsys):
|
|
"""repair.status on a palace dir without chroma.sqlite3 returns a
|
|
structured status (no chromadb client opened, per the design that
|
|
repair-status must work even on corrupted palaces — #1498)."""
|
|
# tmp_path exists, no chroma.sqlite3
|
|
result = repair.status(palace_path=str(tmp_path))
|
|
|
|
assert result["status"] == "uninitialized"
|
|
assert "no chroma.sqlite3" in result["message"]
|
|
captured = capsys.readouterr()
|
|
assert "has no chroma.sqlite3 yet" in captured.out + captured.err
|
|
|
|
|
|
def test_status_returns_empty_when_db_present_no_drawers(tmp_path, capsys):
|
|
"""repair.status on a palace with chroma.sqlite3 but zero drawer rows
|
|
returns a structured 'empty' status, distinguishable from 'unknown' /
|
|
'uninitialized' (#1498). Mocks sqlite_drawer_count to assert the
|
|
return-shape contract; see the real-disk sibling below for the
|
|
no-chromadb-client invariant."""
|
|
make_minimal_chroma_sqlite(tmp_path)
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=0):
|
|
result = repair.status(palace_path=str(tmp_path))
|
|
|
|
assert result["status"] == "empty"
|
|
assert "no drawers yet" in result["message"]
|
|
captured = capsys.readouterr()
|
|
assert "initialized but empty" in captured.out + captured.err
|
|
|
|
|
|
def test_status_empty_palace_never_opens_chromadb_client(tmp_path):
|
|
"""Design invariant from #1498: repair.status on an initialized-but-empty
|
|
palace must NOT open a chromadb client. Opening would materialize HNSW
|
|
segment state files on disk, breaking the promise that repair-status is
|
|
safe to run on corrupted palaces.
|
|
|
|
Real-disk sibling of test_status_returns_empty_when_db_present_no_drawers:
|
|
bootstrap a real chroma.sqlite3 via PersistentClient (creates the DB
|
|
file but no collection), then assert repair.status returns 'empty' and
|
|
no chromadb segment artifacts appeared in the dir."""
|
|
import chromadb
|
|
|
|
chromadb.PersistentClient(path=str(tmp_path))
|
|
before = sorted(p.name for p in tmp_path.iterdir())
|
|
|
|
result = repair.status(palace_path=str(tmp_path))
|
|
|
|
after = sorted(p.name for p in tmp_path.iterdir())
|
|
assert result["status"] == "empty", result
|
|
# repair.status must not create new files; chromadb writes HNSW segment
|
|
# state and *.bin payloads on collection open — none of those should
|
|
# appear here.
|
|
assert before == after, f"repair.status mutated palace on disk: before={before} after={after}"
|
|
|
|
|
|
def test_status_falls_through_to_capacity_when_sqlite_count_unreadable(tmp_path):
|
|
"""When sqlite_drawer_count returns None (schema drift / locked file),
|
|
repair.status must fall through to hnsw_capacity_status instead of
|
|
short-circuiting on 'empty' (#1498)."""
|
|
make_minimal_chroma_sqlite(tmp_path)
|
|
with (
|
|
patch("mempalace.repair.sqlite_drawer_count", return_value=None),
|
|
patch("mempalace.repair.hnsw_capacity_status") as capacity_status,
|
|
):
|
|
capacity_status.side_effect = [
|
|
{
|
|
"sqlite_count": None,
|
|
"hnsw_count": None,
|
|
"divergence": None,
|
|
"diverged": False,
|
|
"status": "unknown",
|
|
"message": "",
|
|
},
|
|
{
|
|
"sqlite_count": None,
|
|
"hnsw_count": None,
|
|
"divergence": None,
|
|
"diverged": False,
|
|
"status": "unknown",
|
|
"message": "",
|
|
},
|
|
]
|
|
result = repair.status(palace_path=str(tmp_path))
|
|
|
|
# Did not short-circuit on 'empty': fell through to capacity check.
|
|
# The healthy/fall-through path returns {drawers, closets} dicts, no top-level "status" key.
|
|
assert "status" not in result or result["status"] != "empty"
|
|
assert "drawers" in result and "closets" in result
|
|
assert capacity_status.called
|
|
|
|
|
|
def test_status_default_uses_configured_drawer_collection(tmp_path):
|
|
# Provide the on-disk preconditions the stratified state helper (#1498)
|
|
# checks before reaching the capacity probe: chroma.sqlite3 file exists
|
|
# and sqlite_drawer_count returns a positive number (palace not empty).
|
|
make_minimal_chroma_sqlite(tmp_path)
|
|
with (
|
|
patch("mempalace.repair._drawers_collection_name", return_value="custom_drawers"),
|
|
patch("mempalace.repair.sqlite_drawer_count", return_value=1),
|
|
patch("mempalace.repair.hnsw_capacity_status") as capacity_status,
|
|
):
|
|
capacity_status.side_effect = [
|
|
{
|
|
"sqlite_count": 1,
|
|
"hnsw_count": 1,
|
|
"divergence": 0,
|
|
"diverged": False,
|
|
"status": "ok",
|
|
"message": "",
|
|
},
|
|
{
|
|
"sqlite_count": 0,
|
|
"hnsw_count": 0,
|
|
"divergence": 0,
|
|
"diverged": False,
|
|
"status": "ok",
|
|
"message": "",
|
|
},
|
|
]
|
|
repair.status(palace_path=str(tmp_path))
|
|
|
|
assert capacity_status.call_args_list[0].args == (str(tmp_path), "custom_drawers")
|
|
assert capacity_status.call_args_list[1].args == (str(tmp_path), "mempalace_closets")
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_aborts_on_truncation_signal(mock_backend_cls, mock_copy, tmp_path):
|
|
"""rebuild_index honors the safety guard: SQLite says 67k, get() returns
|
|
10k → no delete_collection, no upsert, no backup."""
|
|
mock_backend = MagicMock()
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 10_000
|
|
# Single page comes back with 10_000 ids
|
|
mock_col.get.side_effect = [
|
|
{
|
|
"ids": [f"id{i}" for i in range(10_000)],
|
|
"documents": ["x"] * 10_000,
|
|
"metadatas": [{}] * 10_000,
|
|
},
|
|
{"ids": [], "documents": [], "metadatas": []},
|
|
]
|
|
mock_backend.get_collection.return_value = mock_col
|
|
mock_backend_cls.return_value = mock_backend
|
|
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=67_580):
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
# Guard fired: nothing destructive happened
|
|
mock_backend.delete_collection.assert_not_called()
|
|
mock_backend.create_collection.assert_not_called()
|
|
mock_copy.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair.shutil")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_proceeds_with_override(mock_backend_cls, mock_shutil, tmp_path):
|
|
"""Override flag lets repair proceed even when the guard would fire."""
|
|
mock_backend = MagicMock()
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 10_000
|
|
mock_col.get.side_effect = [
|
|
{
|
|
"ids": [f"id{i}" for i in range(10_000)],
|
|
"documents": ["x"] * 10_000,
|
|
"metadatas": [{}] * 10_000,
|
|
},
|
|
{"ids": [], "documents": [], "metadatas": []},
|
|
]
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 10_000
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 10_000
|
|
mock_backend.get_collection.return_value = mock_col
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
mock_backend_cls.return_value = mock_backend
|
|
|
|
with patch("mempalace.repair.sqlite_drawer_count", return_value=67_580):
|
|
repair.rebuild_index(palace_path=str(tmp_path), confirm_truncation_ok=True)
|
|
|
|
assert mock_backend.delete_collection.call_count == 3
|
|
assert mock_backend.create_collection.call_count == 2
|
|
mock_temp_col.upsert.assert_called()
|
|
mock_new_col.upsert.assert_called()
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_stage_failure_leaves_live_collection_untouched(
|
|
mock_backend_cls, mock_copy, tmp_path
|
|
):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 1
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.return_value = mock_temp_col
|
|
|
|
with pytest.raises(repair.RebuildCollectionError) as excinfo:
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
assert excinfo.value.live_replaced is False
|
|
assert mock_copy.call_count == 1
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
]
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_live_failure_restores_backup(mock_backend_cls, mock_copy, tmp_path):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
def _fake_copy2(src, dst, **_):
|
|
with open(dst, "w") as handle:
|
|
handle.write("backup")
|
|
|
|
mock_copy.side_effect = _fake_copy2
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.upsert.side_effect = RuntimeError("live upsert failed")
|
|
active_backend = MagicMock()
|
|
active_backend.get_collection.return_value = mock_col
|
|
active_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
helper_backend = MagicMock()
|
|
mock_backend_cls.side_effect = [active_backend, helper_backend]
|
|
|
|
with pytest.raises(repair.RebuildCollectionError) as excinfo:
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
assert excinfo.value.live_replaced is True
|
|
assert mock_copy.call_count == 2
|
|
assert active_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
]
|
|
active_backend.close_palace.assert_called_once_with(str(tmp_path))
|
|
helper_backend.close_palace.assert_not_called()
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_live_delete_missing_still_restores_backup(
|
|
mock_backend_cls, mock_copy, tmp_path
|
|
):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
def _fake_copy2(src, dst, **_):
|
|
with open(dst, "w") as handle:
|
|
handle.write("backup")
|
|
|
|
mock_copy.side_effect = _fake_copy2
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, RuntimeError("create failed")]
|
|
mock_backend.delete_collection.side_effect = [
|
|
None,
|
|
None,
|
|
None,
|
|
repair.ChromaNotFoundError("missing"),
|
|
]
|
|
|
|
with pytest.raises(repair.RebuildCollectionError) as excinfo:
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
assert excinfo.value.live_replaced is True
|
|
assert mock_copy.call_count == 2
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
]
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_restore_failure_preserves_original_error(
|
|
mock_backend_cls, mock_copy, tmp_path, capsys
|
|
):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
def _copy_side_effect(src, dst, **_):
|
|
# The restore copy reads from the timestamped backup file.
|
|
if ".backup." in str(src):
|
|
raise PermissionError("locked sqlite")
|
|
with open(dst, "w") as handle:
|
|
handle.write("backup")
|
|
|
|
mock_copy.side_effect = _copy_side_effect
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.upsert.side_effect = RuntimeError("live upsert failed")
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
|
|
with pytest.raises(repair.RebuildCollectionError) as excinfo:
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
out = capsys.readouterr().out
|
|
assert "locked sqlite" in out
|
|
assert "Manual restore required" in out
|
|
assert "live upsert failed" in str(excinfo.value)
|
|
|
|
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_collection_via_temp_keeps_original_error_when_cleanup_fails(
|
|
mock_backend_cls,
|
|
):
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, RuntimeError("live build failed")]
|
|
mock_backend.delete_collection.side_effect = [
|
|
None,
|
|
None,
|
|
RuntimeError("cleanup failed"),
|
|
]
|
|
|
|
with pytest.raises(repair.RebuildCollectionError) as excinfo:
|
|
repair._rebuild_collection_via_temp(
|
|
mock_backend,
|
|
"/palace",
|
|
["id1", "id2"],
|
|
["doc1", "doc2"],
|
|
[{"wing": "a"}, {"wing": "b"}],
|
|
batch_size=5000,
|
|
progress=lambda *args, **kwargs: None,
|
|
)
|
|
|
|
assert "live build failed" in str(excinfo.value)
|
|
assert excinfo.value.live_replaced is True
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call("/palace", "mempalace_drawers__repair_tmp"),
|
|
call("/palace", "mempalace_drawers"),
|
|
call("/palace", "mempalace_drawers__repair_tmp"),
|
|
]
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_ignores_temp_cleanup_failure_after_success(
|
|
mock_backend_cls, mock_copy, tmp_path
|
|
):
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
sqlite3.connect(str(sqlite_path)).close()
|
|
|
|
def _fake_copy2(src, dst, **_):
|
|
with open(dst, "w") as handle:
|
|
handle.write("backup")
|
|
|
|
mock_copy.side_effect = _fake_copy2
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 2
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 2
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
mock_backend.delete_collection.side_effect = [
|
|
None,
|
|
None,
|
|
RuntimeError("cleanup failed"),
|
|
]
|
|
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
assert mock_copy.call_count == 1
|
|
assert mock_backend.delete_collection.call_args_list == [
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
call(str(tmp_path), "mempalace_drawers"),
|
|
call(str(tmp_path), "mempalace_drawers__repair_tmp"),
|
|
]
|
|
|
|
|
|
# ── repair_max_seq_id ─────────────────────────────────────────────────
|
|
|
|
|
|
# Realistic poisoned values from the 2026-04-20 incident — from the sysdb-10
|
|
# b'\x11\x11' + 6 ASCII digit format being misread as big-endian u64.
|
|
_POISON_VAL = 1_229_822_654_365_970_487
|
|
|
|
|
|
def _seed_poisoned_max_seq_id(
|
|
palace_path: str,
|
|
*,
|
|
drawers_meta_max: int = 502607,
|
|
closets_meta_max: int = 501418,
|
|
drawers_vec_poison: int = _POISON_VAL,
|
|
drawers_meta_poison: int = _POISON_VAL + 1,
|
|
closets_vec_poison: int = _POISON_VAL + 2,
|
|
closets_meta_poison: int = _POISON_VAL + 3,
|
|
):
|
|
"""Build a minimal palace with poisoned max_seq_id rows.
|
|
|
|
Returns a dict with segment UUIDs and the expected clean values.
|
|
"""
|
|
os.makedirs(palace_path, exist_ok=True)
|
|
db_path = os.path.join(palace_path, "chroma.sqlite3")
|
|
|
|
drawers_coll = "coll-drawers-0000-1111-2222-333344445555"
|
|
closets_coll = "coll-closets-0000-1111-2222-333344445555"
|
|
drawers_vec = "seg-drawers-vec-0000-1111-2222-333344445555"
|
|
drawers_meta = "seg-drawers-meta-0000-1111-2222-33334444555"
|
|
closets_vec = "seg-closets-vec-0000-1111-2222-333344445555"
|
|
closets_meta = "seg-closets-meta-0000-1111-2222-33334444555"
|
|
|
|
with closing(sqlite3.connect(db_path)) as conn:
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE segments(
|
|
id TEXT PRIMARY KEY, type TEXT, scope TEXT, collection TEXT
|
|
);
|
|
CREATE TABLE max_seq_id(segment_id TEXT PRIMARY KEY, seq_id);
|
|
CREATE TABLE embeddings(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
segment_id TEXT,
|
|
embedding_id TEXT,
|
|
seq_id
|
|
);
|
|
CREATE TABLE embeddings_queue(seq_id INTEGER PRIMARY KEY, topic TEXT, id TEXT);
|
|
CREATE TABLE collection_metadata(collection_id TEXT, key TEXT, str_value TEXT);
|
|
"""
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO segments VALUES (?, ?, ?, ?)",
|
|
[
|
|
(drawers_vec, "urn:vector", "VECTOR", drawers_coll),
|
|
(drawers_meta, "urn:metadata", "METADATA", drawers_coll),
|
|
(closets_vec, "urn:vector", "VECTOR", closets_coll),
|
|
(closets_meta, "urn:metadata", "METADATA", closets_coll),
|
|
],
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO max_seq_id(segment_id, seq_id) VALUES (?, ?)",
|
|
[
|
|
(drawers_vec, drawers_vec_poison),
|
|
(drawers_meta, drawers_meta_poison),
|
|
(closets_vec, closets_vec_poison),
|
|
(closets_meta, closets_meta_poison),
|
|
],
|
|
)
|
|
# Populate embeddings so the collection-MAX heuristic has data to work with.
|
|
# drawers METADATA owns the max at drawers_meta_max; closets likewise.
|
|
for i in range(1, drawers_meta_max + 1, max(drawers_meta_max // 5, 1)):
|
|
conn.execute(
|
|
"INSERT INTO embeddings(segment_id, embedding_id, seq_id) VALUES (?, ?, ?)",
|
|
(drawers_meta, f"d-{i}", i),
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO embeddings(segment_id, embedding_id, seq_id) VALUES (?, ?, ?)",
|
|
(drawers_meta, "d-max", drawers_meta_max),
|
|
)
|
|
for i in range(1, closets_meta_max + 1, max(closets_meta_max // 5, 1)):
|
|
conn.execute(
|
|
"INSERT INTO embeddings(segment_id, embedding_id, seq_id) VALUES (?, ?, ?)",
|
|
(closets_meta, f"c-{i}", i),
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO embeddings(segment_id, embedding_id, seq_id) VALUES (?, ?, ?)",
|
|
(closets_meta, "c-max", closets_meta_max),
|
|
)
|
|
conn.commit()
|
|
return {
|
|
"drawers_vec": drawers_vec,
|
|
"drawers_meta": drawers_meta,
|
|
"closets_vec": closets_vec,
|
|
"closets_meta": closets_meta,
|
|
"drawers_meta_max": drawers_meta_max,
|
|
"closets_meta_max": closets_meta_max,
|
|
"poisoned_values": {
|
|
drawers_vec: drawers_vec_poison,
|
|
drawers_meta: drawers_meta_poison,
|
|
closets_vec: closets_vec_poison,
|
|
closets_meta: closets_meta_poison,
|
|
},
|
|
}
|
|
|
|
|
|
def test_max_seq_id_detects_poison_rows(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
|
|
# Add one clean row to confirm the threshold actually filters.
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute(
|
|
"INSERT INTO segments VALUES ('seg-clean', 'urn:vector', 'VECTOR', 'coll-clean')"
|
|
)
|
|
conn.execute("INSERT INTO max_seq_id VALUES ('seg-clean', 1234)")
|
|
conn.commit()
|
|
|
|
found = repair._detect_poisoned_max_seq_ids(db_path)
|
|
ids = {sid for sid, _ in found}
|
|
assert ids == {
|
|
seg["drawers_vec"],
|
|
seg["drawers_meta"],
|
|
seg["closets_vec"],
|
|
seg["closets_meta"],
|
|
}
|
|
for sid, val in found:
|
|
assert val > repair.MAX_SEQ_ID_SANITY_THRESHOLD
|
|
assert "seg-clean" not in ids
|
|
|
|
|
|
def test_max_seq_id_heuristic_uses_collection_max(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
|
|
result = repair.repair_max_seq_id(palace, dry_run=True)
|
|
# Both drawers segments (VECTOR + METADATA) get the drawers collection max.
|
|
assert result["after"][seg["drawers_vec"]] == seg["drawers_meta_max"]
|
|
assert result["after"][seg["drawers_meta"]] == seg["drawers_meta_max"]
|
|
# Both closets segments get the closets collection max.
|
|
assert result["after"][seg["closets_vec"]] == seg["closets_meta_max"]
|
|
assert result["after"][seg["closets_meta"]] == seg["closets_meta_max"]
|
|
|
|
|
|
def test_max_seq_id_from_sidecar_exact_restore(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
|
|
# Craft a sidecar with known clean values that differ from the heuristic's
|
|
# collection-max, so we can prove the sidecar path is preferred.
|
|
sidecar_path = str(tmp_path / "chroma.sqlite3.sidecar")
|
|
clean = {
|
|
seg["drawers_vec"]: 499001,
|
|
seg["drawers_meta"]: 499002,
|
|
seg["closets_vec"]: 498001,
|
|
seg["closets_meta"]: 498002,
|
|
}
|
|
with sqlite3.connect(sidecar_path) as conn:
|
|
conn.execute("CREATE TABLE max_seq_id(segment_id TEXT PRIMARY KEY, seq_id INTEGER)")
|
|
conn.executemany(
|
|
"INSERT INTO max_seq_id VALUES (?, ?)",
|
|
list(clean.items()),
|
|
)
|
|
conn.commit()
|
|
|
|
result = repair.repair_max_seq_id(palace, from_sidecar=sidecar_path, assume_yes=True)
|
|
assert result["segment_repaired"]
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
with sqlite3.connect(db_path) as conn:
|
|
rows = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
for sid, val in clean.items():
|
|
assert rows[sid] == val
|
|
|
|
|
|
def test_max_seq_id_dry_run_no_mutation(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
|
|
with sqlite3.connect(db_path) as conn:
|
|
before = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
|
|
result = repair.repair_max_seq_id(palace, dry_run=True)
|
|
assert result["dry_run"] is True
|
|
assert result["segment_repaired"] == []
|
|
|
|
with sqlite3.connect(db_path) as conn:
|
|
after = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
assert before == after
|
|
# Nothing dropped into the palace dir either (no backup on dry-run).
|
|
assert not any(fn.startswith("chroma.sqlite3.max-seq-id-backup-") for fn in os.listdir(palace))
|
|
assert seg["drawers_vec"] in before # sanity
|
|
|
|
|
|
def test_max_seq_id_segment_filter(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
|
|
result = repair.repair_max_seq_id(palace, segment=seg["drawers_meta"], assume_yes=True)
|
|
assert result["segment_repaired"] == [seg["drawers_meta"]]
|
|
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
with sqlite3.connect(db_path) as conn:
|
|
rows = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
# Filtered segment is fixed; the other three remain poisoned.
|
|
assert rows[seg["drawers_meta"]] == seg["drawers_meta_max"]
|
|
for other in (seg["drawers_vec"], seg["closets_vec"], seg["closets_meta"]):
|
|
assert rows[other] > repair.MAX_SEQ_ID_SANITY_THRESHOLD
|
|
|
|
|
|
def test_max_seq_id_heuristic_decodes_blob_embeddings_seq_id(tmp_path):
|
|
"""`embeddings.seq_id` rows can be BLOB-typed on palaces where chromadb
|
|
1.5.x has been writing seq_ids natively (8-byte big-endian uint64).
|
|
`_compute_heuristic_seq_id` must decode those rather than crashing on
|
|
`int(bytes)` — the recovery feature is meaningless if it can't read
|
|
the storage format it was designed to repair.
|
|
"""
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
|
|
drawers_meta_max = seg["drawers_meta_max"]
|
|
blob_max = drawers_meta_max + 7
|
|
blob_value = blob_max.to_bytes(8, "big")
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute(
|
|
"INSERT INTO embeddings(segment_id, embedding_id, seq_id) VALUES (?, ?, ?)",
|
|
(seg["drawers_meta"], "d-blob-max", blob_value),
|
|
)
|
|
conn.commit()
|
|
|
|
result = repair.repair_max_seq_id(palace, dry_run=True)
|
|
assert result["after"][seg["drawers_vec"]] == blob_max
|
|
assert result["after"][seg["drawers_meta"]] == blob_max
|
|
|
|
|
|
def test_max_seq_id_no_poison_is_noop(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
os.makedirs(palace)
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE segments(
|
|
id TEXT PRIMARY KEY, type TEXT, scope TEXT, collection TEXT
|
|
);
|
|
CREATE TABLE max_seq_id(segment_id TEXT PRIMARY KEY, seq_id);
|
|
CREATE TABLE embeddings(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
segment_id TEXT, embedding_id TEXT, seq_id
|
|
);
|
|
INSERT INTO segments VALUES ('s1', 'urn:vector', 'VECTOR', 'coll');
|
|
INSERT INTO max_seq_id VALUES ('s1', 12345);
|
|
"""
|
|
)
|
|
conn.commit()
|
|
|
|
result = repair.repair_max_seq_id(palace, assume_yes=True)
|
|
assert result["segment_repaired"] == []
|
|
assert result["backup"] is None
|
|
with sqlite3.connect(db_path) as conn:
|
|
rows = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
assert rows == {"s1": 12345}
|
|
|
|
|
|
def test_max_seq_id_backup_created(tmp_path):
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(palace)
|
|
|
|
result = repair.repair_max_seq_id(palace, assume_yes=True)
|
|
assert result["backup"] is not None
|
|
assert os.path.isfile(result["backup"])
|
|
|
|
with sqlite3.connect(result["backup"]) as conn:
|
|
rows = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id").fetchall())
|
|
# Backup preserves the poisoned values from before the repair.
|
|
assert rows[seg["drawers_vec"]] == seg["poisoned_values"][seg["drawers_vec"]]
|
|
assert rows[seg["drawers_meta"]] == seg["poisoned_values"][seg["drawers_meta"]]
|
|
|
|
|
|
def test_max_seq_id_backup_pruned_to_max_backups(tmp_path, monkeypatch):
|
|
"""Old max-seq-id backups beyond MEMPALACE_MAX_BACKUPS are pruned after a repair.
|
|
|
|
Without retention, every repair left a full chroma.sqlite3 copy behind
|
|
that was never cleaned up — the unbounded disk-growth bug this guards.
|
|
"""
|
|
palace = str(tmp_path / "palace")
|
|
_seed_poisoned_max_seq_id(palace)
|
|
|
|
# Pre-seed 4 stale backups with old mtimes so the just-created one is
|
|
# unambiguously the newest.
|
|
for i in range(4):
|
|
stale = os.path.join(palace, f"chroma.sqlite3.max-seq-id-backup-2026010{i}-000000")
|
|
with open(stale, "w") as f:
|
|
f.write("old")
|
|
os.utime(stale, (1_700_000_000 + i, 1_700_000_000 + i))
|
|
|
|
monkeypatch.setenv("MEMPALACE_MAX_BACKUPS", "2")
|
|
|
|
result = repair.repair_max_seq_id(palace, assume_yes=True)
|
|
|
|
backups = sorted(
|
|
fn for fn in os.listdir(palace) if fn.startswith("chroma.sqlite3.max-seq-id-backup-")
|
|
)
|
|
# 4 stale + 1 fresh = 5 written; retention keeps only the 2 newest.
|
|
assert len(backups) == 2
|
|
# The backup created by this repair must be one of the survivors.
|
|
assert os.path.basename(result["backup"]) in backups
|
|
|
|
|
|
def test_max_seq_id_backup_retained_when_pruning_disabled(tmp_path, monkeypatch):
|
|
"""max_backups=0 keeps every backup (opt-out for external retention)."""
|
|
palace = str(tmp_path / "palace")
|
|
_seed_poisoned_max_seq_id(palace)
|
|
|
|
for i in range(3):
|
|
stale = os.path.join(palace, f"chroma.sqlite3.max-seq-id-backup-2026010{i}-000000")
|
|
with open(stale, "w") as f:
|
|
f.write("old")
|
|
os.utime(stale, (1_700_000_000 + i, 1_700_000_000 + i))
|
|
|
|
monkeypatch.setenv("MEMPALACE_MAX_BACKUPS", "0")
|
|
|
|
repair.repair_max_seq_id(palace, assume_yes=True)
|
|
|
|
backups = [
|
|
fn for fn in os.listdir(palace) if fn.startswith("chroma.sqlite3.max-seq-id-backup-")
|
|
]
|
|
assert len(backups) == 4
|
|
|
|
|
|
def test_max_seq_id_rollback_on_verification_failure(tmp_path, monkeypatch):
|
|
"""If the post-update detector still sees poison, raise and leave a backup."""
|
|
palace = str(tmp_path / "palace")
|
|
_seed_poisoned_max_seq_id(palace)
|
|
|
|
real_detect = repair._detect_poisoned_max_seq_ids
|
|
calls = {"n": 0}
|
|
|
|
def flaky_detect(*args, **kwargs):
|
|
calls["n"] += 1
|
|
# First call (pre-repair) returns the real set so the repair proceeds.
|
|
if calls["n"] == 1:
|
|
return real_detect(*args, **kwargs)
|
|
# Second call (post-repair verification) claims poison still exists.
|
|
return [("seg-fake-still-poisoned", repair.MAX_SEQ_ID_SANITY_THRESHOLD + 1)]
|
|
|
|
monkeypatch.setattr(repair, "_detect_poisoned_max_seq_ids", flaky_detect)
|
|
|
|
with pytest.raises(repair.MaxSeqIdVerificationError):
|
|
repair.repair_max_seq_id(palace, assume_yes=True)
|
|
|
|
# A backup file is still present — caller can roll back from it.
|
|
leftover = [fn for fn in os.listdir(palace) if "max-seq-id-backup-" in fn]
|
|
assert leftover
|
|
|
|
|
|
def test_sqlite_integrity_errors_returns_empty_for_healthy_db(tmp_path):
|
|
palace = tmp_path / "palace"
|
|
palace.mkdir()
|
|
db_path = palace / "chroma.sqlite3"
|
|
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute("CREATE TABLE dummy(id INTEGER PRIMARY KEY)")
|
|
conn.commit()
|
|
|
|
assert repair.sqlite_integrity_errors(str(palace)) == []
|
|
|
|
|
|
def test_sqlite_integrity_errors_reports_unreadable_sqlite_file(tmp_path):
|
|
palace = tmp_path / "palace"
|
|
palace.mkdir()
|
|
db_path = palace / "chroma.sqlite3"
|
|
db_path.write_bytes(b"not a sqlite database")
|
|
|
|
errors = repair.sqlite_integrity_errors(str(palace))
|
|
|
|
assert errors
|
|
assert "quick_check failed" in errors[0]
|
|
|
|
|
|
@patch("mempalace.repair._copy_file_no_follow")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_aborts_on_sqlite_integrity_errors_before_delete_collection(
|
|
mock_backend_cls,
|
|
mock_copy,
|
|
tmp_path,
|
|
capsys,
|
|
):
|
|
"""Regression for #1362: fail before Chroma delete_collection on sqlite corruption."""
|
|
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with sqlite3.connect(sqlite_path) as conn:
|
|
conn.execute("CREATE TABLE dummy(id INTEGER PRIMARY KEY)")
|
|
conn.commit()
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 2
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1", "id2"],
|
|
"documents": ["doc1", "doc2"],
|
|
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
|
}
|
|
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
|
|
with patch(
|
|
"mempalace.repair.sqlite_integrity_errors",
|
|
return_value=[
|
|
"Page 4 of B-tree 12345: database disk image is malformed",
|
|
"Page 8 of B-tree 67890: database disk image is malformed",
|
|
],
|
|
):
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
|
|
out = capsys.readouterr().out
|
|
|
|
assert "SQLite-layer corruption detected before repair rebuild" in out
|
|
assert "PRAGMA quick_check" in out
|
|
assert "delete_collection" in out
|
|
assert "Page 4 of B-tree" in out
|
|
|
|
mock_backend.delete_collection.assert_not_called()
|
|
mock_backend.create_collection.assert_not_called()
|
|
mock_copy.assert_not_called()
|
|
|
|
|
|
def test_rebuild_index_runs_sqlite_preflight_before_chromadb_open(tmp_path, capsys):
|
|
"""The SQLite integrity preflight must run BEFORE backend.get_collection.
|
|
|
|
chromadb's rust binding raises pyo3_runtime.PanicException (which is not
|
|
a regular Exception subclass) on a malformed page, so any get_collection
|
|
call against a corrupt SQLite propagates past `except Exception` handlers
|
|
and produces a 30-line stack trace instead of the friendly abort message.
|
|
Regression test for the ordering bug where the preflight was placed after
|
|
the chromadb client open and therefore never reached on the cases it was
|
|
designed to catch (#1364 follow-up).
|
|
"""
|
|
palace = tmp_path / "palace"
|
|
palace.mkdir()
|
|
|
|
# Build a real chromadb palace with one drawer so chroma.sqlite3 exists
|
|
# at full schema size, then mangle several middle pages so PRAGMA
|
|
# quick_check fails with "disk image is malformed". This matches the
|
|
# production failure mode users hit in #1362 / #1364.
|
|
from mempalace.backends.chroma import ChromaBackend
|
|
|
|
backend = ChromaBackend()
|
|
try:
|
|
col = backend.create_collection(str(palace), "mempalace_drawers")
|
|
col.upsert(
|
|
ids=["d1"],
|
|
documents=["doc"],
|
|
metadatas=[{"wing": "w", "room": "r"}],
|
|
)
|
|
finally:
|
|
backend.close()
|
|
|
|
sqlite_path = palace / "chroma.sqlite3"
|
|
pre_size = sqlite_path.stat().st_size
|
|
|
|
# Compute a page-aligned corruption offset that's always inside the
|
|
# existing file. SQLite uses 4 KB pages by default; we mangle 4 pages
|
|
# somewhere in the middle, skipping at least the first 2 pages
|
|
# (header + root) so the file still opens. Without clamping to the
|
|
# actual file size, a seek past EOF on r+b mode would silently
|
|
# extend the file with zero-padding and leave the original pages
|
|
# intact — quick_check would still pass, and the regression guard
|
|
# would skip the bug.
|
|
PAGE = 4096
|
|
CORRUPT_BYTES = 16384 # 4 pages
|
|
HEADER_GUARD = PAGE * 2 # leave header + root pages intact
|
|
assert pre_size >= HEADER_GUARD + CORRUPT_BYTES, (
|
|
f"sqlite db too small to mangle without truncating: {pre_size} bytes"
|
|
)
|
|
# Round (pre_size - CORRUPT_BYTES) down to a page boundary so we
|
|
# mangle whole pages. Cap at offset 40960 (page 10) for stable
|
|
# diagnostics across SQLite versions that may grow the file.
|
|
max_offset = (pre_size - CORRUPT_BYTES) & ~(PAGE - 1)
|
|
corrupt_offset = min(40960, max_offset)
|
|
assert corrupt_offset >= HEADER_GUARD, f"corruption offset {corrupt_offset} too close to header"
|
|
|
|
with open(sqlite_path, "r+b") as f:
|
|
f.seek(corrupt_offset)
|
|
f.write(b"\xde\xad\xbe\xef" * (CORRUPT_BYTES // 4))
|
|
|
|
# No chromadb mocks: rebuild_index must reach sqlite_integrity_errors
|
|
# before any code path that opens a chromadb client. If the preflight
|
|
# comes too late, the test fails with pyo3_runtime.PanicException
|
|
# instead of returning cleanly.
|
|
repair.rebuild_index(palace_path=str(palace))
|
|
|
|
out = capsys.readouterr().out
|
|
assert "SQLite-layer corruption detected before repair rebuild" in out
|
|
assert "PRAGMA quick_check" in out
|
|
assert "disk image is malformed" in out
|
|
|
|
|
|
def test_max_seq_id_preflight_preserves_embeddings_queue(tmp_path):
|
|
"""#1295: default repair preflight must not drop queued writes."""
|
|
|
|
palace = str(tmp_path / "palace")
|
|
seg = _seed_poisoned_max_seq_id(
|
|
palace,
|
|
drawers_meta_max=102,
|
|
closets_meta_max=11,
|
|
)
|
|
db_path = os.path.join(palace, "chroma.sqlite3")
|
|
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.executemany(
|
|
"INSERT INTO embeddings_queue(seq_id, topic, id) VALUES (?, ?, ?)",
|
|
[
|
|
(seq_id, "persistent://default/default/mempalace_drawers", f"queued-{seq_id}")
|
|
for seq_id in range(103, 123)
|
|
],
|
|
)
|
|
conn.commit()
|
|
|
|
result = repair.maybe_repair_poisoned_max_seq_id_before_rebuild(
|
|
palace,
|
|
assume_yes=True,
|
|
)
|
|
|
|
assert result is not None
|
|
assert result["segment_repaired"]
|
|
|
|
with sqlite3.connect(db_path) as conn:
|
|
max_seq_rows = dict(conn.execute("SELECT segment_id, seq_id FROM max_seq_id"))
|
|
queue_count = conn.execute("SELECT COUNT(*) FROM embeddings_queue").fetchone()[0]
|
|
|
|
assert max_seq_rows[seg["drawers_vec"]] == seg["drawers_meta_max"]
|
|
assert max_seq_rows[seg["drawers_meta"]] == seg["drawers_meta_max"]
|
|
assert max_seq_rows[seg["closets_vec"]] == seg["closets_meta_max"]
|
|
assert max_seq_rows[seg["closets_meta"]] == seg["closets_meta_max"]
|
|
|
|
# The old legacy rebuild path can discard queued writes. The preflight
|
|
# repair must leave them on disk for Chroma to drain after the bookmark is
|
|
# unpoisoned.
|
|
assert queue_count == 20
|
|
|
|
|
|
def test_rebuild_index_repairs_poisoned_max_seq_id_before_collection_rebuild(tmp_path, capsys):
|
|
"""A poisoned bookmark should short-circuit before the legacy rebuild path."""
|
|
|
|
palace = str(tmp_path / "palace")
|
|
_seed_poisoned_max_seq_id(palace)
|
|
|
|
with patch("mempalace.repair.ChromaBackend") as mock_backend:
|
|
repair.rebuild_index(palace)
|
|
|
|
out = capsys.readouterr().out
|
|
backend = mock_backend.return_value
|
|
|
|
# repair_max_seq_id may instantiate ChromaBackend to close cached clients
|
|
# after editing sqlite directly. That is safe. The important thing is that
|
|
# rebuild_index must not continue into the legacy Chroma collection read /
|
|
# count / rebuild path after the max_seq_id preflight handles the issue.
|
|
backend.get_collection.assert_not_called()
|
|
|
|
assert "Detected poisoned max_seq_id rows" in out
|
|
assert "non-destructive max_seq_id repair" in out
|
|
|
|
|
|
# ── extract_via_sqlite + rebuild_from_sqlite (#1308) ──────────────────
|
|
#
|
|
# These tests build real chromadb palaces in tmp_path rather than mocking
|
|
# the SQLite layer. The bug class they guard against is "extraction sees
|
|
# different rows than chromadb stored" — the only honest check is to let
|
|
# chromadb actually write rows and then read them back via the SQLite
|
|
# bypass. Mocking the SQLite cursor would defeat the test.
|
|
|
|
|
|
def _seed_palace(palace_path, collection_name, rows):
|
|
"""Build a real chromadb palace at ``palace_path`` and add ``rows``.
|
|
|
|
``rows`` is a list of ``(id, document, metadata)`` tuples.
|
|
"""
|
|
from mempalace.backends.chroma import ChromaBackend
|
|
|
|
backend = ChromaBackend()
|
|
try:
|
|
col = backend.create_collection(str(palace_path), collection_name)
|
|
col.upsert(
|
|
ids=[r[0] for r in rows],
|
|
documents=[r[1] for r in rows],
|
|
metadatas=[r[2] for r in rows],
|
|
)
|
|
finally:
|
|
# Release chromadb's rust-side SQLite/HNSW file locks before the
|
|
# caller proceeds. Without this, an in-place rebuild on Windows
|
|
# fails with WinError 32 on data_level0.bin during the archive
|
|
# rename (cf. PR #1310 test-windows job).
|
|
backend.close()
|
|
|
|
|
|
def test_extract_via_sqlite_returns_all_rows_with_metadata(tmp_path):
|
|
"""Round-trip: a chromadb palace with N upserted rows returns those
|
|
same N rows when read via the SQLite bypass.
|
|
|
|
Catches: anyone who breaks the segments/embeddings/embedding_metadata
|
|
JOIN, swaps the metadata vs vector segment, or changes how the
|
|
document is stored under the ``chroma:document`` key.
|
|
|
|
Also asserts every embedding row underlying the extraction lives in
|
|
a ``segments.scope = 'METADATA'`` segment. Document + metadata rows
|
|
are stored under METADATA in Chroma's segment layout while HNSW
|
|
files live under ``VECTOR``; locking that assumption in here means a
|
|
future refactor that accidentally points the JOIN at ``VECTOR``
|
|
fails this test instead of silently regressing the recovery path.
|
|
"""
|
|
rows = [
|
|
(f"drawer_{i:03d}", f"document body {i}", {"wing": "test_wing", "room": f"r{i % 3}"})
|
|
for i in range(25)
|
|
]
|
|
_seed_palace(tmp_path, "mempalace_drawers", rows)
|
|
|
|
extracted = list(repair.extract_via_sqlite(str(tmp_path), "mempalace_drawers"))
|
|
|
|
assert len(extracted) == 25
|
|
by_id = {emb_id: (doc, meta) for emb_id, doc, meta in extracted}
|
|
assert set(by_id) == {r[0] for r in rows}
|
|
for emb_id, doc, meta in rows:
|
|
got_doc, got_meta = by_id[emb_id]
|
|
assert got_doc == doc, f"document mangled for {emb_id}"
|
|
assert got_meta == meta, f"metadata mangled for {emb_id}: {got_meta!r}"
|
|
|
|
# Lock the segment-scope assumption directly against Chroma's on-disk
|
|
# layout so a future change that points the extraction JOIN at the
|
|
# VECTOR segment cannot pass this test. Query each extracted row's
|
|
# backing segment scope via the same SQLite tables ``extract_via_sqlite``
|
|
# reads from.
|
|
sqlite_path = os.path.join(str(tmp_path), "chroma.sqlite3")
|
|
conn = sqlite3.connect(f"file:{sqlite_path}?mode=ro", uri=True)
|
|
try:
|
|
scopes = {
|
|
scope
|
|
for (scope,) in conn.execute(
|
|
"""
|
|
SELECT DISTINCT s.scope
|
|
FROM embeddings e
|
|
JOIN segments s ON e.segment_id = s.id
|
|
JOIN collections c ON s.collection = c.id
|
|
WHERE c.name = ? AND e.embedding_id IN ({})
|
|
""".format(",".join("?" * len(extracted))),
|
|
("mempalace_drawers", *(emb_id for emb_id, _, _ in extracted)),
|
|
)
|
|
}
|
|
finally:
|
|
conn.close()
|
|
assert scopes == {"METADATA"}, (
|
|
f"extraction is reading from segments scoped {scopes!r}; only "
|
|
"'METADATA' should back the document/metadata rows. If Chroma's "
|
|
"segment layout changed, update extract_via_sqlite's WHERE clause."
|
|
)
|
|
|
|
|
|
def test_extract_via_sqlite_preserves_typed_metadata(tmp_path):
|
|
"""Chromadb stores int / float / bool / string in distinct typed
|
|
columns. Extraction must round-trip the original type, not coerce
|
|
everything to string.
|
|
|
|
Catches: a regression where the SELECT order changes and ints come
|
|
back as None, or where the column-resolution rule prefers the wrong
|
|
column.
|
|
"""
|
|
rows = [
|
|
(
|
|
"drawer_typed",
|
|
"doc",
|
|
{
|
|
"wing": "w",
|
|
"chunk_index": 7, # int
|
|
"score": 0.42, # float
|
|
"is_active": True, # bool
|
|
},
|
|
),
|
|
]
|
|
_seed_palace(tmp_path, "mempalace_drawers", rows)
|
|
|
|
extracted = list(repair.extract_via_sqlite(str(tmp_path), "mempalace_drawers"))
|
|
assert len(extracted) == 1
|
|
_, _, meta = extracted[0]
|
|
|
|
assert meta["chunk_index"] == 7 and isinstance(meta["chunk_index"], int)
|
|
assert meta["score"] == 0.42 and isinstance(meta["score"], float)
|
|
assert meta["is_active"] is True
|
|
assert meta["wing"] == "w"
|
|
|
|
|
|
def test_extract_via_sqlite_unknown_collection_yields_nothing(tmp_path):
|
|
"""Asking for a collection that isn't in the palace must return an
|
|
empty iterator, not silently fall back to another collection's
|
|
metadata segment. Seeds two real collections and queries for a third
|
|
name so a regression that drops the WHERE c.name=? filter would leak
|
|
rows from the seeded collections rather than passing.
|
|
"""
|
|
_seed_palace(tmp_path, "mempalace_drawers", [("d1", "doc", {"wing": "w"})])
|
|
_seed_palace(tmp_path, "mempalace_closets", [("c1", "abbrev", {"wing": "w"})])
|
|
assert list(repair.extract_via_sqlite(str(tmp_path), "not_a_real_collection")) == []
|
|
|
|
|
|
def test_extract_via_sqlite_missing_palace_yields_nothing(tmp_path):
|
|
"""No chroma.sqlite3 → empty iterator, no exception. Callers depend
|
|
on this when probing speculatively."""
|
|
empty = tmp_path / "no_palace_here"
|
|
empty.mkdir()
|
|
assert list(repair.extract_via_sqlite(str(empty), "mempalace_drawers")) == []
|
|
|
|
|
|
def test_rebuild_from_sqlite_roundtrips_via_real_chromadb(tmp_path):
|
|
"""End-to-end: seed source palace, rebuild into a fresh dest, then
|
|
open dest with a fresh ChromaBackend and verify ``count()`` and
|
|
metadata filters return the original rows. Also asserts a closet
|
|
document round-trips so a future regression that re-embeds with the
|
|
wrong EF or swaps drawer/closet content would fail here.
|
|
|
|
This is the single most important regression guard. If
|
|
``rebuild_from_sqlite`` silently drops rows or mangles metadata, no
|
|
other test in this file would catch it because they all stop at the
|
|
extraction layer.
|
|
"""
|
|
from mempalace.backends.chroma import ChromaBackend
|
|
|
|
source = tmp_path / "source"
|
|
dest = tmp_path / "dest"
|
|
|
|
rows = [
|
|
(f"drawer_{i:03d}", f"body {i}", {"wing": "alpha" if i % 2 else "beta", "room": "r0"})
|
|
for i in range(40)
|
|
]
|
|
_seed_palace(source, "mempalace_drawers", rows)
|
|
_seed_palace(
|
|
source,
|
|
"mempalace_closets",
|
|
[("closet_x", "abbrev pointer →drawer_001", {"wing": "alpha"})],
|
|
)
|
|
|
|
counts = repair.rebuild_from_sqlite(str(source), str(dest))
|
|
assert counts == {"mempalace_drawers": 40, "mempalace_closets": 1}
|
|
|
|
backend = ChromaBackend()
|
|
drawers = backend.get_collection(str(dest), "mempalace_drawers")
|
|
assert drawers.count() == 40
|
|
alpha = drawers.get(where={"wing": "alpha"})
|
|
assert len(alpha["ids"]) == 20
|
|
|
|
# Spot-check that document text round-trips for one specific drawer
|
|
# — protects against a regression where extraction or upsert order
|
|
# silently swaps document bodies between IDs.
|
|
one = drawers.get(ids=["drawer_007"], include=["documents", "metadatas"])
|
|
assert one["documents"] == ["body 7"]
|
|
assert one["metadatas"][0]["wing"] == "alpha"
|
|
|
|
# Closets: the AAAK index layer. Re-embedded with the same EF so a
|
|
# known closet ID and its document body must come back intact.
|
|
closets = backend.get_collection(str(dest), "mempalace_closets")
|
|
assert closets.count() == 1
|
|
closet_row = closets.get(ids=["closet_x"], include=["documents", "metadatas"])
|
|
assert closet_row["documents"] == ["abbrev pointer →drawer_001"]
|
|
assert closet_row["metadatas"][0] == {"wing": "alpha"}
|
|
|
|
|
|
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
|
|
interleave a partial rebuild with whatever's already at dest.
|
|
"""
|
|
source = tmp_path / "source"
|
|
dest = tmp_path / "dest"
|
|
_seed_palace(source, "mempalace_drawers", [("d1", "doc", {"wing": "w"})])
|
|
dest.mkdir()
|
|
# Drop a marker file so we can prove the dir wasn't touched.
|
|
(dest / "marker.txt").write_text("preexisting")
|
|
|
|
counts = repair.rebuild_from_sqlite(str(source), str(dest))
|
|
assert counts == {}
|
|
assert (dest / "marker.txt").read_text() == "preexisting"
|
|
assert not (dest / "chroma.sqlite3").exists()
|
|
|
|
|
|
def test_rebuild_from_sqlite_in_place_archives_when_opted_in(tmp_path):
|
|
"""In-place rebuild (source == dest) with ``archive_existing_dest=True``
|
|
must move the original aside to ``<dest>.pre-rebuild-<ts>`` and read
|
|
from the archive — the original drawer rows must survive in the new
|
|
palace, AND the archive itself must still contain the original rows.
|
|
|
|
Catches: a refactor that moves the original out but then reads from
|
|
the now-empty original location, producing an empty rebuild; also
|
|
catches a swap that empties the archive after reading.
|
|
"""
|
|
palace = tmp_path / "palace"
|
|
rows = [(f"d{i}", f"body {i}", {"wing": "w", "room": "r"}) for i in range(15)]
|
|
_seed_palace(palace, "mempalace_drawers", rows)
|
|
|
|
counts = repair.rebuild_from_sqlite(str(palace), str(palace), archive_existing_dest=True)
|
|
assert counts["mempalace_drawers"] == 15
|
|
|
|
archives = [p for p in tmp_path.iterdir() if p.name.startswith("palace.pre-rebuild-")]
|
|
assert len(archives) == 1
|
|
assert (archives[0] / "chroma.sqlite3").exists()
|
|
# Archive must still hold the same row count via the SQLite bypass —
|
|
# proves the archive wasn't silently truncated as a side effect.
|
|
archived_rows = list(repair.extract_via_sqlite(str(archives[0]), "mempalace_drawers"))
|
|
assert len(archived_rows) == 15
|
|
|
|
from mempalace.backends.chroma import ChromaBackend
|
|
|
|
rebuilt = ChromaBackend().get_collection(str(palace), "mempalace_drawers")
|
|
assert rebuilt.count() == 15
|
|
|
|
|
|
def test_rebuild_from_sqlite_in_place_refuses_without_archive_flag(tmp_path):
|
|
"""Source == dest without archive flag must abort untouched. The
|
|
most catastrophic possible regression of this code path is silently
|
|
deleting the only copy of the user's data."""
|
|
palace = tmp_path / "palace"
|
|
_seed_palace(palace, "mempalace_drawers", [("d1", "doc", {"wing": "w"})])
|
|
sqlite_before = (palace / "chroma.sqlite3").stat().st_size
|
|
|
|
counts = repair.rebuild_from_sqlite(str(palace), str(palace))
|
|
assert counts == {}
|
|
# Same file, untouched.
|
|
assert (palace / "chroma.sqlite3").stat().st_size == sqlite_before
|
|
archives = [p for p in tmp_path.iterdir() if "pre-rebuild" in p.name]
|
|
assert archives == []
|
|
|
|
|
|
def test_rebuild_from_sqlite_in_place_archive_failure_leaves_palace_untouched(
|
|
tmp_path, monkeypatch
|
|
):
|
|
"""A file inside the palace held open by another process (MCP server,
|
|
a running mine, another harness) must abort the archive step cleanly,
|
|
leaving the live palace fully intact.
|
|
|
|
Regression test for a real-world incident (2026-07-05/06, Windows 11):
|
|
the archive step used ``shutil.move``, whose fallback for a failed
|
|
``os.rename`` is copytree + rmtree. That rmtree deletes the live
|
|
palace file-by-file until it hits the first locked file, so an
|
|
in-progress mine or a live MCP server holding one file open left the
|
|
palace partially gutted next to a partial archive copy -- twice, on
|
|
two separate nights. os.rename fails atomically up front with nothing
|
|
touched; this test locks that behaviour in so a future change back to
|
|
shutil.move (or an equivalent copy+delete fallback) fails loudly.
|
|
"""
|
|
palace = tmp_path / "palace"
|
|
rows = [("d1", "doc one", {"wing": "w"}), ("d2", "doc two", {"wing": "w"})]
|
|
_seed_palace(palace, "mempalace_drawers", rows)
|
|
sqlite_before = (palace / "chroma.sqlite3").stat().st_size
|
|
entries_before = sorted(p.name for p in palace.iterdir())
|
|
|
|
def _raise(*_args, **_kwargs):
|
|
raise PermissionError("[WinError 32] simulated: file held open by another process")
|
|
|
|
monkeypatch.setattr(repair.os, "rename", _raise)
|
|
|
|
counts = repair.rebuild_from_sqlite(str(palace), str(palace), archive_existing_dest=True)
|
|
|
|
assert counts == {}
|
|
# Palace directory contents and the sqlite file itself are byte-for-byte
|
|
# untouched -- no partial delete, no partial archive left behind.
|
|
assert sorted(p.name for p in palace.iterdir()) == entries_before
|
|
assert (palace / "chroma.sqlite3").stat().st_size == sqlite_before
|
|
archives = [p for p in tmp_path.iterdir() if "pre-rebuild" in p.name]
|
|
assert archives == []
|
|
|
|
|
|
def test_rebuild_from_sqlite_source_missing_chroma_db(tmp_path):
|
|
"""Source dir exists but has no chroma.sqlite3 → returns empty,
|
|
leaves dest untouched."""
|
|
source = tmp_path / "source"
|
|
source.mkdir()
|
|
(source / "stray_file").write_text("not a palace")
|
|
dest = tmp_path / "dest"
|
|
|
|
counts = repair.rebuild_from_sqlite(str(source), str(dest))
|
|
assert counts == {}
|
|
assert not dest.exists()
|
|
|
|
|
|
def test_rebuild_from_sqlite_in_place_validates_source_before_archiving(tmp_path):
|
|
"""In-place + archive_existing_dest=True with a dir that lacks
|
|
chroma.sqlite3 must NOT rename the dir before bailing. An earlier
|
|
revision archived first and validated second, leaving the user with
|
|
a renamed empty dir to manually undo. Catches that ordering bug.
|
|
"""
|
|
palace = tmp_path / "palace"
|
|
palace.mkdir()
|
|
(palace / "marker.txt").write_text("not a real palace")
|
|
|
|
counts = repair.rebuild_from_sqlite(str(palace), str(palace), archive_existing_dest=True)
|
|
assert counts == {}
|
|
# No archive created — original dir still in place with its marker.
|
|
assert palace.exists()
|
|
assert (palace / "marker.txt").read_text() == "not a real palace"
|
|
archives = [p for p in tmp_path.iterdir() if "pre-rebuild" in p.name]
|
|
assert archives == []
|
|
|
|
|
|
def test_rebuild_from_sqlite_raises_on_upsert_failure(tmp_path, monkeypatch):
|
|
"""Mid-batch upsert failure must raise ``RebuildPartialError`` and
|
|
surface the failed collection + archive path so the user can recover.
|
|
Without this, an unattended script gets exit-code-zero on a partial
|
|
rebuild and the user discovers the data loss only when search starts
|
|
returning fewer hits.
|
|
"""
|
|
palace = tmp_path / "palace"
|
|
rows = [(f"d{i}", f"body {i}", {"wing": "w", "room": "r"}) for i in range(5)]
|
|
_seed_palace(palace, "mempalace_drawers", rows)
|
|
|
|
# Make the very first upsert raise so we don't depend on batch
|
|
# boundary behavior. Patching ChromaCollection.upsert (the wrapper
|
|
# mempalace's backend returns) keeps the failure path realistic.
|
|
# ``monkeypatch`` is pytest's built-in fixture that auto-restores
|
|
# the original attribute when the test exits, so we don't need to
|
|
# undo this manually.
|
|
from mempalace.backends.chroma import ChromaCollection
|
|
|
|
def boom(self, **kwargs):
|
|
raise RuntimeError("simulated chromadb upsert failure")
|
|
|
|
monkeypatch.setattr(ChromaCollection, "upsert", boom)
|
|
|
|
with pytest.raises(repair.RebuildPartialError) as excinfo:
|
|
repair.rebuild_from_sqlite(str(palace), str(palace), archive_existing_dest=True)
|
|
|
|
err = excinfo.value
|
|
assert err.failed_collection == "mempalace_drawers"
|
|
assert err.partial_counts.get("mempalace_drawers") == 0
|
|
assert err.archive_path is not None
|
|
assert os.path.isfile(os.path.join(err.archive_path, "chroma.sqlite3"))
|
|
assert err.dest_palace == os.path.abspath(str(palace))
|
|
|
|
|
|
def test_rebuild_from_sqlite_honors_configured_drawer_collection_name(tmp_path, monkeypatch):
|
|
"""A user with a non-default drawers collection name (set via
|
|
``MempalaceConfig().collection_name``) must have THAT collection
|
|
rebuilt — not the hardcoded ``mempalace_drawers``.
|
|
|
|
Catches: a regression where the recovery path silently rebuilds the
|
|
default-name collection on a custom-named palace, leaving the user's
|
|
actual data unrebuilt while reporting "rebuild complete." This is
|
|
the failure mode reviewer mjc flagged on PR #1310 as needing to line
|
|
up with the configured-collection-name work in #1312. Closets stay
|
|
fixed (``mempalace_closets``) by design — the AAAK index references
|
|
drawer IDs by string and is not per-deployment configurable.
|
|
|
|
Strategy: monkeypatch the lazy resolver so the test is hermetic and
|
|
does not depend on the global config file or env state.
|
|
"""
|
|
from mempalace.backends.chroma import ChromaBackend
|
|
|
|
custom_drawers = "custom_drawers_xyz"
|
|
monkeypatch.setattr(repair, "_drawers_collection_name", lambda: custom_drawers)
|
|
|
|
source = tmp_path / "source"
|
|
dest = tmp_path / "dest"
|
|
|
|
drawer_rows = [(f"d{i}", f"body {i}", {"wing": "alpha"}) for i in range(3)]
|
|
closet_rows = [("closet_a", "abbrev →d0", {"wing": "alpha"})]
|
|
_seed_palace(source, custom_drawers, drawer_rows)
|
|
_seed_palace(source, "mempalace_closets", closet_rows)
|
|
|
|
counts = repair.rebuild_from_sqlite(str(source), str(dest))
|
|
|
|
# Rebuilt under the custom name, not under the default "mempalace_drawers".
|
|
assert counts == {custom_drawers: 3, "mempalace_closets": 1}
|
|
|
|
backend = ChromaBackend()
|
|
rebuilt_drawers = backend.get_collection(str(dest), custom_drawers)
|
|
assert rebuilt_drawers.count() == 3
|
|
|
|
# Default-name collection must NOT exist in dest — proves we did not
|
|
# silently fall back to the hardcoded name during rebuild.
|
|
try:
|
|
rebuilt_default = backend.get_collection(str(dest), "mempalace_drawers")
|
|
# If get_collection returns without raising, count() should be 0
|
|
# (chromadb may auto-create on get with some EFs); a non-zero
|
|
# count would mean we wrote rows to the wrong collection.
|
|
assert rebuilt_default.count() == 0, (
|
|
"rebuild leaked rows into the default-name collection on a "
|
|
"custom-name palace — recovery wrote to the wrong collection."
|
|
)
|
|
except Exception:
|
|
pass # Expected: collection wasn't created.
|
|
|
|
|
|
# ── _vacuum_and_rebuild_fts5 ──────────────────────────────────────────
|
|
|
|
|
|
def test_vacuum_and_rebuild_fts5_vacuums_and_rebuilds(tmp_path):
|
|
"""VACUUM runs and FTS5 index is rebuilt when the table is present."""
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
conn.execute(
|
|
"CREATE VIRTUAL TABLE embedding_fulltext_search"
|
|
" USING fts5(string_value, tokenize='unicode61')"
|
|
)
|
|
conn.execute("INSERT INTO embedding_fulltext_search(string_value) VALUES('hello world')")
|
|
conn.commit()
|
|
|
|
repair._vacuum_and_rebuild_fts5(str(tmp_path))
|
|
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
result = conn.execute("PRAGMA integrity_check").fetchall()
|
|
assert result == [("ok",)]
|
|
|
|
|
|
def test_vacuum_and_rebuild_fts5_no_fts5_table(tmp_path):
|
|
"""VACUUM runs without error when embedding_fulltext_search is absent."""
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
conn.execute("CREATE TABLE dummy (id INTEGER PRIMARY KEY)")
|
|
conn.commit()
|
|
|
|
# Must not raise even without the FTS5 table.
|
|
repair._vacuum_and_rebuild_fts5(str(tmp_path))
|
|
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
result = conn.execute("PRAGMA integrity_check").fetchall()
|
|
assert result == [("ok",)]
|
|
|
|
|
|
def test_vacuum_and_rebuild_fts5_missing_sqlite(tmp_path):
|
|
"""Silently skips when chroma.sqlite3 does not exist."""
|
|
repair._vacuum_and_rebuild_fts5(str(tmp_path)) # no file — must not raise
|
|
|
|
|
|
# ── FTS5 inverted-index auto-heal (#1596) ─────────────────────────────
|
|
|
|
|
|
def _make_fts5_palace(tmp_path, *, corrupt: bool) -> str:
|
|
"""Build a palace whose embedding_fulltext_search index is optionally
|
|
corrupted to the malformed-inverted-index quick_check state #1596 hits."""
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
conn.execute(
|
|
"CREATE VIRTUAL TABLE embedding_fulltext_search"
|
|
" USING fts5(string_value, tokenize='unicode61')"
|
|
)
|
|
for i in range(200):
|
|
conn.execute(
|
|
"INSERT INTO embedding_fulltext_search(string_value) VALUES(?)",
|
|
(f"alpha beta gamma row{i} delta epsilon",),
|
|
)
|
|
conn.commit()
|
|
if corrupt:
|
|
# Zero the last index segment leaf: quick_check then reports
|
|
# "malformed inverted index" while the content table stays intact.
|
|
conn.execute(
|
|
"UPDATE embedding_fulltext_search_data SET block=zeroblob(length(block)) "
|
|
"WHERE id=(SELECT max(id) FROM embedding_fulltext_search_data)"
|
|
)
|
|
conn.commit()
|
|
return str(tmp_path)
|
|
|
|
|
|
def test_errors_are_isolated_fts5_classification():
|
|
fts = "malformed inverted index for FTS5 table main.embedding_fulltext_search"
|
|
page = "Page 4 of B-tree 12345: database disk image is malformed"
|
|
assert repair._errors_are_isolated_fts5([fts])
|
|
assert repair._errors_are_isolated_fts5([fts, fts])
|
|
assert not repair._errors_are_isolated_fts5([])
|
|
assert not repair._errors_are_isolated_fts5([page])
|
|
# Any non-FTS5 error in the set means the data itself may be damaged.
|
|
assert not repair._errors_are_isolated_fts5([fts, page])
|
|
|
|
|
|
def test_maybe_autoheal_fts5_index_heals_isolated_corruption(tmp_path):
|
|
palace = _make_fts5_palace(tmp_path, corrupt=True)
|
|
errors = repair.sqlite_integrity_errors(palace)
|
|
assert errors and repair._errors_are_isolated_fts5(errors)
|
|
|
|
remaining = repair.maybe_autoheal_fts5_index(palace, errors, progress=lambda *_: None)
|
|
|
|
assert remaining == []
|
|
# quick_check is clean and full-text search works again.
|
|
assert repair.sqlite_integrity_errors(palace) == []
|
|
with closing(sqlite3.connect(str(tmp_path / "chroma.sqlite3"))) as conn:
|
|
hits = conn.execute(
|
|
"SELECT count(*) FROM embedding_fulltext_search "
|
|
"WHERE embedding_fulltext_search MATCH 'gamma'"
|
|
).fetchone()[0]
|
|
assert hits == 200
|
|
|
|
|
|
def test_maybe_autoheal_fts5_index_leaves_non_fts5_errors_untouched(tmp_path):
|
|
palace = _make_fts5_palace(tmp_path, corrupt=False)
|
|
page_errors = ["Page 4 of B-tree 12345: database disk image is malformed"]
|
|
|
|
# Not isolated FTS5: returned unchanged and the rebuild is never attempted.
|
|
with patch("mempalace.palace.mine_palace_lock") as lock:
|
|
remaining = repair.maybe_autoheal_fts5_index(palace, page_errors, progress=lambda *_: None)
|
|
assert remaining == page_errors
|
|
lock.assert_not_called()
|
|
|
|
|
|
def test_maybe_autoheal_fts5_index_skips_when_palace_is_being_mined(tmp_path):
|
|
from mempalace.palace import MineAlreadyRunning
|
|
|
|
palace = _make_fts5_palace(tmp_path, corrupt=True)
|
|
errors = repair.sqlite_integrity_errors(palace)
|
|
|
|
def _raise(_path):
|
|
raise MineAlreadyRunning("held by pid 999")
|
|
|
|
# A live mine holds the lock: do not race the rebuild — surface and abort.
|
|
with patch("mempalace.palace.mine_palace_lock", side_effect=_raise):
|
|
remaining = repair.maybe_autoheal_fts5_index(palace, errors, progress=lambda *_: None)
|
|
|
|
assert remaining == errors
|
|
# The FTS index is still corrupt because we refused to rebuild under contention.
|
|
assert repair.sqlite_integrity_errors(palace) == errors
|
|
|
|
|
|
def test_rebuild_index_preflight_autoheals_isolated_fts5_then_proceeds(tmp_path, monkeypatch):
|
|
"""The preflight no longer hard-aborts on isolated FTS5 corruption (#1596):
|
|
it rebuilds the index, then continues into the rebuild path."""
|
|
palace = _make_fts5_palace(tmp_path, corrupt=True)
|
|
|
|
called = {}
|
|
|
|
def _fake_max_seq(_palace_path, **_kwargs):
|
|
# Reached only if the preflight did NOT abort — record and stop early
|
|
# so the test doesn't need a real chromadb collection.
|
|
called["reached"] = True
|
|
return {"stopped": True}
|
|
|
|
monkeypatch.setattr(repair, "maybe_repair_poisoned_max_seq_id_before_rebuild", _fake_max_seq)
|
|
|
|
repair.rebuild_index(palace_path=palace, progress=lambda *_: None)
|
|
|
|
assert called.get("reached") is True
|
|
assert repair.sqlite_integrity_errors(palace) == []
|
|
|
|
|
|
@patch("mempalace.repair.shutil")
|
|
@patch("mempalace.repair.ChromaBackend")
|
|
def test_rebuild_index_calls_vacuum(mock_backend_cls, mock_shutil, tmp_path):
|
|
"""rebuild_index closes chroma handles then calls _vacuum_and_rebuild_fts5.
|
|
|
|
ChromaDB's PersistentClient holds an open connection to chroma.sqlite3;
|
|
VACUUM requires an exclusive lock so _close_chroma_handles must be called
|
|
before _vacuum_and_rebuild_fts5.
|
|
"""
|
|
sqlite_path = tmp_path / "chroma.sqlite3"
|
|
with closing(sqlite3.connect(str(sqlite_path))) as conn:
|
|
conn.execute("CREATE TABLE dummy(id INTEGER PRIMARY KEY)")
|
|
conn.commit()
|
|
|
|
mock_col = MagicMock()
|
|
mock_col.count.return_value = 1
|
|
mock_col.get.return_value = {
|
|
"ids": ["id1"],
|
|
"documents": ["doc1"],
|
|
"metadatas": [{"wing": "a"}],
|
|
}
|
|
mock_new_col = MagicMock()
|
|
mock_new_col.count.return_value = 1
|
|
mock_temp_col = MagicMock()
|
|
mock_temp_col.count.return_value = 1
|
|
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
|
mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col]
|
|
|
|
call_order = []
|
|
with (
|
|
patch.object(
|
|
repair, "_close_chroma_handles", side_effect=lambda *a, **kw: call_order.append("close")
|
|
) as mock_close,
|
|
patch.object(
|
|
repair,
|
|
"_vacuum_and_rebuild_fts5",
|
|
side_effect=lambda *a, **kw: call_order.append("vacuum"),
|
|
) as mock_vacuum,
|
|
):
|
|
repair.rebuild_index(palace_path=str(tmp_path))
|
|
mock_close.assert_called_once()
|
|
mock_vacuum.assert_called_once()
|
|
assert call_order == ["close", "vacuum"], "backend must be closed before VACUUM"
|
|
args, kwargs = mock_vacuum.call_args
|
|
assert args[0] == str(tmp_path)
|
|
assert "progress" in kwargs
|
|
|
|
|
|
def test_rebuild_from_sqlite_preserves_knowledge_graph_sidecar(tmp_path):
|
|
"""The from-sqlite repair path must not drop the KG SQLite sidecar."""
|
|
src = tmp_path / "source"
|
|
dest = tmp_path / "dest"
|
|
src.mkdir()
|
|
dest.mkdir()
|
|
|
|
(src / "knowledge_graph.sqlite3").write_text("kg-db", encoding="utf-8")
|
|
(src / "knowledge_graph.sqlite3-wal").write_text("kg-wal", encoding="utf-8")
|
|
(src / "knowledge_graph.sqlite3-shm").write_text("kg-shm", encoding="utf-8")
|
|
|
|
copied = repair._preserve_knowledge_graph_sqlite(str(src), str(dest))
|
|
|
|
assert copied == [
|
|
"knowledge_graph.sqlite3",
|
|
"knowledge_graph.sqlite3-wal",
|
|
"knowledge_graph.sqlite3-shm",
|
|
]
|
|
assert (dest / "knowledge_graph.sqlite3").read_text(encoding="utf-8") == "kg-db"
|
|
assert (dest / "knowledge_graph.sqlite3-wal").read_text(encoding="utf-8") == "kg-wal"
|
|
assert (dest / "knowledge_graph.sqlite3-shm").read_text(encoding="utf-8") == "kg-shm"
|