Tolan/backend/migrations/migrate_special_users_init.py

"""
特殊实验用户：加字段 + 预置 A/B 组 60 个账号
运行（在 backend 目录）: python migrations/migrate_special_users_init.py
"""
import asyncio
import sys
from pathlib import Path
from typing import Any

# 允许从 backend 目录直接运行
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import aiomysql

from config import settings

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


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


async def run_sql_file(cursor: Any, conn: Any, path: Path) -> None:
    sql = path.read_text(encoding="utf-8")
    statements = [s.strip() for s in sql.split(";") if s.strip()]
    for stmt in statements:
        if stmt.upper().startswith("SET NAMES"):
            continue
        await cursor.execute(stmt)
    await conn.commit()


async def count_special_users(cursor: Any) -> list[tuple[str, int]]:
    await cursor.execute(
        """
        SELECT special_group, COUNT(*) AS cnt
        FROM users
        WHERE is_special_user = 1
          AND email LIKE '%@lyrria.com'
          AND (email LIKE 'A%@lyrria.com' OR email LIKE 'B%@lyrria.com')
        GROUP BY special_group
        ORDER BY special_group
        """
    )
    return await cursor.fetchall()


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 column_exists(cursor, "is_special_user"):
                print("🔄 添加 is_special_user / special_group 字段...")
                alter_sql = (MIGRATIONS_DIR / "add_special_user_fields.sql").read_text(encoding="utf-8")
                for stmt in [s.strip() for s in alter_sql.split(";") if s.strip()]:
                    if stmt.upper().startswith("SET NAMES"):
                        continue
                    await cursor.execute(stmt)
                await conn.commit()
                print("✅ 字段添加完成")
            else:
                print("✅ is_special_user 字段已存在，跳过 ALTER")

            print("🔄 预置 A/B 组 60 个账号...")
            seed_path = MIGRATIONS_DIR / "seed_special_users_ab_groups.sql"
            seed_sql = seed_path.read_text(encoding="utf-8")
            for stmt in [s.strip() for s in seed_sql.split(";") if s.strip()]:
                upper = stmt.upper()
                if upper.startswith("SET NAMES") or upper.startswith("SELECT"):
                    if upper.startswith("SELECT"):
                        await cursor.execute(stmt)
                        rows = await cursor.fetchall()
                        print("📊 校验结果:", rows)
                    continue
                await cursor.execute(stmt)
            await conn.commit()
            print("✅ 账号预置完成")

            rows = await count_special_users(cursor)
            print("📊 当前特殊用户数量:", dict(rows) if rows else "无")

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

        traceback.print_exc()
        return False


if __name__ == "__main__":
    ok = asyncio.run(migrate())
    sys.exit(0 if ok else 1)