支持 MySQL 迁移与性能优化

2026-07-25 10:23:19 +08:00
parent df4297bee5
commit 32f9eadfdb
9 changed files with 460 additions and 33 deletions
@@ -0,0 +1,192 @@
+"""将本项目 SQLite 数据安全迁移到 MySQL。"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from collections.abc import Iterable
+from pathlib import Path
+
+from sqlalchemy import JSON, MetaData, create_engine, func, inspect, select
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.schema import Table
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_SOURCE = PROJECT_ROOT / "data" / "web.db"
+IGNORED_SOURCE_TABLES = {"alembic_version"}
+
+
+def parse_args() -> argparse.Namespace:
+    """解析迁移参数。"""
+    parser = argparse.ArgumentParser(description="迁移 SQLite 数据到 MySQL")
+    parser.add_argument(
+        "--source",
+        type=Path,
+        default=DEFAULT_SOURCE,
+        help=f"SQLite 数据库路径，默认：{DEFAULT_SOURCE}",
+    )
+    parser.add_argument(
+        "--target-url",
+        help="MySQL SQLAlchemy 连接串，例如 mysql+pymysql://user:password@host:3306/database",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1000,
+        help="每批写入行数，默认 1000",
+    )
+    args = parser.parse_args()
+    if args.batch_size < 1:
+        parser.error("--batch-size 必须大于 0")
+    return args
+
+
+def _count_rows(connection: Connection, table: Table) -> int:
+    return connection.scalar(select(func.count()).select_from(table)) or 0
+
+
+def _chunks(rows: Iterable[dict], batch_size: int) -> Iterable[list[dict]]:
+    """将行数据分批，控制单次 SQL 的内存与事务体积。"""
+    batch: list[dict] = []
+    for row in rows:
+        batch.append(row)
+        if len(batch) >= batch_size:
+            yield batch
+            batch = []
+    if batch:
+        yield batch
+
+
+def _normalize_value(value: object, target_column) -> object:
+    """将 SQLite 中以文本保存的 JSON 转回结构化值。"""
+    if value is None or not isinstance(target_column.type, JSON):
+        return value
+    if not isinstance(value, str):
+        return value
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"列 {target_column.name} 存在非法 JSON：{value[:120]!r}") from exc
+
+
+def _source_rows(source_connection: Connection, source_table: Table, target_table: Table) -> Iterable[dict]:
+    columns = [column.name for column in target_table.columns]
+    missing_columns = [column for column in columns if column not in source_table.c]
+    if missing_columns:
+        raise RuntimeError(
+            f"源表 {source_table.name} 缺少列：{', '.join(missing_columns)}，"
+            "请先将 SQLite 升级到当前版本后再迁移"
+        )
+
+    primary_key_columns = list(source_table.primary_key.columns)
+    statement = select(*(source_table.c[name] for name in columns))
+    if primary_key_columns:
+        statement = statement.order_by(*(column.asc() for column in primary_key_columns))
+
+    for row in source_connection.execute(statement).mappings():
+        yield {
+            column_name: _normalize_value(row[column_name], target_table.c[column_name])
+            for column_name in columns
+        }
+
+
+def _ensure_target_is_empty(target_engine: Engine, target_tables: list[Table]) -> None:
+    """拒绝向已有业务数据的 MySQL 写入，防止误覆盖。"""
+    with target_engine.connect() as connection:
+        occupied = [
+            table.name
+            for table in target_tables
+            if _count_rows(connection, table) > 0
+        ]
+    if occupied:
+        raise RuntimeError(f"目标 MySQL 已存在业务数据，拒绝迁移：{', '.join(occupied)}")
+
+
+def main() -> int:
+    """执行迁移、计数核验并返回进程退出码。"""
+    args = parse_args()
+    source_path = args.source.expanduser().resolve()
+    target_url = (args.target_url or "").strip()
+
+    if not source_path.is_file():
+        raise FileNotFoundError(f"未找到 SQLite 数据库：{source_path}")
+
+    # 必须在导入数据库模块前设置，Alembic 环境才会使用迁移目标库。
+    if target_url:
+        os.environ["DATABASE_URL"] = target_url
+    from web.backend import models  # noqa: F401
+    from web.backend.database import Base, DATABASE_URL, run_migrations
+
+    target_url = target_url or DATABASE_URL
+    if not target_url.startswith("mysql+"):
+        raise ValueError("目标库必须是 MySQL；可传 --target-url，或设置 DB_HOST/DB_* 环境变量")
+
+    print("正在初始化目标 MySQL 表结构...")
+    run_migrations()
+
+    target_engine = create_engine(target_url, pool_pre_ping=True)
+    source_engine = create_engine(f"sqlite:///{source_path}")
+    target_tables = list(Base.metadata.sorted_tables)
+    target_table_names = {table.name for table in target_tables}
+
+    try:
+        source_table_names = set(inspect(source_engine).get_table_names())
+        unexpected_tables = source_table_names - target_table_names - IGNORED_SOURCE_TABLES
+        if unexpected_tables:
+            raise RuntimeError(
+                "源 SQLite 存在当前程序无法识别的表："
+                f"{', '.join(sorted(unexpected_tables))}，为避免遗漏已停止迁移"
+            )
+
+        source_metadata = MetaData()
+        source_metadata.reflect(
+            bind=source_engine,
+            only=sorted(source_table_names & target_table_names),
+        )
+        tables_to_copy = [
+            table for table in target_tables if table.name in source_metadata.tables
+        ]
+
+        _ensure_target_is_empty(target_engine, tables_to_copy)
+
+        with source_engine.connect() as source_connection, target_engine.begin() as target_connection:
+            for target_table in tables_to_copy:
+                source_table = source_metadata.tables[target_table.name]
+                source_count = _count_rows(source_connection, source_table)
+                if not source_count:
+                    print(f"{target_table.name}: 0 行，跳过")
+                    continue
+
+                print(f"{target_table.name}: 正在迁移 {source_count} 行...")
+                rows = _source_rows(source_connection, source_table, target_table)
+                for batch in _chunks(rows, args.batch_size):
+                    target_connection.execute(target_table.insert(), batch)
+
+        with source_engine.connect() as source_connection, target_engine.connect() as target_connection:
+            for target_table in tables_to_copy:
+                source_table = source_metadata.tables[target_table.name]
+                source_count = _count_rows(source_connection, source_table)
+                target_count = _count_rows(target_connection, target_table)
+                if source_count != target_count:
+                    raise RuntimeError(
+                        f"{target_table.name} 行数核验失败："
+                        f"SQLite={source_count}，MySQL={target_count}"
+                    )
+                print(f"{target_table.name}: 核验通过（{target_count} 行）")
+    finally:
+        source_engine.dispose()
+        target_engine.dispose()
+
+    print("迁移完成，SQLite 源文件未修改。")
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except Exception as exc:
+        print(f"迁移失败：{exc}", file=sys.stderr)
+        sys.exit(1)