shudao-main/shudao-chat-py/run_migration.py

#!/usr/bin/env python3
"""
数据库迁移脚本：为user表添加points字段
"""
import pymysql
import sys

# 数据库配置
DB_CONFIG = {
    'host': '172.16.29.101',
    'port': 21000,
    'user': 'root',
    'password': '88888888',
    'database': 'shudao',
    'charset': 'utf8mb4'
}

# 迁移SQL语句列表
MIGRATION_STATEMENTS = [
    # 1. 添加积分字段
    "ALTER TABLE `user` ADD COLUMN `points` INT DEFAULT 0 COMMENT '积分余额' AFTER `role`",
    # 2. 为现有用户初始化积分
    "UPDATE `user` SET `points` = 100 WHERE `points` IS NULL"
]

def run_migration():
    """执行数据库迁移"""
    connection = None
    try:
        print("正在连接数据库...")
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        
        print("开始执行迁移...")
        
        # 逐条执行SQL语句
        statements = MIGRATION_STATEMENTS
        
        for i, statement in enumerate(statements, 1):
            if statement:
                print(f"\n执行语句 {i}:")
                print(f"  {statement[:100]}...")
                cursor.execute(statement)
                print(f"  ✓ 成功")
        
        connection.commit()
        
        # 验证结果
        print("\n验证迁移结果...")
        cursor.execute("SELECT id, username, points FROM `user` LIMIT 5")
        results = cursor.fetchall()
        
        print("\n用户积分数据:")
        print("ID\t用户名\t\t积分")
        print("-" * 40)
        for row in results:
            print(f"{row[0]}\t{row[1]}\t\t{row[2]}")
        
        print("\n✓ 迁移成功完成！")
        return True
        
    except pymysql.err.OperationalError as e:
        error_code, error_msg = e.args
        if error_code == 1060:  # Duplicate column name
            print(f"\n⚠ 字段已存在，无需重复添加: {error_msg}")
            return True
        else:
            print(f"\n✗ 数据库操作错误: {error_msg}")
            return False
            
    except Exception as e:
        print(f"\n✗ 迁移失败: {str(e)}")
        if connection:
            connection.rollback()
        return False
        
    finally:
        if connection:
            cursor.close()
            connection.close()
            print("\n数据库连接已关闭")

if __name__ == '__main__':
    success = run_migration()
    sys.exit(0 if success else 1)