refactor(cli): collection.delete(where=) instead of nuke-and-rebuild

jphein · claude · jphein · commit e9a59de75bb1 · 2026-04-26T09:50:24.000-07:00
Addresses @igorls's review on MemPalace#1087: 1. **Premise was wrong (MemPalace#521 doesn't apply to delete-by-where).** The original implementation extracted survivors, rmtree'd the palace, and re-inserted into a fresh collection — predicated on the idea that collection.delete(where=...) would trigger the same updatePoint / repairConnectionsForUpdate race as the upsert path in MemPalace#521. It doesn't: chromadb's filter-delete bypasses the HNSW update codepath that races. The simpler approach is correct. 2. **Embedding function preserved.** No more bypassing ChromaBackend.create_collection / _resolve_embedding_function — we never recreate the collection at all. 3. **No data loss on interrupt.** No rmtree means no window where a crash leaves the palace empty. 4. **Routes through the backend.** Uses ChromaBackend.get_collection instead of `import chromadb; PersistentClient(...)` directly. No more `del col, client` reliance on refcount finalization. 5. **Reuses confirm_destructive_action.** Replaces the custom `input(...)` prompt with the migrate.py helper. End-to-end test added: real chromadb palace, real backend, real delete, verify count + surviving ids. Catches the embedding-function regression that was the load-bearing concern in the review. Closes MemPalace#848 narrowly (wing/room delete). Source-file / query / dry-run modes from MemPalace#848's broader ask are deliberately NOT in scope here — filing as follow-ups if there's appetite. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
diff --git a/mempalace/cli.py b/mempalace/cli.py
@@ -287,30 +287,29 @@ def cmd_migrate(args):
 def cmd_purge(args):
     """Delete drawers by wing and/or room.
 
-    Extracts the drawers to *keep*, nukes the palace directory, and
-    re-inserts them into a fresh ChromaDB. This avoids HNSW ghost entries
-    that ChromaDB's in-place ``collection.delete()`` leaves behind, which
-    can cause segfaults on subsequent queries or inserts (see #521 for the
-    underlying hnswlib ``updatePoint`` race on modified-file re-mines).
-
-    Note: ``--room`` without ``--wing`` purges that room across ALL wings.
-
-    Not idempotent — running purge twice on the same criteria will print
-    "No drawers found" the second time. The operation is destructive; the
-    ``--yes`` flag skips the interactive confirmation.
+    Uses ``collection.delete(where=...)`` — chromadb's filter-delete path
+    doesn't go through ``updatePoint`` / ``repairConnectionsForUpdate``,
+    which is the upsert-only race from #521 that an earlier draft of this
+    command tried to side-step with a nuke-and-rebuild. The simpler path
+    works without losing drawers if the process is interrupted, without
+    re-embedding the survivors under a default model, and without
+    bypassing the backend abstraction.
+
+    ``--room`` without ``--wing`` purges that room across ALL wings.
+    Not idempotent — running purge twice on the same criteria prints
+    "No drawers found" the second time.
     """
-    import chromadb
-    import shutil
+    from .backends.chroma import ChromaBackend
+    from .migrate import confirm_destructive_action, contains_palace_database
 
-    palace_path = os.path.expanduser(args.palace) if args.palace else MempalaceConfig().palace_path
-    try:
-        client = chromadb.PersistentClient(path=palace_path)
-        col = client.get_collection("mempalace_drawers")
-    except Exception:
+    palace_path = os.path.abspath(
+        os.path.expanduser(args.palace) if args.palace else MempalaceConfig().palace_path
+    )
+
+    if not os.path.isdir(palace_path) or not contains_palace_database(palace_path):
         print(f"\n  No palace found at {palace_path}")
         return
 
-    where = {}
     if args.wing and args.room:
         where = {"$and": [{"wing": args.wing}, {"room": args.room}]}
     elif args.wing:
@@ -321,80 +320,51 @@ def cmd_purge(args):
         print("  Error: specify --wing and/or --room")
         return
 
