hg-agents/scripts/migrate_memory_redis_to_pg.py

"""Redis → PostgreSQL 记忆数据迁移脚本

将 Redis 中现有记忆数据（消息、摘要）迁移到 PostgreSQL memory_messages 表。

使用方式:
  python scripts/migrate_memory_redis_to_pg.py [--dry-run] [--user-id UUID]

选项:
  --dry-run    仅扫描不实际写入
  --user-id    只迁移指定用户的数据
  --batch-size 每批写入条数 (默认100)
"""

import asyncio
import json
import uuid
import sys
import os
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))

from redis.asyncio import Redis
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from config import settings


async def scan_redis_keys(redis: Redis, pattern: str) -> list[str]:
    keys = []
    cursor = 0
    while True:
        cursor, batch = await redis.scan(cursor, match=pattern, count=200)
        keys.extend(batch)
        if cursor == 0:
            break
    return keys


def parse_key_info(key: str) -> dict | None:
    """从 Redis key 中解析 user_id, flow_id, session_id"""
    parts = key.split(":")
    info = {}

    uid_idx = None
    fid_idx = None
    sid_idx = None

    for i, p in enumerate(parts):
        try:
            val = uuid.UUID(p)
            if uid_idx is None:
                uid_idx = i
                info["user_id"] = val
            elif sid_idx is None:
                sid_idx = i
                info["session_id"] = val
            elif fid_idx is None:
                fid_idx = i
                info["flow_id"] = val
        except ValueError:
            if p == "messages" and uid_idx is not None and sid_idx is not None and fid_idx is None:
                fid_idx = -1
            continue

    if not info.get("user_id"):
        return None

    return info


async def migrate_old_format_keys(redis: Redis, engine, dry_run: bool):
    """迁移旧格式: mem:{uid}:{fid}:{sid}:messages (List) 和 mem:{uid}:{fid}:{sid}:meta (Hash)"""
    print(">>> 扫描旧格式 Redis 记忆键 ...")
    keys = await scan_redis_keys(redis, "mem:*:messages")
    keys += await scan_redis_keys(redis, "mem:*:meta")
    print(f"    找到 {len(keys)} 个旧格式键")

    migrated = 0
    session_info = {}

    for key in keys:
        key_type = await redis.type(key)
        info = parse_key_info(key)
        if not info:
            print(f"    [跳过] 无法解析键: {key}")
            continue

        uid = info.get("user_id")
        sid = info.get("session_id")

        if key_type == "hash":
            meta = await redis.hgetall(key)
            session_info[str(sid)] = meta
        elif key_type == "list":
            messages = await redis.lrange(key, 0, -1)
            for raw in messages:
                try:
                    msg = json.loads(raw)
                    role = msg.get("role", "user")
                    content = msg.get("content", "")
                    ts_str = msg.get("ts", msg.get("timestamp", ""))
                    created_at = datetime.fromisoformat(ts_str) if ts_str else datetime.utcnow()

                    if not dry_run:
                        async with AsyncSession(engine) as session:
                            await session.execute(
                                text("""
                                    INSERT INTO memory_messages (user_id, flow_id, session_id, role, content, created_at)
                                    VALUES (:uid, :fid, :sid, :role, :content, :ts)
                                    ON CONFLICT DO NOTHING
                                """),
                                {
                                    "uid": uid,
                                    "fid": info.get("flow_id"),
                                    "sid": sid,
                                    "role": role,
                                    "content": content,
                                    "ts": created_at,
                                },
                            )
                            await session.commit()
                    migrated += 1
                except (json.JSONDecodeError, Exception) as e:
                    print(f"    [错误] 解析消息失败: {raw[:80]}... -> {e}")

    return migrated


async def migrate_new_format_cache(redis: Redis, engine, dry_run: bool):
    """迁移新格式缓存: mem:cache:msgs:{uid}:{sid} (String JSON array)"""
    print(">>> 扫描新格式 Redis 消息缓存 ...")
    keys = await scan_redis_keys(redis, "mem:cache:msgs:*")
    print(f"    找到 {len(keys)} 个缓存键")

    migrated = 0
    for key in keys:
        parts = key.split(":")
        if len(parts) < 5:
            continue
        try:
            uid = uuid.UUID(parts[3])
            sid = uuid.UUID(parts[4])
        except (ValueError, IndexError):
            continue

        raw = await redis.get(key)
        if not raw:
            continue
        try:
            messages = json.loads(raw)
        except json.JSONDecodeError:
            continue

        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            ts_str = msg.get("ts", "")
            created_at = datetime.fromisoformat(ts_str) if ts_str else datetime.utcnow()

            if not dry_run:
                async with AsyncSession(engine) as session:
                    await session.execute(
                        text("""
                            INSERT INTO memory_messages (user_id, session_id, role, content, created_at)
                            VALUES (:uid, :sid, :role, :content, :ts)
                            ON CONFLICT DO NOTHING
                        """),
                        {"uid": uid, "sid": sid, "role": role, "content": content, "ts": created_at},
                    )
                    await session.commit()
            migrated += 1

    return migrated


