Maas-Model-platform/backend/migrations/run_migrations_054_056.py

"""
执行迁移脚本 054 ~ 056

054: 为 models_new 表添加 group_name 和 encrypted_api_key 字段
055: 将 asr_task.result_text 从 VARCHAR(500) 改为 TEXT
056: 将 asr_task.file_url / result_url 从 VARCHAR(500) 改为 TEXT

使用方法（在项目根目录执行）：
    python backend/migrations/run_migrations_054_056.py

或在 backend/ 目录下执行：
    python migrations/run_migrations_054_056.py
"""
import sys
import os
from pathlib import Path

# 将 backend/ 加入 Python 路径，以便复用 DATABASE_URL
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

import psycopg2
from app.database import DATABASE_URL

# ==================== 迁移文件列表（按顺序执行）====================
MIGRATIONS_DIR = Path(__file__).parent
MIGRATION_FILES = [
    "054_add_group_name_and_encrypted_api_key.sql",
    "055_alter_asr_task_result_text_to_text.sql",
    "056_alter_asr_task_url_fields_to_text.sql",
]


def get_dsn(database_url: str) -> str:
    """将 SQLAlchemy URL 转换为 psycopg2 DSN"""
    # 兼容 postgresql+psycopg2:// 和 postgresql:// 两种格式
    url = database_url.replace("postgresql+psycopg2://", "postgresql://")
    return url


def run_sql_file(conn, filepath: Path) -> None:
    """执行单个 SQL 文件"""
    sql = filepath.read_text(encoding="utf-8")
    with conn.cursor() as cur:
        cur.execute(sql)


def main():
    print("=" * 60)
    print("迁移执行脚本：054 ~ 056")
    print("=" * 60)
    print()

    # 列出将要执行的文件
    print("将执行以下迁移文件：")
    for name in MIGRATION_FILES:
        path = MIGRATIONS_DIR / name
        status = "✓ 存在" if path.exists() else "✗ 文件不存在！"
        print(f"  [{status}] {name}")
    print()

    # 检查文件是否都存在
    missing = [n for n in MIGRATION_FILES if not (MIGRATIONS_DIR / n).exists()]
    if missing:
        print("错误：以下文件不存在，请检查路径：")
        for n in missing:
            print(f"  {MIGRATIONS_DIR / n}")
        sys.exit(1)

    confirm = input("确认执行以上迁移？(yes/no): ").strip().lower()
    if confirm != "yes":
        print("已取消。")
        sys.exit(0)

    print()

    # 连接数据库
    dsn = get_dsn(DATABASE_URL)
    print("正在连接数据库...")
    try:
        conn = psycopg2.connect(dsn)
        conn.autocommit = False
        print("✓ 数据库连接成功")
    except Exception as e:
        print(f"✗ 数据库连接失败：{e}")
        sys.exit(1)

    print()

    # 逐个执行迁移
    try:
        for name in MIGRATION_FILES:
            path = MIGRATIONS_DIR / name
            print(f"执行：{name} ...")
            try:
                run_sql_file(conn, path)
                print(f"  ✓ 成功")
            except Exception as e:
                conn.rollback()
                print(f"  ✗ 失败：{e}")
                print()
                print("已回滚所有未提交的变更。")
                sys.exit(1)

        conn.commit()
        print()
        print("=" * 60)
        print("✓ 全部迁移执行成功，已提交。")
        print("=" * 60)

    finally:
        conn.close()


if __name__ == "__main__":
    main()