yunke/yunke-server/scripts/migrate_sqlite_to_mysql.py

#!/usr/bin/env python
"""SQLite → MySQL 数据迁移(本地 smoke.db → 服务器 MySQL,部署方式B)。

在目标环境(MySQL 可达)里跑:
    python scripts/migrate_sqlite_to_mysql.py --sqlite ./smoke.db \
        --mysql "mysql+asyncmy://yunke:pwd@mysql:3306/yunke"

* 逐表 DELETE 清空后全量导入(幂等,可重跑)
* 复用 repos 的 _build_column_values 强转(bool/int/日期字符串→ms,2026-08-21
  的 3 个 MySQL 方言修复同样生效)
* JSON 列(raw_json 等)SQLite 里是文本,插入前 json.loads 回对象
* message(67万行)按 msg_id 键集分批,其余表全量读(最大 media 4.9万行)
* 跑完打印每表行数对照
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import sqlite3
import time

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("migrate")

BATCH = 500

# 迁移顺序(先组织后业务)
TABLES = [
    "department",
    "staff",
    "wechat_account",
    "whatsapp_account",
    "customer",
    "chat_group",
    "group_member",
    "message",
    "phone_call",
    "sms",
    "voice_clip",
    "media",
    "sync_run",
    "sync_state",
]


def _models_by_table() -> dict[str, type]:
    from app.db import models

    return {
        obj.__tablename__: obj
        for obj in vars(models).values()
        if isinstance(obj, type) and hasattr(obj, "__tablename__")
    }


def _json_cols(model: type) -> set[str]:
    return {c.name for c in model.__table__.columns if isinstance(c.type, sa.JSON)}


def _coerce_row(model: type, row: dict, json_cols: set[str]) -> dict:
    """SQLite 行 → MySQL 可绑定的值(复用 repos 强转 + JSON 文本解析)。"""
    from app.storage.repos import _build_column_values

    out = _build_column_values(model, row)
    for c in json_cols:
        v = out.get(c)
        if isinstance(v, str):
            try:
                out[c] = json.loads(v)
            except (ValueError, TypeError):
                out[c] = None
    return out


async def migrate_table(conn, model: type, sqlite_con, table: str) -> tuple[int, int]:
    """清空 MySQL 表并从 SQLite 全量导入。返回 (sqlite_rows, mysql_rows)。"""
    from app.storage.repos import _build_column_values  # noqa: F401 (coerce 内用)

    pk_cols = [c.name for c in model.__table__.primary_key.columns]
    json_cols = _json_cols(model)

    n_src = sqlite_con.execute(f"SELECT count(*) FROM {table}").fetchone()[0]
    await conn.execute(sa.text(f"DELETE FROM {table}"))

    inserted = 0
    if table == "message":
        # 67 万行:按第一主键(msg_id)键集分批,避免 OFFSET 全扫
        keyset_col = pk_cols[0]
        last_key = None
        while True:
            if last_key is None:
                rows = sqlite_con.execute(
                    f"SELECT * FROM {table} ORDER BY {keyset_col} LIMIT {BATCH}"
                ).fetchall()
            else:
                rows = sqlite_con.execute(
                    f"SELECT * FROM {table} WHERE {keyset_col} > ? ORDER BY {keyset_col} LIMIT {BATCH}",
                    (last_key,),
                ).fetchall()
            if not rows:
                break
            batch = [_coerce_row(model, dict(r), json_cols) for r in rows]
            if batch:
                await conn.execute(sa.insert(model), batch)
                inserted += len(batch)
            last_key = rows[-1][keyset_col]
            if inserted % 50_000 < BATCH:
                log.info(f"  {table}: {inserted}/{n_src}")
    else:
        rows = sqlite_con.execute(f"SELECT * FROM {table}").fetchall()
        for i in range(0, len(rows), BATCH):
            batch = [_coerce_row(model, dict(r), json_cols) for r in rows[i:i + BATCH]]
            if batch:
                await conn.execute(sa.insert(model), batch)
                inserted += len(batch)

    return n_src, inserted


async def main() -> None:
    parser = argparse.ArgumentParser(description="SQLite → MySQL 迁移")
    parser.add_argument("--sqlite", required=True, help="smoke.db 路径")
    parser.add_argument("--mysql", required=True, help="目标 MySQL async URL")
    args = parser.parse_args()

    t0 = time.time()
    sqlite_con = sqlite3.connect(args.sqlite)
    sqlite_con.row_factory = sqlite3.Row
    models_by_table = _models_by_table()

    engine = create_async_engine(args.mysql, pool_pre_ping=True, pool_recycle=3600)
    results: list[tuple[str, int, int]] = []
    async with engine.begin() as conn:
        for table in TABLES:
            model = models_by_table.get(table)
            if model is None:
                log.warning(f"跳过(无模型): {table}")
                continue
            t1 = time.time()
            n_src, n_ins = await migrate_table(conn, model, sqlite_con, table)
            results.append((table, n_src, n_ins))
            log.info(f"{table}: {n_src} → {n_ins} ({time.time()-t1:.1f}s)")
    await engine.dispose()
    sqlite_con.close()

    log.info("=" * 50)
    ok = True
    for table, n_src, n_ins in results:
        mark = "✓" if n_src == n_ins else "✗"
        if n_src != n_ins:
            ok = False
        log.info(f"{mark} {table}: {n_src:,} → {n_ins:,}")
    log.info(f"总耗时 {time.time()-t0:.0f}s — {'全部一致 ✓' if ok else '存在差异 ✗'}")


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