ai-LiaoWangweb-app/migrate_db.py

"""
数据库迁移：为已有的 user 表添加 SSO 新增字段
"""
import os
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), '.env'))

from app import create_app, db

def migrate():
    app = create_app()
    with app.app_context():
        conn = db.engine.connect()

        # 检查并添加 sso_sub
        result = conn.execute(
            db.text(
                "SELECT column_name FROM information_schema.columns "
                "WHERE table_name='user' AND column_name='sso_sub'"
            )
        )
        if not result.fetchone():
            conn.execute(db.text('ALTER TABLE "user" ADD COLUMN sso_sub VARCHAR(256)'))
            conn.execute(db.text('ALTER TABLE "user" ADD CONSTRAINT user_sso_sub_key UNIQUE (sso_sub)'))
            print("Added column sso_sub + unique constraint")
        else:
            print("sso_sub already exists")

        cols = [
            ("real_name", "VARCHAR(100)"),
            ("roles", "TEXT"),
            ("email", "VARCHAR(120)"),
            ("phone", "VARCHAR(30)"),
            ("avatar_url", "VARCHAR(500)"),
        ]
        for col, col_type in cols:
            result = conn.execute(
                db.text(
                    f"SELECT column_name FROM information_schema.columns "
                    f"WHERE table_name='user' AND column_name='{col}'"
                )
            )
            if not result.fetchone():
                conn.execute(db.text(f'ALTER TABLE "user" ADD COLUMN {col} {col_type}'))
                print(f"Added column {col}")
            else:
                print(f"{col} already exists")

        conn.commit()
        print("Migration complete.")

if __name__ == "__main__":
    migrate()