Tolan/backend/migrations/migrate_shop_items_special_scene_trigger.py

"""
迁移脚本：为 shop_items 增加 has_special_scene_trigger 并回填

执行：
python migrations/migrate_shop_items_special_scene_trigger.py
或：
python migrations/migrate_shop_items_special_scene_trigger.py --force
"""

import asyncio
import sys
from pathlib import Path
from typing import Any

import aiomysql

backend_dir = Path(__file__).resolve().parents[1]
if str(backend_dir) not in sys.path:
    sys.path.insert(0, str(backend_dir))

from config import settings


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


async def migrate() -> bool:
    conn: aiomysql.Connection | None = None
    try:
        conn = 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",
            autocommit=False,
        )
        async with conn.cursor() as cursor:
            has_column = await _column_exists(conn, "shop_items", "has_special_scene_trigger")
            if not has_column:
                await cursor.execute(
                    """
                    ALTER TABLE `shop_items`
                      ADD COLUMN `has_special_scene_trigger` TINYINT(1)
                      NOT NULL DEFAULT 0
                      COMMENT '是否有特殊触发场景(1=有,0=无)'
                      AFTER `scene_anchor_type`
                    """
                )
                print("✅ 已新增列 shop_items.has_special_scene_trigger")
            else:
                print("ℹ️ 列 shop_items.has_special_scene_trigger 已存在，跳过新增")

            await cursor.execute(
                """
                UPDATE `shop_items`
                SET `has_special_scene_trigger` = 1
                WHERE `item_code` IN ('coffee_mug', 'books', 'snacks')
                """
            )
            print(f"✅ 已回填特殊触发场景商品：{cursor.rowcount or 0} 行")

            await cursor.execute(
                """
                UPDATE `shop_items`
                SET `has_special_scene_trigger` = 0
                WHERE `item_code` NOT IN ('coffee_mug', 'books', 'snacks')
                """
            )
            print(f"✅ 已回填非特殊触发场景商品：{cursor.rowcount or 0} 行")

        await conn.commit()
        print("🎉 has_special_scene_trigger 迁移完成")
        return True
    except Exception as exc:  # noqa: BLE001
        if conn is not None:
            try:
                await conn.rollback()
            except Exception:
                pass
        print(f"❌ 迁移失败: {exc}")
        return False
    finally:
        if conn is not None:
            conn.close()


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续执行特殊触发场景字段迁移? (y/n): ").strip().lower()

    if confirm != "y":
        sys.exit(0)

    ok = asyncio.run(migrate())
    sys.exit(0 if ok else 1)