157 lines
4.8 KiB
Python
157 lines
4.8 KiB
Python
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)
|