artemis/titles/idac/schema/factory.py

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