Tolan/backend/migrations/migrate_add_area_code_to_users.py

"""
数据库迁移脚本：为 users 表添加 area_code 字段
运行: python migrate_add_area_code_to_users.py
"""
import asyncio
import sys
import aiomysql
from config import settings


async def migrate_add_area_code_column() -> bool:
    """为 users 表添加 area_code 字段"""
    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 表不存在，请先创建表")
                conn.close()
                return False

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

            if has_area_code:
                print("✅ area_code 字段已存在，无需迁移")
                conn.close()
                return True

            print("🔄 添加 area_code 字段...")
            await cursor.execute(
                """
                ALTER TABLE users
                ADD COLUMN area_code VARCHAR(10) NULL COMMENT '手机区号（如 +86）' AFTER user_type
                """
            )
            print("✅ area_code 字段添加成功")

            await conn.commit()
            print("✅ 迁移完成")

        conn.close()
        return True

    except Exception as e:  # noqa: BLE001
        import traceback

        print(f"❌ 迁移失败: {e}")
        traceback.print_exc()
        return False


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

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

    success = asyncio.run(migrate_add_area_code_column())

    if not success:
        sys.exit(1)