Tolan/backend/migrations/migrate_add_has_shared_avatar.py

"""
数据库迁移脚本：为 user_ui_states 表添加 has_shared_avatar 字段
运行: python migrate_add_has_shared_avatar.py
"""
import asyncio
import sys

import aiomysql

from config import settings


async def migrate_add_has_shared_avatar() -> bool:
  """为 user_ui_states 表添加 has_shared_avatar 字段"""
  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",
    )

    async with conn.cursor() as cursor:
      # 检查字段是否已存在
      await cursor.execute(
        """
        SELECT COUNT(*)
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_SCHEMA = %s
          AND TABLE_NAME = 'user_ui_states'
          AND COLUMN_NAME = 'has_shared_avatar'
        """,
        (settings.MYSQL_DATABASE,),
      )
      result = await cursor.fetchone()
      if result and result[0] > 0:
        conn.close()
        return True

      # 添加字段，默认 false
      await cursor.execute(
        """
        ALTER TABLE user_ui_states
        ADD COLUMN has_shared_avatar TINYINT(1) NOT NULL DEFAULT 0 COMMENT '是否已分享头像奖励'
        AFTER avatar_c_unlocked
        """
      )
      await conn.commit()

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

    traceback.print_exc()
    return False


if __name__ == "__main__":
  if len(sys.argv) > 1 and sys.argv[1] == "--force":
    confirm = "y"
  else:
    confirm = input("\n是否继续执行迁移? (y/n): ")

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

  ok = asyncio.run(migrate_add_has_shared_avatar())
  if not ok:
    sys.exit(1)