Tolan/backend/migrate_add_verification_codes_table.py

"""
数据库迁移脚本：创建验证码表
运行: python migrate_add_verification_codes_table.py
"""
import asyncio
import sys
from typing import Any

import aiomysql

from config import settings


async def migrate_create_verification_codes_table() -> bool:
    """创建验证码表"""
    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.tables
                WHERE table_schema = %s AND table_name = 'verification_codes'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row: Any = await cursor.fetchone()
            table_exists = bool(exists_row and exists_row[0] > 0)

            if not table_exists:
                print("🔄 创建 verification_codes 表...")
                await cursor.execute(
                    """
                    CREATE TABLE verification_codes (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        code_id VARCHAR(64) NOT NULL UNIQUE COMMENT '验证码ID',
                        
                        -- 接收方信息（手机号或邮箱，根据发送类型二选一）
                        phone VARCHAR(20) NULL COMMENT '手机号',
                        email VARCHAR(200) NULL COMMENT '邮箱',
                        
                        -- 验证码信息
                        code VARCHAR(10) NOT NULL COMMENT '验证码',
                        send_type ENUM('domestic_sms', 'international_sms', 'email') NOT NULL COMMENT '发送类型：国内短信、国外短信、邮箱',
                        
                        -- 状态信息
                        is_sent TINYINT(1) NOT NULL DEFAULT 0 COMMENT '是否已发送',
                        is_valid TINYINT(1) NOT NULL DEFAULT 1 COMMENT '是否有效',
                        
                        -- 时间信息
                        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                        sent_at TIMESTAMP NULL COMMENT '发送时间',
                        
                        -- 扩展信息
                        metadata JSON NULL COMMENT '扩展信息（发送结果、错误信息等）',
                        
                        INDEX idx_code_id (code_id),
                        INDEX idx_verification_code_phone (phone, is_valid, created_at),
                        INDEX idx_verification_code_email (email, is_valid, created_at),
                        INDEX idx_verification_code_send_type (send_type, is_sent, is_valid)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                    COMMENT='验证码表'
                    """
                )
                await conn.commit()
                print("✅ verification_codes 表创建完成")
            else:
                print("✅ verification_codes 表已存在")

        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_create_verification_codes_table())
    if not ok:
        sys.exit(1)