Tolan/backend/migrations/migrate_add_area_code_to_verification_codes.py

"""
数据库迁移脚本：为 verification_codes 表添加 area_code 字段
运行: python migrate_add_area_code_to_verification_codes.py
"""
import asyncio
import sys
from typing import Any

import aiomysql

from config import settings


async def migrate_add_area_code_field() -> bool:
    """为 verification_codes 表添加 area_code 字段"""
    try:
        conn: aiomysql.Connection = 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:
            # 检查字段是否已存在
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.columns
                WHERE table_schema = %s 
                AND table_name = 'verification_codes' 
                AND column_name = 'area_code'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row: Any = await cursor.fetchone()
            field_exists = bool(exists_row and exists_row[0] > 0)

            if not field_exists:
                print("🔄 为 verification_codes 表添加 area_code 字段...")
                await cursor.execute(
                    """
                    ALTER TABLE verification_codes
                    ADD COLUMN area_code VARCHAR(10) NULL COMMENT '手机区号（如 +86）'
                    AFTER code_id
                    """
                )
                await conn.commit()
                print("✅ area_code 字段添加完成")
            else:
                print("✅ area_code 字段已存在")

        conn.close()
        return True
    except Exception:
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续执行迁移? (y/n): ")

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

    ok = asyncio.run(migrate_add_area_code_field())
    if not ok:
        sys.exit(1)