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 } 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: return f"\n|\tRegion: {self.region.value.upper()}\tFlavor: {self.flavor.name.lower()}\tAlert: {self.alert_type.name.lower()}\t|" 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 = [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)