Pine
ac4c1aabe4
- scripts/migrate_sqlite_to_mysql.py：离线迁移（MySQL 建表→alembic stamp head→ 逐表拷贝[NOT NULL 列类型化兜底]→自增回填→行数校验），支持 --drop-existing 重建与增量续跑 - tests/test_training_persist.py：夹具改 monkeypatch DATABASE_URL + reset_engine， 建表统一走 Base.metadata（原 DB_PATH/SCHEMA 随 sqlite3 直连移除）
151 lines
7.0 KiB
Python
151 lines
7.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""SQLite → MySQL 一次性迁移脚本（离线执行，非运行时注入）。
|
||
|
||
用法：
|
||
# 1) 在 .env 配置 PINEAGENTS_DEMO_DATABASE_URL 指向 MySQL（或用 --mysql 覆盖）
|
||
# 2) 跑迁移（默认目标 = config.DATABASE_URL，源 = SQLite 数据目录 app.db）
|
||
uv run python scripts/migrate_sqlite_to_mysql.py
|
||
uv run python scripts/migrate_sqlite_to_mysql.py --drop-existing # 危险：先清空 MySQL 目标表
|
||
|
||
流程：
|
||
1. 按平台 Base.metadata（含 app.pay.models）在 MySQL create_all 建全量表
|
||
（MySQL 方言钩子：String 无长度→VARCHAR(255)，Text→MEDIUMTEXT）
|
||
2. alembic stamp head（对齐迁移版本，后续 alembic upgrade 可用）
|
||
3. 从 SQLite 逐表拷贝（metadata 拓扑序，外键安全），自增表回填 AUTO_INCREMENT
|
||
4. 逐表行数校验，输出报告
|
||
|
||
幂等：表已存在且非空时默认跳过该表（增量续跑）；--drop-existing 强制重建。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||
|
||
from app import config # noqa: E402 加载 .env
|
||
from app.infrastructure.db import Base, make_sync_engine # noqa: E402
|
||
|
||
|
||
def _sqlite_url() -> str:
|
||
p = Path(config.DATA_DIR) / "app.db"
|
||
return f"sqlite:///{p}"
|
||
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--sqlite", default=_sqlite_url(), help="源 SQLite URL")
|
||
ap.add_argument("--mysql", default=config.DATABASE_URL, help="目标 MySQL URL（默认取 DATABASE_URL）")
|
||
ap.add_argument("--drop-existing", action="store_true", help="先 DROP 目标表再重建（危险）")
|
||
args = ap.parse_args()
|
||
|
||
if args.mysql.startswith("sqlite"):
|
||
print("[!] 目标仍是 SQLite——请先在 .env 配置 PINEAGENTS_DEMO_DATABASE_URL 指向 MySQL")
|
||
return 2
|
||
|
||
from sqlalchemy import Boolean, Integer, Numeric, Text, inspect, text
|
||
from app.infrastructure import models # noqa: F401 平台+培训 ORM 注册
|
||
import app.pay.models # noqa: F401 支付子应用 ORM 注册
|
||
|
||
src = make_sync_engine(args.sqlite)
|
||
dst = make_sync_engine(args.mysql)
|
||
insp = inspect(dst)
|
||
|
||
# 1) 建表 ---------------------------------------------------------------
|
||
existing = set(insp.get_table_names())
|
||
to_create = [t for t in Base.metadata.sorted_tables if t.name not in existing]
|
||
if args.drop_existing:
|
||
with dst.begin() as conn:
|
||
conn.execute(text("SET FOREIGN_KEY_CHECKS=0"))
|
||
for t in Base.metadata.sorted_tables:
|
||
conn.execute(text(f"DROP TABLE IF EXISTS `{t.name}`"))
|
||
conn.execute(text("SET FOREIGN_KEY_CHECKS=1"))
|
||
existing, to_create = set(), list(Base.metadata.sorted_tables)
|
||
print(f"[1/4] 已清空并重建目标（{len(to_create)} 表）")
|
||
else:
|
||
print(f"[1/4] MySQL 建表：新建 {len(to_create)} / 共 {len(Base.metadata.tables)}（已有 {len(existing)}）")
|
||
if to_create:
|
||
Base.metadata.create_all(dst, tables=to_create)
|
||
|
||
# 2) alembic stamp head -------------------------------------------------
|
||
from alembic.config import Config as AlembicConfig
|
||
from alembic import command
|
||
ac = AlembicConfig(str(Path(__file__).resolve().parent.parent / "alembic.ini"))
|
||
ac.set_main_option("sqlalchemy.url", args.mysql)
|
||
try:
|
||
command.stamp(ac, "head")
|
||
print("[2/4] alembic stamp head 完成")
|
||
except Exception as exc: # noqa: BLE001
|
||
print(f"[2/4] stamp 跳过：{exc}")
|
||
|
||
# 3) 拷贝数据 -----------------------------------------------------------
|
||
from sqlalchemy import MetaData, Table
|
||
|
||
copied, skipped = [], []
|
||
with src.connect() as sconn:
|
||
with dst.begin() as dconn:
|
||
dconn.execute(text("SET FOREIGN_KEY_CHECKS=0"))
|
||
try:
|
||
for t in Base.metadata.sorted_tables:
|
||
# 行数（源为空则跳过）
|
||
n_src = sconn.execute(text(f"SELECT COUNT(*) FROM `{t.name}`")).scalar() or 0
|
||
n_dst = dconn.execute(text(f"SELECT COUNT(*) FROM `{t.name}`")).scalar() or 0
|
||
if n_src == 0:
|
||
skipped.append((t.name, "源为空"))
|
||
continue
|
||
if n_dst >= n_src:
|
||
skipped.append((t.name, f"目标已有 {n_dst} 行"))
|
||
continue
|
||
rows = [dict(r._mapping) for r in sconn.execute(text(f"SELECT * FROM `{t.name}`"))]
|
||
# NOT NULL 兜底归一：SQLite 历史数据可能存在 NULL（迁移期 NOT NULL 未强制），
|
||
# MySQL 严格模式会拒收 → 按列类型补零值（str→"" / int→0 / bool→False / float→0.0）
|
||
for row in rows:
|
||
for c in t.columns:
|
||
if row.get(c.name) is None and not c.nullable:
|
||
if isinstance(c.type, Boolean):
|
||
row[c.name] = False
|
||
elif isinstance(c.type, Integer):
|
||
row[c.name] = 0
|
||
elif isinstance(c.type, Numeric):
|
||
row[c.name] = 0.0
|
||
elif isinstance(c.type, (Text, __import__("sqlalchemy").String)):
|
||
row[c.name] = ""
|
||
row[c.name] = 0.0
|
||
if rows:
|
||
dconn.execute(t.insert(), rows)
|
||
# 自增回填（MySQL ALTER 不允许子查询，先取 MAX 再写回）
|
||
ai_col = next((c.name for c in t.columns
|
||
if isinstance(c.type, Integer) and c.autoincrement
|
||
and c.primary_key), None)
|
||
if ai_col and rows:
|
||
mx = dconn.execute(text(f"SELECT IFNULL(MAX(`{ai_col}`),0) FROM `{t.name}`")).scalar() or 0
|
||
dconn.execute(text(f"ALTER TABLE `{t.name}` AUTO_INCREMENT = {int(mx) + 1}"))
|
||
copied.append((t.name, len(rows)))
|
||
finally:
|
||
dconn.execute(text("SET FOREIGN_KEY_CHECKS=1"))
|
||
|
||
# 4) 校验 ---------------------------------------------------------------
|
||
print("[3/4] 拷贝完成：")
|
||
for name, n in copied:
|
||
print(f" ✓ {name}: {n} 行")
|
||
for name, why in skipped:
|
||
print(f" - {name}: 跳过（{why}）")
|
||
|
||
bad = []
|
||
with src.connect() as sconn, dst.connect() as dconn:
|
||
for t in Base.metadata.sorted_tables:
|
||
a = sconn.execute(text(f"SELECT COUNT(*) FROM `{t.name}`")).scalar() or 0
|
||
b = dconn.execute(text(f"SELECT COUNT(*) FROM `{t.name}`")).scalar() or 0
|
||
if a != b:
|
||
bad.append((t.name, a, b))
|
||
print("[4/4] 行数校验：" + ("全部一致 ✓" if not bad else "不一致 ✗"))
|
||
for name, a, b in bad:
|
||
print(f" ✗ {name}: sqlite={a} mysql={b}")
|
||
return 1 if bad else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|