-    total = col.count()
+    backend = ChromaBackend()
+    try:
+        col = backend.get_collection(palace_path, "mempalace_drawers")
+    except Exception as e:
+        print(f"\n  Error reading palace: {e}")
+        return
 
-    # Count matching drawers
-    match_ids = set()
-    offset = 0
-    while True:
-        batch = col.get(limit=10000, offset=offset, where=where, include=[])
-        if not batch["ids"]:
-            break
-        match_ids.update(batch["ids"])
-        offset += len(batch["ids"])
+    # Probe match count via a `where=`-filtered get with no payload.
+    # ChromaCollection.get returns the typed result; we only need ids.
+    try:
+        matched = col.get(where=where, include=[])
+    except Exception as e:
+        print(f"\n  Error querying drawers: {e}")
+        return
+
+    match_ids = matched.get("ids") if isinstance(matched, dict) else getattr(matched, "ids", [])
+    match_ids = match_ids or []
+    match_count = len(match_ids)
 
-    if not match_ids:
-        label = f"wing={args.wing}" if args.wing else ""
-        if args.room:
-            label = f"{label} room={args.room}" if label else f"room={args.room}"
+    label_parts = []
+    if args.wing:
+        label_parts.append(f"wing={args.wing}")
+    if args.room:
+        label_parts.append(f"room={args.room}")
+    label = " ".join(label_parts)
+
+    if match_count == 0:
         print(f"\n  No drawers found matching {label}\n")
         return
 
-    label = f"wing={args.wing}" if args.wing else ""
-    if args.room:
-        label = f"{label} room={args.room}" if label else f"room={args.room}"
-    keep_count = total - len(match_ids)
-    print(f"\n  Found {len(match_ids):,} drawers matching {label}")
-    print(f"  Will keep {keep_count:,} drawers, rebuild index")
+    print(f"\n  Found {match_count:,} drawers matching {label}")
 
     if not args.yes:
-        confirm = input(f"  Purge {len(match_ids):,} drawers? [y/N] ").strip().lower()
-        if confirm not in ("y", "yes"):
-            print("  Aborted.")
+        if not confirm_destructive_action(f"Purge of {match_count:,} drawers", palace_path):
             return
 
-    # Extract drawers to keep (everything NOT matching the filter)
-    print("  Extracting drawers to keep...")
-    keep_ids, keep_docs, keep_metas = [], [], []
-    offset = 0
-    batch_size = 5000
-    while offset < total:
-        batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"])
-        if not batch["ids"]:
-            break
-        for i, doc_id in enumerate(batch["ids"]):
-            if doc_id not in match_ids:
-                keep_ids.append(doc_id)
-                keep_docs.append(batch["documents"][i])
-                keep_metas.append(batch["metadatas"][i])
-        offset += len(batch["ids"])
-    print(f"  Extracted {len(keep_ids):,} drawers to keep")
-
-    # Release client before nuking — ChromaDB holds open file handles
-    # (WAL journal, HNSW mmap) that block rmtree on Windows and some Linux FS.
-    del col, client
-
-    # Nuke and rebuild with clean HNSW index
-    palace_path = palace_path.rstrip(os.sep)
-    print("  Rebuilding palace...")
-    shutil.rmtree(palace_path)
-    os.makedirs(palace_path, mode=0o700)
-
-    new_client = chromadb.PersistentClient(path=palace_path)
-    new_col = new_client.create_collection("mempalace_drawers", metadata={"hnsw:space": "cosine"})
-
-    filed = 0
-    for i in range(0, len(keep_ids), batch_size):
-        end = min(i + batch_size, len(keep_ids))
-        new_col.add(
-            documents=keep_docs[i:end],
-            ids=keep_ids[i:end],
-            metadatas=keep_metas[i:end],
-        )
-        filed += end - i
-        print(f"  Re-filed {filed:,} / {len(keep_ids):,}...", flush=True)
+    print("  Deleting matching drawers...")
+    try:
+        col.delete(where=where)
+    except Exception as e:
+        print(f"\n  Delete failed: {e}\n")
+        return
 