async def migrate_summaries(redis: Redis, engine, dry_run: bool):
    """迁移 Redis 摘要到 PostgreSQL memory_atoms"""
    print(">>> 扫描 Redis 摘要缓存 ...")
    keys = await scan_redis_keys(redis, "mem:summary:*")
    print(f"    找到 {len(keys)} 个摘要键")

    migrated = 0
    for key in keys:
        parts = key.split(":")
        if len(parts) < 4:
            continue
        try:
            uid = uuid.UUID(parts[2])
            sid = uuid.UUID(parts[3])
        except (ValueError, IndexError):
            continue

        summary = await redis.get(key)
        if not summary or len(summary.strip()) < 10:
            continue

        if not dry_run:
            async with AsyncSession(engine) as session:
                result = await session.execute(
                    text("SELECT id FROM memory_atoms WHERE user_id = :uid AND source_session_id = :sid AND atom_type = 'summary'"),
                    {"uid": uid, "sid": sid},
                )
                existing = result.fetchone()
                if existing:
                    await session.execute(
                        text("UPDATE memory_atoms SET content = :content, updated_at = NOW() WHERE id = :id"),
                        {"content": summary, "id": existing[0]},
                    )
                else:
                    await session.execute(
                        text("""
                            INSERT INTO memory_atoms (user_id, atom_type, content, priority, source_session_id, created_at, updated_at)
                            VALUES (:uid, 'summary', :content, 60, :sid, NOW(), NOW())
                        """),
                        {"uid": uid, "content": summary, "sid": sid},
                    )
                await session.commit()
        migrated += 1

    return migrated


async def migrate_session_list(redis: Redis, engine, dry_run: bool):
    """迁移 mem:{uid}:sessions Set 中的会话列表"""
    print(">>> 扫描会话列表 (mem:*:sessions)...")
    keys = await scan_redis_keys(redis, "mem:*:sessions")
    print(f"    找到 {len(keys)} 个会话列表键")

    migrated = 0
    for key in keys:
        parts = key.split(":")
        if len(parts) < 3:
            continue
        try:
            uid = uuid.UUID(parts[1])
        except ValueError:
            continue

        sessions = await redis.smembers(key)
        for sid_str in sessions:
            try:
                sid = uuid.UUID(sid_str)
            except ValueError:
                continue

            if not dry_run:
                async with AsyncSession(engine) as session:
                    await session.execute(
                        text("""
                            INSERT INTO memory_sessions (user_id, session_id, last_active_at, created_at)
                            VALUES (:uid, :sid, NOW(), NOW())
                            ON CONFLICT (user_id, session_id) DO NOTHING
                        """),
                        {"uid": uid, "sid": sid},
                    )
                    await session.commit()
            migrated += 1

    return migrated


async def main():
    dry_run = "--dry-run" in sys.argv
    batch_size = 100

    for i, arg in enumerate(sys.argv):
        if arg == "--batch-size" and i + 1 < len(sys.argv):
            batch_size = int(sys.argv[i + 1])

    print("=" * 60)
    print("Redis → PostgreSQL 记忆数据迁移")
    print(f"模式: {'试运行(不写入)' if dry_run else '正式迁移'}")
    print(f"数据库: {settings.DATABASE_URL}")
    print(f"Redis: {settings.REDIS_URL}")
    print("=" * 60)

    engine = create_async_engine(settings.DATABASE_URL, echo=False)
    redis = Redis.from_url(settings.REDIS_URL, decode_responses=True)

    try:
        await redis.ping()
        print("Redis 连接成功")

        total = 0
        total += await migrate_old_format_keys(redis, engine, dry_run)
        total += await migrate_new_format_cache(redis, engine, dry_run)
        total += await migrate_summaries(redis, engine, dry_run)
        total += await migrate_session_list(redis, engine, dry_run)

        print(f"\n迁移完成！共处理 {total} 条记录")
        if dry_run:
            print("提示: 使用不带 --dry-run 参数运行以实际写入数据")
    finally:
        await redis.aclose()
        await engine.dispose()


if __name__ == "__main__":
    asyncio.run(main())