alert_schema.py (5330B)
1 from typing import List 2 3 from aiodynamo.client import Table 4 from aiodynamo.errors import ItemNotFound 5 from aiodynamo.expressions import F 6 7 from token_bot.persistant_database import AlertType 8 from token_bot.token_database.flavor import Flavor 9 from token_bot.token_database.region import Region 10 import token_bot.persistant_database as pdb 11 12 13 class Alert: 14 def __init__( 15 self, alert: pdb.AlertType, flavor: Flavor, region: Region, price: int = 0 16 ) -> None: 17 # Validate price for SPECIFIC_PRICE alerts 18 if alert == AlertType.SPECIFIC_PRICE: 19 if price is None or price <= 0: 20 raise ValueError("SPECIFIC_PRICE alerts require a positive price value") 21 22 # AlertType is the Primary Key 23 self._alert_type: pdb.AlertType = alert 24 # Flavor (Retail, Classic) is the Sort Key 25 self._flavor: Flavor = flavor 26 self._region: Region = region 27 self._price: int = price 28 self._loaded: bool = False 29 self._users: List[pdb.User] = [] 30 31 @classmethod 32 def from_item(cls, primary_key: int, sort_key: str, users: List[int]) -> "Alert": 33 alert_type = pdb.AlertType(primary_key) 34 flavor_repr, region_repr, price_repr = sort_key.split("-") 35 flavor = Flavor(int(flavor_repr)) 36 region = Region(region_repr) 37 price = int(price_repr) 38 return cls(alert_type, flavor, region, price) 39 40 @classmethod 41 def from_str(cls, string_trinity: str) -> "Alert": 42 alert_repr, flavor_repr, region_repr, price_repr = string_trinity.split("-") 43 if len(string_trinity.split("-")) != 4: 44 raise ValueError 45 alert = pdb.AlertType(int(alert_repr)) 46 flavor = Flavor(int(flavor_repr)) 47 region = Region(region_repr) 48 price = int(price_repr) 49 return cls(alert, flavor, region, price) 50 51 @property 52 def primary_key(self) -> int: 53 return self.alert_type.value 54 55 @property 56 def primary_key_name(self) -> str: 57 return "alert" 58 59 @property 60 def sort_key(self) -> str: 61 return f"{self.flavor.value}-{self.region.value}-{self.price}" 62 63 @property 64 def sort_key_name(self) -> str: 65 return "flavor-region-price" 66 67 @property 68 def key(self) -> dict[str, str | int]: 69 return { 70 self.primary_key_name: self.primary_key, 71 self.sort_key_name: self.sort_key, 72 } 73 74 @property 75 def alert_type(self) -> pdb.AlertType: 76 return self._alert_type 77 78 @property 79 def flavor(self) -> Flavor: 80 return self._flavor 81 82 @property 83 def region(self) -> Region: 84 return self._region 85 86 @property 87 def price(self) -> int: 88 return self._price 89 90 @property 91 def users(self) -> List[pdb.User]: 92 return self._users 93 94 def __str__(self): 95 return f"{self.alert_type.value}-{self.flavor.value}-{self.region.value}-{self.price}" 96 97 def __eq__(self, other): 98 return ( 99 self.alert_type == other.alert_type 100 and self.flavor == other.flavor 101 and self.region == other.region 102 and self.price == other.price 103 ) 104 105 def to_human_string(self): 106 if self.alert_type == AlertType.SPECIFIC_PRICE: 107 price_gold = self.price 108 return f"Custom Price: {format(price_gold, ',')}g" 109 else: 110 alert_type_str = " ".join(self.alert_type.name.split("_")) 111 return f"{alert_type_str.title()}" 112 113 async def _lazy_load(self, table: Table, consistent: bool = False) -> None: 114 if consistent or not self._loaded: 115 await self.get(table, consistent=consistent) 116 117 async def _append_user(self, table: Table, user: pdb.User) -> None: 118 self.users.append(user) 119 await self.put(table) 120 121 async def _remove_user(self, table: Table, user: pdb.User) -> None: 122 self.users.remove(user) 123 await self.put(table) 124 125 async def put(self, table: Table) -> None: 126 user_ids = [str(user.user_id) for user in self.users] 127 await table.put_item( 128 item={ 129 self.primary_key_name: self.primary_key, 130 self.sort_key_name: self.sort_key, 131 "users": user_ids, 132 } 133 ) 134 135 async def get(self, table: Table, consistent: bool = False) -> bool: 136 try: 137 response = await table.get_item(key=self.key, consistent_read=consistent) 138 except ItemNotFound: 139 return False 140 self._users = [pdb.User(int(user_id)) for user_id in response["users"]] 141 self._loaded = True 142 return True 143 144 async def get_users(self, table: Table, consistent: bool = False) -> List[pdb.User]: 145 await self._lazy_load(table, consistent=consistent) 146 147 return self.users 148 149 async def add_user( 150 self, table: Table, user: pdb.User, consistent: bool = False 151 ) -> None: 152 await self._lazy_load(table, consistent=consistent) 153 154 if user not in self.users: 155 await self._append_user(table=table, user=user) 156 157 async def remove_user( 158 self, table: Table, user: pdb.User, consistent: bool = True 159 ) -> None: 160 await self._lazy_load(table, consistent=consistent) 161 162 if user in self.users: 163 await self._remove_user(table=table, user=user)