-    print(f"\n  Purged {len(match_ids):,} drawers. Remaining: {new_col.count():,}\n")
+    remaining = col.count()
+    print(f"\n  Purged {match_count:,} drawers. Remaining: {remaining:,}\n")
 
 
 def cmd_status(args):
@@ -888,7 +858,7 @@ def main():
     # purge
     p_purge = sub.add_parser(
         "purge",
-        help="Delete drawers by wing and/or room (nuke-and-reinsert to avoid HNSW ghost entries)",
+        help="Delete drawers by wing and/or room (filtered delete via chromadb)",
     )
     p_purge.add_argument("--wing", help="Wing to purge")
     p_purge.add_argument("--room", help="Room to purge (across all wings if --wing omitted)")
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -59,74 +59,106 @@ def _make_purge_args(**overrides):
 
 
 @patch("mempalace.cli.MempalaceConfig")
-def test_cmd_purge_no_palace_found(mock_config_cls, capsys):
+def test_cmd_purge_no_palace_found(mock_config_cls, capsys, tmp_path):
     """Purge prints a clear message when the palace doesn't exist."""
-    mock_config_cls.return_value.palace_path = "/nonexistent/palace"
-    args = _make_purge_args(wing="any")
-    mock_chromadb = MagicMock()
-    mock_chromadb.PersistentClient.side_effect = Exception("no palace")
-    with patch.dict("sys.modules", {"chromadb": mock_chromadb, "shutil": MagicMock()}):
-        cmd_purge(args)
+    missing = tmp_path / "nonexistent"
+    mock_config_cls.return_value.palace_path = str(missing)
+    args = _make_purge_args(wing="any", palace=str(missing))
+    cmd_purge(args)
     out = capsys.readouterr().out
     assert "No palace found" in out
 
 
 @patch("mempalace.cli.MempalaceConfig")
-def test_cmd_purge_requires_filter(mock_config_cls, capsys):
+def test_cmd_purge_requires_filter(mock_config_cls, capsys, tmp_path):
     """Purge refuses to run without --wing or --room (no mass-delete safety valve)."""
-    mock_config_cls.return_value.palace_path = "/fake/palace"
-    args = _make_purge_args()  # no wing, no room
-    mock_col = MagicMock()
-    mock_client = MagicMock()
-    mock_client.get_collection.return_value = mock_col
-    mock_chromadb = MagicMock()
-    mock_chromadb.PersistentClient.return_value = mock_client
-    with patch.dict("sys.modules", {"chromadb": mock_chromadb, "shutil": MagicMock()}):
-        cmd_purge(args)
+    palace = tmp_path / "palace"
+    palace.mkdir()
+    (palace / "chroma.sqlite3").write_text("")
+    mock_config_cls.return_value.palace_path = str(palace)
+    args = _make_purge_args(palace=str(palace))  # no wing, no room
+    cmd_purge(args)
     out = capsys.readouterr().out
     assert "Error: specify --wing and/or --room" in out
-    mock_col.count.assert_not_called()  # didn't touch data
 
 
 @patch("mempalace.cli.MempalaceConfig")
-def test_cmd_purge_no_matches(mock_config_cls, capsys):
-    """When the filter matches zero drawers, purge exits cleanly without rebuilding."""
-    mock_config_cls.return_value.palace_path = "/fake/palace"
-    args = _make_purge_args(wing="empty-wing")
+def test_cmd_purge_no_matches(mock_config_cls, capsys, tmp_path):
+    """When the filter matches zero drawers, purge exits cleanly."""
+    palace = tmp_path / "palace"
+    palace.mkdir()
+    (palace / "chroma.sqlite3").write_text("")
+    mock_config_cls.return_value.palace_path = str(palace)
+    args = _make_purge_args(wing="empty-wing", palace=str(palace))
+
     mock_col = MagicMock()
-    mock_col.count.return_value = 100
     mock_col.get.return_value = {"ids": []}
