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.")