Tolan/backend/migrations/migrate_ab_subtitle_settings_init.py

"""
A/B 组字幕开关配置表：建表 + 种子数据（A=0 关，B=1 开）
运行（在 backend 目录）: python migrations/migrate_ab_subtitle_settings_init.py
"""
import asyncio
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import aiomysql

from config import settings

MIGRATIONS_DIR = Path(__file__).resolve().parent


async def table_exists(cursor: Any) -> bool:
    await cursor.execute(
        """
        SELECT COUNT(*) FROM information_schema.tables
        WHERE table_schema = %s AND table_name = 'ab_subtitle_settings'
        """,
        (settings.MYSQL_DATABASE,),
    )
    row = await cursor.fetchone()
    return bool(row and row[0] > 0)


async def migrate() -> bool:
    try:
        conn: aiomysql.Connection = 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:
            if not await table_exists(cursor):
                print("🔄 创建 ab_subtitle_settings 表...")
                ddl = (MIGRATIONS_DIR / "add_ab_subtitle_settings_table.sql").read_text(encoding="utf-8")
                for stmt in [s.strip() for s in ddl.split(";") if s.strip()]:
                    if stmt.upper().startswith("SET NAMES"):
                        continue
                    await cursor.execute(stmt)
                await conn.commit()
                print("✅ ab_subtitle_settings 表创建完成")
            else:
                print("✅ ab_subtitle_settings 表已存在")
                await cursor.execute(
                    """
                    INSERT INTO ab_subtitle_settings (group_key, subtitle_on) VALUES
                    ('A', 0),
                    ('B', 1)
                    ON DUPLICATE KEY UPDATE group_key = VALUES(group_key)
                    """
                )
                await conn.commit()
                print("✅ A/B 组种子数据已确保存在")

            await cursor.execute(
                "SELECT group_key, subtitle_on, updated_at FROM ab_subtitle_settings ORDER BY group_key"
            )
            rows = await cursor.fetchall()
            print("📊 当前配置:", rows)

        conn.close()
        return True
    except Exception:
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":
    confirm = input("\n是否继续执行迁移? (y/n): ") if len(sys.argv) <= 1 or sys.argv[1] != "--force" else "y"
    if confirm.lower() != "y":
        sys.exit(0)
    ok: Any = asyncio.run(migrate())
    sys.exit(0 if ok else 1)