-    mock_client = MagicMock()
-    mock_client.get_collection.return_value = mock_col
-    mock_chromadb = MagicMock()
-    mock_chromadb.PersistentClient.return_value = mock_client
-    mock_shutil = MagicMock()
-    with patch.dict("sys.modules", {"chromadb": mock_chromadb, "shutil": mock_shutil}):
+    mock_backend = MagicMock()
+    mock_backend.return_value.get_collection.return_value = mock_col
+    with patch("mempalace.backends.chroma.ChromaBackend", mock_backend):
         cmd_purge(args)
     out = capsys.readouterr().out
     assert "No drawers found matching" in out
-    mock_shutil.rmtree.assert_not_called()  # nothing to rebuild
+    mock_col.delete.assert_not_called()
 
 
 @patch("mempalace.cli.MempalaceConfig")
-def test_cmd_purge_wing_and_room_uses_and_filter(mock_config_cls):
+def test_cmd_purge_wing_and_room_uses_and_filter(mock_config_cls, tmp_path):
     """Purge builds a $and filter when both --wing and --room are set."""
-    mock_config_cls.return_value.palace_path = "/fake/palace"
-    args = _make_purge_args(wing="myproj", room="drafts")
+    palace = tmp_path / "palace"
+    palace.mkdir()
+    (palace / "chroma.sqlite3").write_text("")
+    mock_config_cls.return_value.palace_path = str(palace)
+    args = _make_purge_args(wing="myproj", room="drafts", palace=str(palace))
+
     mock_col = MagicMock()
-    mock_col.count.return_value = 0
     mock_col.get.return_value = {"ids": []}
-    mock_client = MagicMock()
-    mock_client.get_collection.return_value = mock_col
-    mock_chromadb = MagicMock()
-    mock_chromadb.PersistentClient.return_value = mock_client
-    with patch.dict("sys.modules", {"chromadb": mock_chromadb, "shutil": MagicMock()}):
+    mock_backend = MagicMock()
+    mock_backend.return_value.get_collection.return_value = mock_col
+    with patch("mempalace.backends.chroma.ChromaBackend", mock_backend):
         cmd_purge(args)
-    # First col.get call builds the match-id set, and its `where` should be $and
     first_call = mock_col.get.call_args_list[0]
     assert first_call.kwargs["where"] == {"$and": [{"wing": "myproj"}, {"room": "drafts"}]}
 
 
+def test_cmd_purge_deletes_via_where_clause(tmp_path):
+    """End-to-end: purge filters via collection.delete(where=...) — preserves
+    embedding function, no rmtree, no rebuild. Regression for igorls' #1087
+    review concerns 1, 2, 3."""
+    palace = tmp_path / "palace"
+    palace.mkdir()
+    # Real chromadb palace via the backend so the embedding function is
+    # whatever the backend resolves (the actual concern from the review).
+    from mempalace.backends.chroma import ChromaBackend
+
+    backend = ChromaBackend()
+    col = backend.get_collection(str(palace), "mempalace_drawers", create=True)
+    col.add(
+        ids=["k1", "k2", "p1", "p2"],
+        documents=["keep one", "keep two", "purge one", "purge two"],
+        metadatas=[
+            {"wing": "keep", "room": "r"},
+            {"wing": "keep", "room": "r"},
+            {"wing": "purge-me", "room": "r"},
+            {"wing": "purge-me", "room": "r"},
+        ],
+    )
+    assert col.count() == 4
+
+    args = _make_purge_args(wing="purge-me", palace=str(palace))
+    with patch("mempalace.cli.MempalaceConfig") as mock_config_cls:
+        mock_config_cls.return_value.palace_path = str(palace)
+        cmd_purge(args)
+
+    # Re-open through the backend to confirm survivors and that the index
+    # still works (no nuke, embedding function intact).
+    col2 = backend.get_collection(str(palace), "mempalace_drawers", create=False)
+    assert col2.count() == 2
+    surviving = col2.get(include=["metadatas"])
+    surviving_ids = surviving.get("ids") if isinstance(surviving, dict) else surviving.ids
+    assert set(surviving_ids) == {"k1", "k2"}
+
+
 # ── cmd_search ─────────────────────────────────────────────────────────
 
 
