Tolan/backend/migrations/migrate_add_qq_unionid.py

"""
数据库迁移脚本：为 users 表添加 qq_unionid 字段
运行: python migrate_add_qq_unionid.py
"""
import asyncio
import sys

import aiomysql

from config import settings


async def migrate_add_qq_unionid():
    """为 users 表添加 qq_unionid 字段"""
    conn = None
    try:
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset="utf8mb4",
        )

        async with conn.cursor() as cursor:
            # 检查 users 表是否存在
            await cursor.execute(
                """
                SELECT COUNT(*)
                FROM information_schema.tables 
                WHERE table_schema = %s AND table_name = 'users'
                """,
                (settings.MYSQL_DATABASE,),
            )
            table_exists = (await cursor.fetchone())[0] > 0
            if not table_exists:
                print("❌ users 表不存在，请先创建表")
                await conn.ensure_closed()
                return False

            # 检查是否已存在 qq_unionid 列
            await cursor.execute(
                """
                SELECT COUNT(*)
                FROM information_schema.columns 
                WHERE table_schema = %s 
                AND table_name = 'users' 
                AND column_name = 'qq_unionid'
                """,
                (settings.MYSQL_DATABASE,),
            )
            has_unionid = (await cursor.fetchone())[0] > 0

            if has_unionid:
                print("✅ qq_unionid 字段已存在，无需迁移")
                await conn.ensure_closed()
                return True

            print("🔄 添加 qq_unionid 字段...")
            await cursor.execute(
                """
                ALTER TABLE users 
                ADD COLUMN qq_unionid VARCHAR(128) UNIQUE NULL COMMENT 'QQ UnionID' 
                AFTER qq_openid
                """
            )
            await cursor.execute(
                """
                CREATE INDEX idx_qq_unionid ON users(qq_unionid)
                """
            )
            await conn.commit()
            print("✅ qq_unionid 字段添加成功")
            await conn.ensure_closed()
            return True
    except Exception as exc:
        import traceback

        print(f"❌ 迁移失败: {exc}")
        traceback.print_exc()
        if conn and not conn.closed:
            await conn.ensure_closed()
        return False


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续迁移 users 表添加 qq_unionid 字段? (y/n): ")

    if confirm.lower() != "y":
        sys.exit(0)

    success = asyncio.run(migrate_add_qq_unionid())
    if not success:
        sys.exit(1)