from datetime import datetime from typing import Dict, Optional from sqlalchemy import ( Table, Column, UniqueConstraint, and_, ) from sqlalchemy.types import Integer, TIMESTAMP from sqlalchemy.schema import ForeignKey from sqlalchemy.engine import Row from sqlalchemy.sql import func, select from sqlalchemy.dialects.mysql import insert from core.data.schema import BaseData, metadata lottery = Table( "idac_user_lottery", metadata, Column("id", Integer, primary_key=True, nullable=False), Column("user", ForeignKey("aime_user.id", ondelete="cascade", onupdate="cascade")), Column("version", Integer, nullable=False), Column("saved_value", Integer, nullable=False), Column("lottery_count", Integer, nullable=False), Column("create_date", TIMESTAMP, server_default=func.now()), UniqueConstraint("user", "version", name="idac_user_lottery_uk"), mysql_charset="utf8mb4", ) class IDACFactoryData(BaseData): async def get_lottery(self, aime_id: int, version: int) -> Optional[Row]: sql = select(lottery).where( and_( lottery.c.user == aime_id, lottery.c.version == version ) ) result = await self.execute(sql) if result is None: return None return result.fetchone() async def put_lottery( self, aime_id: int, version: int, saved_value: int, lottery_count: int, create_date: datetime ) -> Optional[int]: lottery_data = { "user": aime_id, "version": version, "saved_value": saved_value, "lottery_count": lottery_count, "create_date": create_date } sql = insert(lottery).values(**lottery_data) conflict = sql.on_duplicate_key_update(**lottery_data) result = await self.execute(conflict) if result is None: self.logger.warn(f"put_lottery: Failed to update! aime_id: {aime_id}") return None return result.lastrowid