wow-token-app-bot/token_bot/persistant_database/alert_schema.py

from typing import List

from aiodynamo.client import Table
from aiodynamo.errors import ItemNotFound
from aiodynamo.expressions import F

from token_bot.persistant_database import AlertType
from token_bot.token_database.flavor import Flavor
from token_bot.token_database.region import Region
import token_bot.persistant_database as pdb


class Alert:
    def __init__(
        self, alert: pdb.AlertType, flavor: Flavor, region: Region, price: int = 0
    ) -> None:
        # AlertType is the Primary Key
        self._alert_type: pdb.AlertType = alert
        # Flavor (Retail, Classic) is the Sort Key
        self._flavor: Flavor = flavor
        self._region: Region = region
        self._price: int = price
        self._loaded: bool = False
        self._users: List[pdb.User] = []

    @classmethod
    def from_item(cls, primary_key: int, sort_key: str, users: List[int]) -> "Alert":
        alert_type = pdb.AlertType(primary_key)
        flavor_repr, region_repr, price_repr = sort_key.split("-")
        flavor = Flavor(int(flavor_repr))
        region = Region(region_repr)
        price = int(price_repr)
        return cls(alert_type, flavor, region, price)

    @classmethod
    def from_str(cls, string_trinity: str) -> "Alert":
        alert_repr, flavor_repr, region_repr, price_repr = string_trinity.split("-")
        if len(string_trinity.split("-")) != 4:
            raise ValueError
        alert = pdb.AlertType(int(alert_repr))
        flavor = Flavor(int(flavor_repr))
        region = Region(region_repr)
        price = int(price_repr)
        return cls(alert, flavor, region, price)

    @property
    def primary_key(self) -> int:
        return self.alert_type.value

    @property
    def primary_key_name(self) -> str:
        return "alert"

    @property
    def sort_key(self) -> str:
        return f"{self.flavor.value}-{self.region.value}-{self.price}"

    @property
    def sort_key_name(self) -> str:
        return "flavor-region-price"

    @property
    def key(self) -> dict[str, str | int]:
        return {
            self.primary_key_name: self.primary_key,
            self.sort_key_name: self.sort_key,
        }

    @property
    def alert_type(self) -> pdb.AlertType:
        return self._alert_type

    @property
    def flavor(self) -> Flavor:
        return self._flavor

    @property
    def region(self) -> Region:
        return self._region

    @property
    def price(self) -> int:
        return self._price

    @property
    def users(self) -> List[pdb.User]:
        return self._users

    def __str__(self):
        return f"{self.alert_type.value}-{self.flavor.value}-{self.region.value}-{self.price}"

    def __eq__(self, other):
        return (
            self.alert_type == other.alert_type
            and self.flavor == other.flavor
            and self.price == other.price
        )

    def to_human_string(self):
        if self.alert_type == AlertType.SPECIFIC_PRICE:
            raise NotImplementedError
        else:
            alert_type_str = " ".join(self.alert_type.name.split("_"))
            return f"{alert_type_str.title()}"

    async def _lazy_load(self, table: Table, consistent: bool = False) -> None:
        if consistent or not self._loaded:
            await self.get(table, consistent=consistent)

    async def _append_user(self, table: Table, user: pdb.User) -> None:
        self.users.append(user)
        await self.put(table)

    async def _remove_user(self, table: Table, user: pdb.User) -> None:
        self.users.remove(user)
        await self.put(table)

    async def put(self, table: Table) -> None:
        user_ids = [str(user.user_id) for user in self.users]
        await table.put_item(
            item={
                self.primary_key_name: self.primary_key,
                self.sort_key_name: self.sort_key,
                "users": user_ids,
            }
        )

    async def get(self, table: Table, consistent: bool = False) -> bool:
        try:
            response = await table.get_item(key=self.key, consistent_read=consistent)
        except ItemNotFound:
            return False
        self._users = [pdb.User(int(user_id)) for user_id in response["users"]]
        self._loaded = True
        return True

    async def get_users(self, table: Table, consistent: bool = False) -> List[pdb.User]:
        await self._lazy_load(table, consistent=consistent)

        return self.users

    async def add_user(
        self, table: Table, user: pdb.User, consistent: bool = False
    ) -> None:
        await self._lazy_load(table, consistent=consistent)

        if user not in self.users:
            await self._append_user(table=table, user=user)

    async def remove_user(
        self, table: Table, user: pdb.User, consistent: bool = True
    ) -> None:
        await self._lazy_load(table, consistent=consistent)

        if user in self.users:
            await self._remove_user(table=table, user=user)