Tolan/backend/migrations/migrate_add_bot_name.py

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


async def migrate_add_bot_name():
    """为 users 表添加 bot_name 字段"""
    
    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 = 'bot_name'
            """, (settings.MYSQL_DATABASE,))
            
            result = await cursor.fetchone()
            if result[0] > 0:
                conn.close()
                return True
            
            # 添加 bot_name 字段
            await cursor.execute("""
                ALTER TABLE users 
                ADD COLUMN bot_name VARCHAR(100) COMMENT '用户给机器人取的名字' 
                AFTER user_type
            """)
            
            await conn.commit()
        
        conn.close()
        
        return True
        
    except Exception as e:
        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_bot_name())
    
    if not success:
        sys.exit(1)