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.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
        }

    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.SPECIFIC_PRICE:
            raise NotImplementedError
        else:
            return f"\n|Region: {self.region.value.upper()}\tFlavor: {self.flavor.value.upper()}\tAlert: {self.alert_type.name.lower()}|"

    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:
        update_expression = F("users").delete({user.user_id})
        await table.update_item(
            key=self.key,
            update_expression=update_expression
        )

    async def put(self, table: Table) -> None:
        user_ids = [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
        if 'Item' in response:
            self._users = [pdb.User(int(user_id)) for user_id in response['Item']['users']['NS']]
            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)