Tolan/backend/migrations/migrate_llm_tasks_init.py

"""
LLM 任务表：建表 + 种子 Task 01 + lumi_first_line 列迁移
运行（在 backend 目录）: python migrations/migrate_llm_tasks_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
from data.ab_task_01_prompt import TASK_01_NAME, TASK_01_PROMPT, TASK_01_FIRST_LINE

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 = 'llm_tasks'
        """,
        (settings.MYSQL_DATABASE,),
    )
    row = await cursor.fetchone()
    return bool(row and row[0] > 0)


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 = 'llm_tasks' AND column_name = %s
        """,
        (settings.MYSQL_DATABASE, column_name),
    )
    row = await cursor.fetchone()
    return bool(row and row[0] > 0)


async def seed_exists(cursor: Any) -> bool:
    await cursor.execute("SELECT COUNT(*) FROM llm_tasks WHERE task_name = %s", (TASK_01_NAME,))
    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("🔄 创建 llm_tasks 表...")
                ddl = (MIGRATIONS_DIR / "add_llm_tasks_table.sql").read_text(encoding="utf-8")
                create_part = ddl.split("INSERT INTO", 1)[0]
                for stmt in [s.strip() for s in create_part.split(";") if s.strip()]:
                    if stmt.upper().startswith("SET NAMES"):
                        continue
                    await cursor.execute(stmt)
                await conn.commit()
                print("✅ llm_tasks 表创建完成")
            else:
                print("✅ llm_tasks 表已存在")

            if not await column_exists(cursor, "lumi_first_line"):
                print("🔄 添加 lumi_first_line 列...")
                alter_sql = (MIGRATIONS_DIR / "add_llm_tasks_lumi_first_line.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("✅ lumi_first_line 列已添加")
            else:
                print("✅ lumi_first_line 列已存在")

            if not await seed_exists(cursor):
                print("🔄 插入 Task 01 种子数据...")
                await cursor.execute(
                    "INSERT INTO llm_tasks (task_name, task_prompt, lumi_first_line) VALUES (%s, %s, %s)",
                    (TASK_01_NAME, TASK_01_PROMPT, TASK_01_FIRST_LINE),
                )
                await conn.commit()
                print("✅ Task 01 已写入")
            else:
                print("✅ Task 01 已存在，跳过插入")
                await cursor.execute(
                    """
                    UPDATE llm_tasks
                    SET lumi_first_line = %s
                    WHERE task_name = %s
                      AND (lumi_first_line IS NULL OR lumi_first_line = '')
                    """,
                    (TASK_01_FIRST_LINE, TASK_01_NAME),
                )
                await conn.commit()

            await cursor.execute(
                "SELECT id, task_name, lumi_first_line, created_at FROM llm_tasks ORDER BY created_at DESC LIMIT 1"
            )
            latest = await cursor.fetchone()
            print("📊 当前最新任务:", latest)

        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)