Tolan/backend/migrations/migrate_add_user_level.py

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


async def migrate_add_user_level():
    """为 users 表添加 user_level 字段"""
    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:
            await cursor.execute(
                """
                SELECT COUNT(*)
                FROM INFORMATION_SCHEMA.COLUMNS
                WHERE TABLE_SCHEMA = %s
                  AND TABLE_NAME = 'users'
                  AND COLUMN_NAME = 'user_level'
                """,
                (settings.MYSQL_DATABASE,),
            )
            result = await cursor.fetchone()
            if result and result[0] > 0:
                conn.close()
                return True

            await cursor.execute(
                """
                ALTER TABLE users
                ADD COLUMN user_level INT NOT NULL DEFAULT 0 COMMENT '用户等级(0-5)'
                AFTER total_speech_duration_ms
                """
            )
            await conn.commit()

        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)

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