diff --git a/mempalace/backends/chroma.py b/mempalace/backends/chroma.py index 0730a7b..c8c8c14 100644 --- a/mempalace/backends/chroma.py +++ b/mempalace/backends/chroma.py @@ -2,6 +2,7 @@ import contextlib import datetime as _dt +import json import logging import os import pickle @@ -665,6 +666,7 @@ def _pin_hnsw_threads(collection) -> None: _BLOB_FIX_MARKER = ".blob_seq_ids_migrated" +_COLLECTION_TYPE_MARKER = ".collection_type_fixed" def _valid_dimensionality(value: object) -> bool: @@ -883,6 +885,67 @@ def _fix_blob_seq_ids(palace_path: str) -> None: logger.exception("Could not write migration marker %s", marker) +def _fix_missing_collection_type(palace_path: str) -> None: + """Add ``_type`` to ``collections.config_json_str`` where absent. + + chromadb <= 1.5.8 writes ``config_json_str = '{}'`` (empty JSON) when + creating collections. chromadb 1.5.9 switched from the permissive + ``load_collection_configuration_from_json_str`` to + ``CollectionConfigurationInternal.from_json`` which requires a ``_type`` + key — its absence raises ``KeyError: '_type'`` on palace open. + + This migration adds the missing marker so both old and new chromadb + versions can load the collection. The value + ``"CollectionConfigurationInternal"`` matches what ``to_json()`` writes + for freshly-created collections. + + Same lifecycle constraints as :func:`_fix_blob_seq_ids`: must run + BEFORE ``PersistentClient`` is created. + """ + db_path = os.path.join(palace_path, "chroma.sqlite3") + if not os.path.isfile(db_path): + return + marker = os.path.join(palace_path, _COLLECTION_TYPE_MARKER) + if os.path.isfile(marker): + return + try: + with sqlite3.connect(db_path) as conn: + try: + rows = conn.execute("SELECT id, config_json_str FROM collections").fetchall() + except sqlite3.OperationalError: + return + updates = [] + for coll_id, config_str in rows: + if not config_str: + config_str = "{}" + try: + config = json.loads(config_str) + except (json.JSONDecodeError, TypeError): + continue + if not isinstance(config, dict): + continue + if "_type" not in config: + config["_type"] = "CollectionConfigurationInternal" + updates.append((json.dumps(config), coll_id)) + if updates: + conn.executemany( + "UPDATE collections SET config_json_str = ? WHERE id = ?", + updates, + ) + logger.info( + "Fixed %d collection(s) missing _type in config_json_str", + len(updates), + ) + conn.commit() + except Exception: + logger.exception("Could not fix collection config_json_str in %s", db_path) + return + try: + Path(marker).touch() + except OSError: + logger.exception("Could not write migration marker %s", marker) + + # --------------------------------------------------------------------------- # Collection adapter # --------------------------------------------------------------------------- @@ -1366,15 +1429,18 @@ class ChromaBackend(BaseBackend): """Run the pre-open safety pass shared by :meth:`make_client` and :meth:`_client`. - Three steps, all required before constructing a ``PersistentClient``: + Four steps, all required before constructing a ``PersistentClient``: - 1. ``_fix_blob_seq_ids`` — repairs the BLOB seq_id quirk that bites + 1. ``_fix_missing_collection_type`` — adds the ``_type`` marker to + ``collections.config_json_str`` that chromadb 1.5.9+ requires + but <= 1.5.8 never wrote (#1611). + 2. ``_fix_blob_seq_ids`` — repairs the BLOB seq_id quirk that bites certain chromadb migrations. - 2. ``quarantine_invalid_hnsw_metadata`` — renames aside any HNSW + 3. ``quarantine_invalid_hnsw_metadata`` — renames aside any HNSW ``index_metadata.pickle`` that fails to load, so chromadb opens against an empty index instead of crashing on the unloadable pickle (#1266 / PR #1285). - 3. ``quarantine_stale_hnsw`` — also gated by :attr:`_quarantined_paths` + 4. ``quarantine_stale_hnsw`` — also gated by :attr:`_quarantined_paths` so it fires once per palace per process. This is the SIGSEGV prevention path for stale HNSW segments (see #1121, #1132, #1263); wiring it through this helper means CLI mining, search, repair, @@ -1385,6 +1451,7 @@ class ChromaBackend(BaseBackend): re-open a palace. The ``_quarantined_paths`` gate prevents thrash on hot paths (e.g. ``_client()`` is called on every backend operation). """ + _fix_missing_collection_type(palace_path) _fix_blob_seq_ids(palace_path) if palace_path not in ChromaBackend._quarantined_paths: quarantine_invalid_hnsw_metadata(palace_path) @@ -1393,7 +1460,7 @@ class ChromaBackend(BaseBackend): @staticmethod def make_client(palace_path: str): - """Create a fresh ``PersistentClient`` (fixes BLOB seq_ids first). + """Create a fresh ``PersistentClient`` (runs pre-open safety pass first). Deprecated-ish: exposed for legacy long-lived callers that manage their own client cache. New code should obtain a collection through diff --git a/tests/test_backends.py b/tests/test_backends.py index e0c3b2b..95ad69f 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -23,6 +23,7 @@ from mempalace.backends.chroma import ( ChromaCollection, _HNSW_MISSING_METADATA_DATA_FLOOR, _fix_blob_seq_ids, + _fix_missing_collection_type, _pin_hnsw_threads, _segment_appears_healthy, quarantine_invalid_hnsw_metadata, @@ -614,6 +615,186 @@ def test_fix_blob_seq_ids_skips_sqlite_when_marker_present(tmp_path): mock_connect.assert_not_called() +# ── _fix_missing_collection_type ───────────────────────────────────────── + + +def test_fix_collection_type_adds_type(tmp_path): + """Legacy config_json_str '{}' gets _type added.""" + import json + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute( + "INSERT INTO collections (id, config_json_str) VALUES (?, ?)", + ("col-1", "{}"), + ) + conn.commit() + + _fix_missing_collection_type(str(tmp_path)) + + with closing(sqlite3.connect(str(db_path))) as conn: + row = conn.execute("SELECT config_json_str FROM collections WHERE id = 'col-1'").fetchone() + config = json.loads(row[0]) + assert config["_type"] == "CollectionConfigurationInternal" + + +def test_fix_collection_type_preserves_existing(tmp_path): + """Config that already has _type is left unchanged.""" + import json + + original = json.dumps({"_type": "CollectionConfigurationInternal", "extra": 1}) + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute( + "INSERT INTO collections (id, config_json_str) VALUES (?, ?)", + ("col-1", original), + ) + conn.commit() + + _fix_missing_collection_type(str(tmp_path)) + + with closing(sqlite3.connect(str(db_path))) as conn: + row = conn.execute("SELECT config_json_str FROM collections WHERE id = 'col-1'").fetchone() + assert row[0] == original + + +def test_fix_collection_type_noop_without_db(tmp_path): + """No error when palace has no chroma.sqlite3, no marker written.""" + from mempalace.backends.chroma import _COLLECTION_TYPE_MARKER + + _fix_missing_collection_type(str(tmp_path)) + assert not (tmp_path / _COLLECTION_TYPE_MARKER).exists() + + +def test_fix_collection_type_writes_marker(tmp_path): + """Marker is written after a successful migration.""" + from mempalace.backends.chroma import _COLLECTION_TYPE_MARKER + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute( + "INSERT INTO collections (id, config_json_str) VALUES (?, ?)", + ("col-1", "{}"), + ) + conn.commit() + + marker = tmp_path / _COLLECTION_TYPE_MARKER + assert not marker.exists() + + _fix_missing_collection_type(str(tmp_path)) + + assert marker.is_file() + + +def test_fix_collection_type_skips_with_marker(tmp_path): + """When the marker exists, sqlite3 is not opened.""" + from unittest.mock import patch + + from mempalace.backends.chroma import _COLLECTION_TYPE_MARKER + + db_path = tmp_path / "chroma.sqlite3" + db_path.write_bytes(b"sentinel") + (tmp_path / _COLLECTION_TYPE_MARKER).touch() + + with patch("mempalace.backends.chroma.sqlite3.connect") as mock_connect: + _fix_missing_collection_type(str(tmp_path)) + + mock_connect.assert_not_called() + + +def test_fix_collection_type_writes_marker_when_already_has_type(tmp_path): + """Marker written even when all collections already have _type (noop case).""" + import json + + from mempalace.backends.chroma import _COLLECTION_TYPE_MARKER + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute( + "INSERT INTO collections (id, config_json_str) VALUES (?, ?)", + ("col-1", json.dumps({"_type": "CollectionConfigurationInternal"})), + ) + conn.commit() + + marker = tmp_path / _COLLECTION_TYPE_MARKER + assert not marker.exists() + + _fix_missing_collection_type(str(tmp_path)) + + assert marker.is_file(), "marker must be written even when no collections needed fixing" + + +def test_fix_collection_type_multi_collection_mixed(tmp_path): + """Multiple collections: NULL, empty, and already-valid configs.""" + import json + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-null", None)) + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-empty", "{}")) + conn.execute( + "INSERT INTO collections VALUES (?, ?)", + ("col-ok", json.dumps({"_type": "CollectionConfigurationInternal"})), + ) + conn.commit() + + _fix_missing_collection_type(str(tmp_path)) + + with closing(sqlite3.connect(str(db_path))) as conn: + rows = { + r[0]: json.loads(r[1]) if r[1] else None + for r in conn.execute("SELECT id, config_json_str FROM collections") + } + assert rows["col-null"]["_type"] == "CollectionConfigurationInternal" + assert rows["col-empty"]["_type"] == "CollectionConfigurationInternal" + assert rows["col-ok"] == {"_type": "CollectionConfigurationInternal"} + + +def test_fix_collection_type_skips_non_dict_json(tmp_path): + """Non-dict JSON (array, null literal) is skipped without error.""" + import json + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-arr", "[]")) + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-null", "null")) + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-ok", "{}")) + conn.commit() + + _fix_missing_collection_type(str(tmp_path)) + + with closing(sqlite3.connect(str(db_path))) as conn: + rows = dict(conn.execute("SELECT id, config_json_str FROM collections").fetchall()) + assert rows["col-arr"] == "[]" + assert rows["col-null"] == "null" + assert json.loads(rows["col-ok"])["_type"] == "CollectionConfigurationInternal" + + +def test_fix_collection_type_skips_malformed_json(tmp_path): + """Malformed JSON in one row does not prevent fixing other rows.""" + import json + + db_path = tmp_path / "chroma.sqlite3" + with closing(sqlite3.connect(str(db_path))) as conn: + conn.execute("CREATE TABLE collections (id TEXT PRIMARY KEY, config_json_str TEXT)") + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-bad", "{corrupt")) + conn.execute("INSERT INTO collections VALUES (?, ?)", ("col-ok", "{}")) + conn.commit() + + _fix_missing_collection_type(str(tmp_path)) + + with closing(sqlite3.connect(str(db_path))) as conn: + rows = dict(conn.execute("SELECT id, config_json_str FROM collections").fetchall()) + assert rows["col-bad"] == "{corrupt" + assert json.loads(rows["col-ok"])["_type"] == "CollectionConfigurationInternal" + + # ── quarantine_stale_hnsw ───────────────────────────────────────────────── @@ -1218,6 +1399,9 @@ def test_chroma_backend_preflights_metadata_before_persistent_client(tmp_path, m return inner + monkeypatch.setattr( + "mempalace.backends.chroma._fix_missing_collection_type", _record("collection_type") + ) monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) monkeypatch.setattr( "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") @@ -1235,6 +1419,7 @@ def test_chroma_backend_preflights_metadata_before_persistent_client(tmp_path, m backend._client(str(palace)) assert calls == [ + ("collection_type", str(palace)), ("blob", str(palace)), ("invalid", str(palace)), ("stale", str(palace)), @@ -1255,6 +1440,9 @@ def test_chroma_backend_stale_quarantine_is_cold_start_only_on_refresh(tmp_path, return inner monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + monkeypatch.setattr( + "mempalace.backends.chroma._fix_missing_collection_type", _record("collection_type") + ) monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) monkeypatch.setattr( "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") @@ -1276,9 +1464,11 @@ def test_chroma_backend_stale_quarantine_is_cold_start_only_on_refresh(tmp_path, backend._client(str(palace)) assert calls == [ + ("collection_type", str(palace)), ("blob", str(palace)), ("invalid", str(palace)), ("stale", str(palace)), + ("collection_type", str(palace)), ("blob", str(palace)), ] @@ -1297,6 +1487,9 @@ def test_chroma_backend_requarantines_after_inode_replacement(tmp_path, monkeypa return inner monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + monkeypatch.setattr( + "mempalace.backends.chroma._fix_missing_collection_type", _record("collection_type") + ) monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) monkeypatch.setattr( "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") @@ -1318,9 +1511,11 @@ def test_chroma_backend_requarantines_after_inode_replacement(tmp_path, monkeypa backend._client(str(palace)) assert calls == [ + ("collection_type", str(palace)), ("blob", str(palace)), ("invalid", str(palace)), ("stale", str(palace)), + ("collection_type", str(palace)), ("blob", str(palace)), ("invalid", str(palace)), ("stale", str(palace)),