From 5dfe8531543f674022b30722209733c7b0ce309e Mon Sep 17 00:00:00 2001
From: JunghwanNA <70629228+shaun0927@users.noreply.github.com>
Date: Thu, 16 Apr 2026 12:11:18 +0900
Subject: [PATCH] fix: guard against data loss in repair, migrate, and CLI
 rebuild

- repair.py: wrap upsert loop in try/except; restore from backup on
  failure instead of leaving a partially rebuilt collection
- migrate.py: replace non-atomic rmtree+move with rename-aside swap
  so a crash between the two calls does not destroy both copies
- cli.py: use offset += len(batch["ids"]) with empty-batch guard
  instead of fixed offset += batch_size to prevent skipping drawers
---
 mempalace/cli.py     |  4 +++-
 mempalace/migrate.py | 15 ++++++++++++---
 mempalace/repair.py  | 26 +++++++++++++++++++-------
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/mempalace/cli.py b/mempalace/cli.py
index fb2f0ae..d185f61 100644
--- a/mempalace/cli.py
+++ b/mempalace/cli.py
@@ -270,10 +270,12 @@ def cmd_repair(args):
     offset = 0
     while offset < total:
         batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"])
+        if not batch["ids"]:
+            break
         all_ids.extend(batch["ids"])
         all_docs.extend(batch["documents"])
         all_metas.extend(batch["metadatas"])
-        offset += batch_size
+        offset += len(batch["ids"])
     print(f"  Extracted {len(all_ids)} drawers")
 
     # Backup and rebuild
diff --git a/mempalace/migrate.py b/mempalace/migrate.py
index 2eebb61..c804211 100644
--- a/mempalace/migrate.py
+++ b/mempalace/migrate.py
@@ -229,10 +229,19 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
     del col
     del fresh_backend
 
-    # Swap: remove old palace, move new one into place
+    # Swap: rename old palace aside, then move new one into place.
+    # This avoids a window where both old and new are missing.
     print("  Swapping old palace for migrated version...")
-    shutil.rmtree(palace_path)
-    shutil.move(temp_palace, palace_path)
+    stale_path = palace_path + ".old"
+    if os.path.exists(stale_path):
+        shutil.rmtree(stale_path)
+    os.rename(palace_path, stale_path)
+    try:
+        os.rename(temp_palace, palace_path)
+    except OSError:
+        # os.rename fails across filesystems; fall back to move
+        shutil.move(temp_palace, palace_path)
+    shutil.rmtree(stale_path, ignore_errors=True)
 
     print("\n  Migration complete.")
     print(f"  Drawers migrated: {final_count}")
diff --git a/mempalace/repair.py b/mempalace/repair.py
index 9a9aa88..1b4f271 100644
--- a/mempalace/repair.py
+++ b/mempalace/repair.py
@@ -266,13 +266,25 @@ def rebuild_index(palace_path=None):
     new_col = backend.create_collection(palace_path, COLLECTION_NAME)
 
     filed = 0
-    for i in range(0, len(all_ids), batch_size):
-        batch_ids = all_ids[i : i + batch_size]
-        batch_docs = all_docs[i : i + batch_size]
-        batch_metas = all_metas[i : i + batch_size]
-        new_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas)
-        filed += len(batch_ids)
-        print(f"  Re-filed {filed}/{len(all_ids)} drawers...")
+    try:
+        for i in range(0, len(all_ids), batch_size):
+            batch_ids = all_ids[i : i + batch_size]
+            batch_docs = all_docs[i : i + batch_size]
+            batch_metas = all_metas[i : i + batch_size]
+            new_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas)
+            filed += len(batch_ids)
+            print(f"  Re-filed {filed}/{len(all_ids)} drawers...")
+    except Exception as e:
+        print(f"\n  ERROR during rebuild: {e}")
+        print(f"  Only {filed}/{len(all_ids)} drawers were re-filed.")
+        if os.path.exists(backup_path):
+            print(f"  Restoring from backup: {backup_path}")
+            backend.delete_collection(palace_path, COLLECTION_NAME)
+            shutil.copy2(backup_path, sqlite_path)
+            print("  Backup restored. Palace is back to pre-repair state.")
+        else:
+            print("  No backup available. Re-mine from source files to recover.")
+        raise
 
     print(f"\n  Repair complete. {filed} drawers rebuilt.")
     print("  HNSW index is now clean with cosine distance metric.")