Files
wehub-resource-sync c3749daf48
Tests / test-linux (3.13) (push) Failing after 0s
Tests / test-linux (3.11) (push) Failing after 1s
Tests / lint (push) Failing after 0s
Tests / test-linux (3.9) (push) Failing after 1s
Docker / build (push) Failing after 1s
Docker / build-gpu (push) Failing after 2s
Tests / test-windows (push) Has been cancelled
Tests / test-macos (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:03:03 +08:00

592 lines
22 KiB
Python

#!/usr/bin/env python3
"""
mempalace migrate — Recover a palace created with a different ChromaDB version.
Reads documents and metadata directly from the palace's SQLite database
(bypassing ChromaDB's API, which fails on version-mismatched palaces),
then re-imports everything into a fresh palace using the currently installed
ChromaDB version.
Since mempalace 3.2.0 (chromadb>=1.5.4), chromadb automatically migrates
0.4.1+ databases on first open — no manual migration needed for upgrades.
Use this command only when downgrading chromadb (e.g. rolling back to an
older mempalace release) or if automatic migration fails.
Usage:
mempalace migrate # migrate default palace
mempalace migrate --palace /path/to/palace # migrate specific palace
mempalace migrate --dry-run # show what would be migrated
"""
import errno
import glob
import os
import shutil
import sqlite3
import tempfile
import uuid
from collections import defaultdict
from contextlib import closing
from datetime import datetime
from .backups import prune_backups
from .config import MempalaceConfig
def _restore_stale_palace(palace_path: str, stale_path: str) -> None:
"""Roll back a failed swap.
shutil.move() can partially create palace_path before raising, which
would make a bare os.replace(stale_path, palace_path) fail (dest exists).
Clear any partial destination first, then restore. Best-effort: if the
restore itself fails, log both paths so the operator can recover by hand.
"""
try:
if os.path.lexists(palace_path):
shutil.rmtree(palace_path, ignore_errors=True)
os.replace(stale_path, palace_path)
except Exception as err:
print(
f" CRITICAL: rollback failed — original palace at {stale_path}, "
f"partial migration data at {palace_path}. Restore manually. "
f"({err})"
)
def extract_drawers_from_sqlite(db_path: str) -> list:
"""Read all drawers directly from ChromaDB's SQLite, bypassing the API.
Works regardless of which ChromaDB version created the database.
Returns list of dicts with 'id', 'document', and 'metadata' keys.
The connection is wrapped in ``contextlib.closing`` so an exception
during extraction does not leak the SQLite handle. On Windows that
would leave a file lock on ``chroma.sqlite3`` and prevent the rest
of the migration from touching the palace directory.
"""
with closing(sqlite3.connect(db_path)) as conn:
conn.row_factory = sqlite3.Row
# Get all embedding IDs and their documents
rows = conn.execute(
"""
SELECT e.embedding_id,
MAX(CASE WHEN em.key = 'chroma:document' THEN em.string_value END) as document
FROM embeddings e
JOIN embedding_metadata em ON em.id = e.id
GROUP BY e.embedding_id
"""
).fetchall()
drawers = []
for row in rows:
embedding_id = row["embedding_id"]
document = row["document"]
if not document:
continue
# Get metadata for this embedding
meta_rows = conn.execute(
"""
SELECT em.key, em.string_value, em.int_value, em.float_value, em.bool_value
FROM embedding_metadata em
JOIN embeddings e ON e.id = em.id
WHERE e.embedding_id = ?
AND em.key NOT LIKE 'chroma:%'
""",
(embedding_id,),
).fetchall()
metadata = {}
for mr in meta_rows:
key = mr["key"]
if mr["string_value"] is not None:
metadata[key] = mr["string_value"]
elif mr["int_value"] is not None:
metadata[key] = mr["int_value"]
elif mr["float_value"] is not None:
metadata[key] = mr["float_value"]
elif mr["bool_value"] is not None:
metadata[key] = bool(mr["bool_value"])
drawers.append(
{
"id": embedding_id,
"document": document,
"metadata": metadata,
}
)
return drawers
def detect_chromadb_version(db_path: str) -> str:
"""Detect which ChromaDB version created the database by checking schema."""
conn = sqlite3.connect(db_path)
try:
# 1.x has schema_str column in collections table
cols = [r[1] for r in conn.execute("PRAGMA table_info(collections)").fetchall()]
if "schema_str" in cols:
return "1.x"
# 0.6.x has embeddings_queue but no schema_str
tables = [
r[0]
for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
]
if "embeddings_queue" in tables:
return "0.6.x"
return "unknown"
finally:
conn.close()
def contains_palace_database(path: str) -> bool:
"""Return True when path looks like a MemPalace ChromaDB directory."""
return os.path.isfile(os.path.join(path, "chroma.sqlite3"))
def confirm_destructive_action(
operation_name: str, palace_path: str, assume_yes: bool = False
) -> bool:
"""Require confirmation before destructive palace operations."""
if assume_yes:
return True
print(f"\n {operation_name} will replace data in: {palace_path}")
print(" A backup will be created first, then the palace will be rebuilt.")
try:
answer = input(" Continue? [y/N]: ").strip().lower()
except EOFError:
print(" Aborted. Re-run with --yes to confirm destructive changes.")
return False
if answer not in {"y", "yes"}:
print(" Aborted.")
return False
return True
def _result_ids(result) -> list:
"""Return ids from either the backend typed result or raw Chroma dict."""
if isinstance(result, dict):
return list(result.get("ids") or [])
return list(getattr(result, "ids", []) or [])
def collection_write_roundtrip_works(col) -> bool:
"""Return True only if the collection can upsert, read, and delete.
Some ChromaDB 0.6.x -> 1.5.x migrated collections remain readable while
writes and deletes silently no-op. A plain ``count()`` probe misses that
failure mode, so migrate must verify an actual write round-trip before
deciding that no rebuild is needed.
"""
probe_id = f"_mempalace_migrate_probe_{uuid.uuid4().hex}"
probe_doc = "mempalace migrate write round-trip probe"
probe_meta = {
"wing": "_mempalace_probe",
"room": "_mempalace_probe",
"source_file": "mempalace_migrate_probe",
"chunk_index": 0,
}
try:
col.upsert(
ids=[probe_id],
documents=[probe_doc],
metadatas=[probe_meta],
)
after_upsert = col.get(ids=[probe_id], include=[])
if probe_id not in _result_ids(after_upsert):
return False
col.delete(ids=[probe_id])
after_delete = col.get(ids=[probe_id], include=[])
if probe_id in _result_ids(after_delete):
return False
return True
except Exception:
return False
def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
"""Migrate a palace to the currently installed ChromaDB version."""
from .backends.chroma import ChromaBackend
palace_path = os.path.abspath(os.path.expanduser(palace_path))
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isdir(palace_path) or not contains_palace_database(palace_path):
print(f"\n No palace database found at {db_path}")
return False
print(f"\n{'=' * 60}")
print(" MemPalace Migrate")
print(f"{'=' * 60}\n")
print(f" Palace: {palace_path}")
print(f" Database: {db_path}")
print(f" DB size: {os.path.getsize(db_path) / 1024 / 1024:.1f} MB")
# Detect version
source_version = detect_chromadb_version(db_path)
target_version = ChromaBackend.backend_version()
print(f" Source: ChromaDB {source_version}")
print(f" Target: ChromaDB {target_version}")
# Try reading and writing with current chromadb first.
#
# A plain count() is not enough: some 0.6.x -> 1.5.x migrated collections
# are readable but silently drop upsert/delete operations. In that state,
# migrate must rebuild from SQLite instead of returning "No migration needed."
try:
col = ChromaBackend().get_collection(palace_path, "mempalace_drawers")
count = col.count()
if collection_write_roundtrip_works(col):
print(f"\n Palace is already readable and writable by chromadb {target_version}.")
print(f" {count} drawers found. No migration needed.")
return True
print(
f"\n Palace is readable by chromadb {target_version}, but write/delete verification failed."
)
print(" Rebuilding from SQLite to restore native write/delete behavior...")
except Exception:
print(f"\n Palace is NOT readable by chromadb {target_version}.")
print(" Extracting from SQLite directly...")
# Extract all drawers via raw SQL
drawers = extract_drawers_from_sqlite(db_path)
print(f" Extracted {len(drawers)} drawers from SQLite")
if not drawers:
print(" Nothing to migrate.")
return True
# Show summary
wings = defaultdict(lambda: defaultdict(int))
for d in drawers:
w = d["metadata"].get("wing", "?")
r = d["metadata"].get("room", "?")
wings[w][r] += 1
print("\n Summary:")
for wing, rooms in sorted(wings.items()):
total = sum(rooms.values())
print(f" WING: {wing} ({total} drawers)")
for room, count in sorted(rooms.items(), key=lambda x: -x[1]):
print(f" ROOM: {room:30} {count:5}")
if dry_run:
print("\n DRY RUN — no changes made.")
print(f" Would migrate {len(drawers)} drawers.")
return True
if not confirm_destructive_action("Migration", palace_path, assume_yes=confirm):
return False
# Backup the old palace
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
backup_path = f"{palace_path}.pre-migrate.{timestamp}"
print(f"\n Backing up to {backup_path}...")
shutil.copytree(palace_path, backup_path, symlinks=True)
# Enforce backup retention so repeated migrations cannot fill the disk
# with full-palace copies. The backup we just created is the newest, so
# it survives; only older ``.pre-migrate.*`` siblings beyond the limit
# are removed. Best-effort — never let cleanup fail the migration.
prune_backups(
glob.escape(palace_path) + ".pre-migrate.*",
MempalaceConfig().max_backups,
log=print,
)
# Build fresh palace in a temp directory (avoids chromadb reading old state).
# Wrap the whole import-and-swap dance in try/finally so the temp dir is
# cleaned up if any of the chromadb writes, the verify count, or the
# rename fails — without try/finally a crashed migration leaves a partial
# palace dir under the system temp root that the user has to find by hand.
temp_palace = tempfile.mkdtemp(prefix="mempalace_migrate_")
try:
print(f" Creating fresh palace in {temp_palace}...")
fresh_backend = ChromaBackend()
col = fresh_backend.get_or_create_collection(temp_palace, "mempalace_drawers")
# Re-import in batches
batch_size = 500
imported = 0
for i in range(0, len(drawers), batch_size):
batch = drawers[i : i + batch_size]
col.add(
ids=[d["id"] for d in batch],
documents=[d["document"] for d in batch],
metadatas=[d["metadata"] for d in batch],
)
imported += len(batch)
print(f" Imported {imported}/{len(drawers)} drawers...")
# Verify before swapping
final_count = col.count()
del col
del fresh_backend
# 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...")
stale_path = palace_path + ".old"
if os.path.exists(stale_path):
shutil.rmtree(stale_path)
os.replace(palace_path, stale_path)
try:
os.replace(temp_palace, palace_path)
except OSError as e:
# EXDEV = temp lives on a different filesystem; fall back to copy+delete.
# Anything else is a real error — don't mask it with shutil.move.
if getattr(e, "errno", None) != errno.EXDEV:
_restore_stale_palace(palace_path, stale_path)
raise
try:
shutil.move(temp_palace, palace_path)
except Exception:
_restore_stale_palace(palace_path, stale_path)
raise
shutil.rmtree(stale_path, ignore_errors=True)
finally:
# On the happy path os.replace/shutil.move consumed temp_palace, so
# the directory no longer exists at the temp location — the existence
# guard makes this a no-op then. On any failure path it actually
# removes the orphan.
if os.path.exists(temp_palace):
shutil.rmtree(temp_palace, ignore_errors=True)
print("\n Migration complete.")
print(f" Drawers migrated: {final_count}")
print(f" Backup at: {backup_path}")
if final_count != len(drawers):
print(f" WARNING: Expected {len(drawers)}, got {final_count}")
print(f"\n{'=' * 60}\n")
return True
# ---------------------------------------------------------------------------
# Wing-name normalization migration (#1675 follow-up)
# ---------------------------------------------------------------------------
#
# normalize_wing_name now strips leading/trailing separators, so a path-encoded
# dirname like ``-home-user-proj`` derives ``home_user_proj`` instead of
# ``_home_user_proj``. Palaces built before that rule filed drawers under the
# old, leading-underscore wing, which the new derivation no longer matches —
# searches and diary reads under the new name miss the old memories.
#
# This migration re-keys the ``wing`` metadata field on drawers and closets to
# the normalized form, merging collisions. Drawer/closet IDs embed the wing as
# an opaque prefix that is never decoded back into a wing (verified: nothing
# splits a wing out of an ID; mining idempotency keys on ``source_file``), so
# the IDs are left untouched — closet ``→drawer_id`` pointers stay valid and
# future mining still skips already-mined files. Tunnels resolve via existing
# read-time normalization and need no rewrite. The pass is idempotent.
def _normalized_wing_target(wing):
"""Return the normalized wing if it differs from ``wing``, else ``None``.
``None`` means "no migration needed" — either the value is not a non-empty
string, normalization is a no-op, or it would normalize to empty.
"""
from .config import normalize_wing_name
if not isinstance(wing, str) or not wing:
return None
# Apply the full normalization and explicitly strip leading/trailing
# separators. The strip is this migration's whole purpose (#1675); doing it
# here rather than relying on normalize_wing_name keeps the migration correct
# even when run against a build whose normalize_wing_name predates #1675, and
# matches the post-#1675 derivation exactly.
target = normalize_wing_name(wing).strip("_")
if not target or target == wing:
return None
return target
def plan_wing_renames(items):
"""Pure planner over ``(id, metadata)`` pairs.
Returns ``(summary, updates)`` where ``summary`` is ``{(old, new): count}``
and ``updates`` is ``[(id, new_metadata), ...]`` for only the records whose
wing changes. Metadata is copied; only the ``wing`` key is rewritten.
"""
summary = defaultdict(int)
updates = []
for rec_id, meta in items:
meta = dict(meta or {})
target = _normalized_wing_target(meta.get("wing"))
if target is None:
continue
summary[(meta["wing"], target)] += 1
meta["wing"] = target
updates.append((rec_id, meta))
return summary, updates
def _iter_collection_items(col, batch_size=1000):
"""Yield ``(id, metadata)`` for every record in a backend collection."""
total = col.count()
offset = 0
while offset < total:
batch = col.get(limit=batch_size, offset=offset, include=["metadatas"])
ids = batch.ids if hasattr(batch, "ids") else batch["ids"]
metas = batch.metadatas if hasattr(batch, "metadatas") else batch["metadatas"]
if not ids:
break
for rec_id, meta in zip(ids, metas):
yield rec_id, meta
offset += len(ids)
def _apply_wing_updates(col, updates, batch_size=500):
"""Re-label the ``wing`` metadata field in place for the planned updates."""
for i in range(0, len(updates), batch_size):
chunk = updates[i : i + batch_size]
col.update(ids=[u[0] for u in chunk], metadatas=[u[1] for u in chunk])
def _plan_topics_by_wing_renames():
"""Return ``{old_wing: new_wing}`` for ``topics_by_wing`` keys to normalize."""
try:
from .miner import _load_known_entities_raw
reg = _load_known_entities_raw()
except Exception:
return {}
tbw = reg.get("topics_by_wing")
if not isinstance(tbw, dict):
return {}
renames = {}
for wing in list(tbw.keys()):
target = _normalized_wing_target(wing)
if target is not None:
renames[wing] = target
return renames
def _apply_topics_by_wing_renames(renames):
"""Re-key ``topics_by_wing`` in known_entities.json, merging on collision."""
if not renames:
return
import json
from .miner import _ENTITY_REGISTRY_PATH, _load_known_entities_raw
try:
reg = _load_known_entities_raw()
except Exception:
return
tbw = reg.get("topics_by_wing")
if not isinstance(tbw, dict):
return
for old, new in renames.items():
if old not in tbw:
continue
old_topics = tbw.pop(old) or []
if new in tbw:
merged = list(tbw[new])
for topic in old_topics:
if topic not in merged:
merged.append(topic)
tbw[new] = merged
else:
tbw[new] = old_topics
reg["topics_by_wing"] = tbw
os.makedirs(os.path.dirname(_ENTITY_REGISTRY_PATH), exist_ok=True)
fd, tmp = tempfile.mkstemp(dir=os.path.dirname(_ENTITY_REGISTRY_PATH), suffix=".tmp")
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(reg, f, ensure_ascii=False, indent=2)
os.replace(tmp, _ENTITY_REGISTRY_PATH)
except Exception:
if os.path.exists(tmp):
os.remove(tmp)
raise
def migrate_wing_names(palace_path: str, dry_run: bool = False, confirm: bool = False) -> bool:
"""Normalize legacy wing names in ``palace_path`` (strip leading/trailing
separators), so palaces built before #1675 keep their memories discoverable.
Re-keys the ``wing`` metadata on drawers and closets in place (IDs untouched)
and the ``topics_by_wing`` registry, merging collisions. Idempotent.
Returns True if anything was (or, in dry-run, would be) migrated.
"""
from .palace import get_closets_collection, get_collection
try:
drawers = get_collection(palace_path, create=False)
except Exception as exc:
print(f" No drawer collection found at {palace_path} ({exc}).")
return False
d_items = list(_iter_collection_items(drawers))
all_wings = {(m or {}).get("wing") for _, m in d_items if (m or {}).get("wing")}
d_summary, d_updates = plan_wing_renames(d_items)
closets = None
c_summary, c_updates = defaultdict(int), []
try:
closets = get_closets_collection(palace_path, create=False)
c_summary, c_updates = plan_wing_renames(_iter_collection_items(closets))
except Exception:
closets = None
topic_renames = _plan_topics_by_wing_renames()
if not d_updates and not c_updates and not topic_renames:
print(" All wing names are already normalized — nothing to migrate.")
return False
print("\n Wing-name migration plan:")
merged = defaultdict(lambda: [0, 0])
for key, count in d_summary.items():
merged[key][0] = count
for key, count in c_summary.items():
merged[key][1] = count
for (old, new), (d_count, c_count) in sorted(merged.items()):
note = " (MERGE into existing wing)" if new in all_wings else ""
print(f" {old!r} -> {new!r}: {d_count} drawer(s), {c_count} closet(s){note}")
if topic_renames:
print(f" topics_by_wing: {len(topic_renames)} key(s) re-keyed")
if dry_run:
print("\n DRY RUN — no changes made.\n")
return True
if not confirm:
try:
resp = input(" Apply this wing-name migration? [y/N] ").strip().lower()
except EOFError:
resp = ""
if resp not in ("y", "yes"):
print(" Aborted.")
return False
_apply_wing_updates(drawers, d_updates)
if closets is not None and c_updates:
_apply_wing_updates(closets, c_updates)
_apply_topics_by_wing_renames(topic_renames)
parts = [f"{len(d_updates)} drawer(s)"]
if c_updates:
parts.append(f"{len(c_updates)} closet(s)")
if topic_renames:
parts.append(f"{len(topic_renames)} topic key(s)")
print(f"\n Migrated {', '.join(parts)}.\n")
return True