fix: auto-repair BLOB seq_ids from chromadb 0.6→1.5 migration (#664)
Note from code review: (1) silent exception swallow on migration failure means caller proceeds with potentially corrupt DB — consider returning a boolean or re-raising in a follow-up. (2) No blob length validation before int.from_bytes — malformed rows could produce wrong seq_id values. Both are edge cases; the fix is still valuable for the common chromadb 0.6→1.5 migration path.
This commit is contained in:
+49
-1
@@ -1,7 +1,9 @@
|
||||
import sqlite3
|
||||
|
||||
import chromadb
|
||||
import pytest
|
||||
|
||||
from mempalace.backends.chroma import ChromaBackend, ChromaCollection
|
||||
from mempalace.backends.chroma import ChromaBackend, ChromaCollection, _fix_blob_seq_ids
|
||||
|
||||
|
||||
class _FakeCollection:
|
||||
@@ -78,3 +80,49 @@ def test_chroma_backend_create_true_creates_directory_and_collection(tmp_path):
|
||||
|
||||
client = chromadb.PersistentClient(path=str(palace_path))
|
||||
client.get_collection("mempalace_drawers")
|
||||
|
||||
|
||||
def test_fix_blob_seq_ids_converts_blobs_to_integers(tmp_path):
|
||||
"""Simulate a ChromaDB 0.6.x database with BLOB seq_ids and verify repair."""
|
||||
db_path = tmp_path / "chroma.sqlite3"
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.execute("CREATE TABLE embeddings (rowid INTEGER PRIMARY KEY, seq_id)")
|
||||
conn.execute("CREATE TABLE max_seq_id (rowid INTEGER PRIMARY KEY, seq_id)")
|
||||
# Insert BLOB seq_ids like ChromaDB 0.6.x would
|
||||
blob_42 = (42).to_bytes(8, byteorder="big")
|
||||
blob_99 = (99).to_bytes(8, byteorder="big")
|
||||
conn.execute("INSERT INTO embeddings (seq_id) VALUES (?)", (blob_42,))
|
||||
conn.execute("INSERT INTO max_seq_id (seq_id) VALUES (?)", (blob_99,))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
_fix_blob_seq_ids(str(tmp_path))
|
||||
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
row = conn.execute("SELECT seq_id, typeof(seq_id) FROM embeddings").fetchone()
|
||||
assert row == (42, "integer")
|
||||
row = conn.execute("SELECT seq_id, typeof(seq_id) FROM max_seq_id").fetchone()
|
||||
assert row == (99, "integer")
|
||||
conn.close()
|
||||
|
||||
|
||||
def test_fix_blob_seq_ids_noop_without_blobs(tmp_path):
|
||||
"""No error when seq_ids are already integers."""
|
||||
db_path = tmp_path / "chroma.sqlite3"
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.execute("CREATE TABLE embeddings (rowid INTEGER PRIMARY KEY, seq_id INTEGER)")
|
||||
conn.execute("INSERT INTO embeddings (seq_id) VALUES (42)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
_fix_blob_seq_ids(str(tmp_path))
|
||||
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
row = conn.execute("SELECT seq_id, typeof(seq_id) FROM embeddings").fetchone()
|
||||
assert row == (42, "integer")
|
||||
conn.close()
|
||||
|
||||
|
||||
def test_fix_blob_seq_ids_noop_without_database(tmp_path):
|
||||
"""No error when palace has no chroma.sqlite3."""
|
||||
_fix_blob_seq_ids(str(tmp_path)) # should not raise
|
||||
|
||||
Reference in New Issue
Block a user