feat: native rename_wing backend operation + CLI command (#154) (#154)

jphein · claude · web-flow · commit d045f83e50df · 2026-05-23T11:40:25.000-07:00
The MCP rename_wing tool was unusably slow on postgres — it inherited
the base update() which re-embeds every document via upsert(). A
metadata-only wing rename should never re-embed.

Changes:
- Add rename_wing() default to BaseCollection (get/update loop)
- Override in PostgresCollection with single SQL UPDATE (atomic, no
  re-embedding, milliseconds instead of minutes for 20K+ drawers)
- Add update() override in PostgresCollection for metadata-only
  updates (skips re-embedding when no documents/embeddings change)
- Simplify MCP tool_rename_wing() to delegate to col.rename_wing()
- Add CLI `mempalace rename-wing --from X --to Y` subcommand with
  daemon routing and --dry-run support
- Tests for backend, MCP tool, and CLI (7 new tests)

Co-authored-by: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/mempalace/backends/base.py b/mempalace/backends/base.py
@@ -299,6 +299,42 @@ def update(
             embeddings=embeddings,
         )
 
+    def rename_wing(
+        self, *, from_wing: str, to_wing: str, batch_size: int = 500
+    ) -> dict:
+        """Rename all drawers from one wing to another.
+
+        Default implementation iterates in batches using metadata-only
+        ``update()`` calls.  Backends with native bulk-update support
+        (e.g. PostgreSQL) should override with an atomic implementation.
+
+        Returns ``{"renamed": int, "errors": int}``.
+        """
+        renamed = 0
+        errors = 0
+        while True:
+            batch = self.get(
+                where={"wing": from_wing},
+                include=["metadatas"],
+                limit=batch_size,
+                offset=0,
+            )
+            if not batch.ids:
+                break
+            new_metas = []
+            for meta in batch.metadatas:
+                m = dict(meta)
+                m["wing"] = to_wing
+                new_metas.append(m)
+            try:
+                self.update(ids=batch.ids, metadatas=new_metas)
+                renamed += len(batch.ids)
+            except Exception:
+                errors += len(batch.ids)
+            if len(batch.ids) < batch_size:
+                break
+        return {"renamed": renamed, "errors": errors}
+
 
 # ---------------------------------------------------------------------------
 # Backend contract
diff --git a/mempalace/backends/postgres.py b/mempalace/backends/postgres.py
@@ -549,6 +549,64 @@ def delete(
             params,
         )
 
+    def update(
+        self,
+        *,
+        ids: list[str],
+        documents: Optional[list[str]] = None,
+        metadatas: Optional[list[dict]] = None,
+        embeddings: Optional[list[list[float]]] = None,
+    ) -> None:
+        if documents is None and metadatas is None and embeddings is None:
+            raise ValueError("update requires at least one of documents, metadatas, embeddings")
+        if documents is not None or embeddings is not None:
+            super().update(
+                ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings,
+            )
+            return
+
+        n = len(ids)
+        if metadatas is not None and len(metadatas) != n:
+            raise ValueError(f"metadatas length {len(metadatas)} does not match ids length {n}")
+        self._ensure_setup(create=True)
+
+        cur = self._get_conn().cursor()
+        for i, doc_id in enumerate(ids):
+            meta = dict(metadatas[i]) if metadatas else {}
+            raw_wing = meta.pop("wing", None)
+            raw_room = meta.pop("room", None)
+            set_parts = []
+            params: list[Any] = []
+            if raw_wing is not None:
+                set_parts.append(self._sql.SQL("wing = %s"))
+                params.append(_metadata_value(raw_wing))
+            if raw_room is not None:
+                set_parts.append(self._sql.SQL("room = %s"))
+                params.append(_metadata_value(raw_room))
+            set_parts.append(self._sql.SQL("metadata = metadata || %s::jsonb"))
+            params.append(json.dumps(meta))
+            params.append(doc_id)
+            cur.execute(
+                self._sql.SQL("UPDATE {} SET {} WHERE id = %s").format(
+                    self._table_id,
+                    self._sql.SQL(", ").join(set_parts),
+                ),
+                params,
+            )
+
+    def rename_wing(
+        self, *, from_wing: str, to_wing: str, batch_size: int = 500
+    ) -> dict:
+        self._ensure_setup(create=True)
+        cur = self._get_conn().cursor()
+        cur.execute(
+            self._sql.SQL("UPDATE {} SET wing = %s WHERE wing = %s").format(
+                self._table_id,
+            ),
+            [to_wing, from_wing],
+        )
+        return {"renamed": cur.rowcount, "errors": 0}
+
     def count(self) -> int:
         self._ensure_setup(create=True)
         cur = self._get_conn().cursor()
diff --git a/mempalace/cli.py b/mempalace/cli.py
@@ -1422,6 +1422,91 @@ def cmd_purge(args):
     print(f"\n  Purged {match_count:,} drawers. Remaining: {remaining:,}\n")
 
 
+def cmd_rename_wing(args):
+    want_json = getattr(args, "json", False)
+    from_wing = args.from_wing
+    to_wing = args.to_wing
+    dry_run = getattr(args, "dry_run", False)
+    batch_size = getattr(args, "batch_size", 500)
+
+    if _daemon_strict():
+        if dry_run:
+            try:
+                data = _call_daemon_tool("mempalace_list_drawers", {
+                    "wing": from_wing, "limit": 1,
+                })
+            except DaemonError as e:
+                if want_json:
+                    _emit_json({"error": str(e)})
+                else:
+                    print(f"\n  ERROR: {e}", file=sys.stderr)
+                sys.exit(2)
+            total = data.get("total", 0)
+            if want_json:
+                _emit_json({"dry_run": True, "from_wing": from_wing, "to_wing": to_wing, "count": total})
+            else:
+                print(f"\n  Dry run: {total:,} drawers would be renamed from '{from_wing}' to '{to_wing}'\n")
+            return
+
+        try:
+            data = _call_daemon_tool("mempalace_rename_wing", {
+                "from_wing": from_wing,
+                "to_wing": to_wing,
+                "batch_size": batch_size,
+            })
+        except DaemonError as e:
+            if want_json:
+                _emit_json({"error": str(e)})
+            else:
+                print(f"\n  ERROR: {e}", file=sys.stderr)
+            sys.exit(2)
+
+        if want_json:
+            _emit_json(data)
+        else:
+            renamed = data.get("renamed", 0)
+            errors = data.get("errors", 0)
+            print(f"\n  Renamed {renamed:,} drawers: '{from_wing}' -> '{to_wing}'")
+            if errors:
+                print(f"  Errors: {errors:,}")
+            print()
+        return
+
+    from .backends.chroma import ChromaBackend
+    from .backends.base import PalaceRef
+
+    palace_path = os.path.abspath(
+        os.path.expanduser(args.palace) if getattr(args, "palace", None) else MempalaceConfig().palace_path
+    )
+    backend = ChromaBackend()
+    try:
+        col = backend.get_collection(
+            palace=PalaceRef(id=palace_path, local_path=palace_path),
+            collection_name="mempalace_drawers",
+        )
+    except Exception as e:
+        print(f"\n  Error reading palace: {e}")
+        sys.exit(1)
+
+    if dry_run:
+        matched = col.get(where={"wing": from_wing}, include=[])
+        count = len(matched.ids) if hasattr(matched, "ids") else len(matched.get("ids", []))
+        if want_json:
+            _emit_json({"dry_run": True, "from_wing": from_wing, "to_wing": to_wing, "count": count})
+        else:
+            print(f"\n  Dry run: {count:,} drawers would be renamed from '{from_wing}' to '{to_wing}'\n")
+        return
+
+    result = col.rename_wing(from_wing=from_wing, to_wing=to_wing, batch_size=batch_size)
+    if want_json:
+        _emit_json({"success": True, "from_wing": from_wing, "to_wing": to_wing, **result})
+    else:
+        print(f"\n  Renamed {result['renamed']:,} drawers: '{from_wing}' -> '{to_wing}'")
+        if result["errors"]:
+            print(f"  Errors: {result['errors']:,}")
+        print()
+
+
 def cmd_replay(args):
     """Drain ``~/.mempalace/pending/*.jsonl`` by re-issuing each request to the daemon.
 
@@ -2622,6 +2707,15 @@ def main():
     )
     p_purge.add_argument("--yes", "-y", action="store_true", help="Skip confirmation prompt")
 
+    p_rename_wing = sub.add_parser(
+        "rename-wing",
+        help="Rename all drawers from one wing to another (atomic on postgres)",
+    )
+    p_rename_wing.add_argument("--from", dest="from_wing", required=True, help="Source wing name")
+    p_rename_wing.add_argument("--to", dest="to_wing", required=True, help="Target wing name")
+    p_rename_wing.add_argument("--dry-run", action="store_true", help="Count matching drawers without renaming")
+    p_rename_wing.add_argument("--batch-size", type=int, default=500, help="Batch size for non-postgres backends (default: 500)")
+
     # ── rooms — manage the canonical room set (hybrid-search-taxonomy follow-up) ────
     p_rooms = sub.add_parser(
         "rooms",
@@ -2761,6 +2855,7 @@ def _nonneg_int(value: str) -> int:
         "migrate": cmd_migrate,
         "migrate-to-postgres": cmd_migrate_to_postgres,
         "purge": cmd_purge,
+        "rename-wing": cmd_rename_wing,
         "rooms": cmd_rooms,
         "status": cmd_status,
         "mined": cmd_mined,
diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py
@@ -2097,61 +2097,26 @@ def tool_rename_wing(from_wing: str, to_wing: str, batch_size: int = 500):
         return _no_palace()
 
     try:
-        total_result = col.get(where={"wing": from_wing}, include=[])
-        total = len(total_result["ids"])
-        if total == 0:
-            return {"success": True, "renamed": 0, "message": f"no drawers in wing '{from_wing}'"}
-
-        renamed = 0
-        errors = 0
-
-        while True:
-            batch = col.get(
-                where={"wing": from_wing},
-                include=["metadatas"],
-                limit=batch_size,
-                offset=0,
-            )
-            if not batch["ids"]:
-                break
-
-            update_ids = []
-            update_metas = []
-            for i, did in enumerate(batch["ids"]):
-                meta = dict(_safe_meta(batch["metadatas"][i]))
-                meta["wing"] = to_wing
-                update_ids.append(did)
-                update_metas.append(meta)
-
-            try:
-                col.update(ids=update_ids, metadatas=update_metas)
-                renamed += len(update_ids)
-            except Exception:
-                errors += len(update_ids)
-
-            if len(batch["ids"]) < batch_size:
-                break
-
+        result = col.rename_wing(
+            from_wing=from_wing, to_wing=to_wing, batch_size=batch_size,
+        )
         _metadata_cache = None
 
         _wal_log(
             "rename_wing",
             {
                 "from_wing": from_wing,
                 "to_wing": to_wing,
-                "renamed": renamed,
-                "errors": errors,
+                **result,
             },
         )
 
-        logger.info(f"Renamed wing: {from_wing} -> {to_wing} ({renamed} drawers)")
+        logger.info(f"Renamed wing: {from_wing} -> {to_wing} ({result['renamed']} drawers)")
         return {
             "success": True,
             "from_wing": from_wing,
             "to_wing": to_wing,
-            "renamed": renamed,
-            "errors": errors,
-            "total": total,
+            **result,
         }
     except Exception as e:
         return {"success": False, "error": str(e)}
diff --git a/tests/test_backends.py b/tests/test_backends.py
@@ -782,6 +782,56 @@ def test_base_collection_update_default_rejects_mismatched_lengths():
         BaseCollection.update(collection, ids=["1", "2"], metadatas=[{"k": 9}])
 
 
+def test_base_collection_rename_wing_default(tmp_path):
+    """The ABC default rename_wing() moves drawers between wings via update()."""
+    palace_path = tmp_path / "palace"
+    backend = ChromaBackend()
+    col = backend.get_collection(
+        palace=PalaceRef(id=str(palace_path), local_path=str(palace_path)),
+        collection_name="mempalace_drawers",
+        create=True,
+    )
+    col.add(
+        ids=["d1", "d2", "d3"],
+        documents=["doc one", "doc two", "doc three"],
+        metadatas=[
+            {"wing": "old_wing", "room": "r1"},
+            {"wing": "old_wing", "room": "r2"},
+            {"wing": "other_wing", "room": "r1"},
+        ],
+    )
+
+    result = col.rename_wing(from_wing="old_wing", to_wing="new_wing")
+
+    assert result["renamed"] == 2
+    assert result["errors"] == 0
+
+    remaining_old = col.get(where={"wing": "old_wing"}, include=["metadatas"])
+    assert len(remaining_old.ids) == 0
+
+    renamed = col.get(where={"wing": "new_wing"}, include=["metadatas"])
+    assert len(renamed.ids) == 2
+    assert set(renamed.ids) == {"d1", "d2"}
+
+    untouched = col.get(where={"wing": "other_wing"}, include=["metadatas"])
+    assert len(untouched.ids) == 1
+    assert untouched.ids[0] == "d3"
+
+
+def test_base_collection_rename_wing_empty_source(tmp_path):
+    """rename_wing() with no matching drawers returns zero renamed, no errors."""
+    palace_path = tmp_path / "palace"
+    backend = ChromaBackend()
+    col = backend.get_collection(
+        palace=PalaceRef(id=str(palace_path), local_path=str(palace_path)),
+        collection_name="mempalace_drawers",
+        create=True,
+    )
+    result = col.rename_wing(from_wing="nonexistent", to_wing="target")
+    assert result["renamed"] == 0
+    assert result["errors"] == 0
+
+
 def test_chroma_backend_accepts_palace_ref_kwarg(tmp_path):
     palace_path = tmp_path / "palace"
     backend = ChromaBackend()
diff --git a/tests/test_cli.py b/tests/test_cli.py
diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py
