from flask_migrate import Migrate
from flask_prompt_master import create_app, db
import pymysql

def upgrade_database():
    """添加微信相关字段并修改字段约束"""
    try:
        conn = pymysql.connect(
            host='localhost',
            user='root',
            password='123456',
            database='food_db',
            charset='utf8mb4'
        )
        cursor = conn.cursor()
        
        # 修改现有字段约束
        modify_statements = [
            "ALTER TABLE user MODIFY mobile VARCHAR(20) NULL",
            "ALTER TABLE user MODIFY email VARCHAR(100) NULL",
            "ALTER TABLE user MODIFY avatar VARCHAR(64) NULL"
        ]
        
        print("\n=== 修改字段约束 ===")
        for statement in modify_statements:
            try:
                cursor.execute(statement)
                print(f"执行成功: {statement}")
            except Exception as e:
                print(f"执行失败: {statement}")
                print(f"错误: {str(e)}")
        
        # 添加新字段
        add_statements = [
            "ALTER TABLE user ADD COLUMN IF NOT EXISTS openid VARCHAR(64) UNIQUE",
            "ALTER TABLE user ADD COLUMN IF NOT EXISTS session_key VARCHAR(64)",
            "ALTER TABLE user ADD COLUMN IF NOT EXISTS unionid VARCHAR(64) UNIQUE",
            "ALTER TABLE user ADD COLUMN IF NOT EXISTS wx_nickname VARCHAR(100)",
            "ALTER TABLE user ADD COLUMN IF NOT EXISTS wx_avatar VARCHAR(255)"
        ]
        
        print("\n=== 添加新字段 ===")
        for statement in add_statements:
            try:
                cursor.execute(statement)
                print(f"执行成功: {statement}")
            except Exception as e:
                print(f"执行失败: {statement}")
                print(f"错误: {str(e)}")
        
        # 提交事务
        conn.commit()
        print("\n=== 数据库升级完成 ===")
        
    except Exception as e:
        print(f"数据库升级失败: {str(e)}")
        if 'conn' in locals():
            conn.rollback()
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

if __name__ == '__main__':
    upgrade_database()