101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
from typing import List, Dict
|
|
|
|
from aiodynamo.client import Table
|
|
from aiodynamo.errors import ItemNotFound
|
|
|
|
import token_bot.persistant_database as pdb
|
|
from token_bot.token_database.region import Region
|
|
|
|
|
|
class User:
|
|
def __init__(
|
|
self,
|
|
user_id: int,
|
|
region: Region = None,
|
|
subscribed_alerts: List["pdb.Alert"] = None,
|
|
) -> None:
|
|
self.user_id: int = user_id
|
|
self._loaded: bool = False
|
|
self.region: Region = region
|
|
self.subscribed_alerts: List[pdb.Alert] = subscribed_alerts
|
|
|
|
def __eq__(self, other):
|
|
return self.user_id == other.user_id
|
|
|
|
def __hash__(self):
|
|
return hash(self.user_id)
|
|
|
|
@classmethod
|
|
def from_item(
|
|
cls, primary_key: int, region: Region, subscribed_alerts: List[str]
|
|
) -> "User":
|
|
alerts = [pdb.Alert.from_str(alert_str) for alert_str in subscribed_alerts]
|
|
return cls(primary_key, region, alerts)
|
|
|
|
@property
|
|
def primary_key(self) -> str:
|
|
return str(self.user_id)
|
|
|
|
@property
|
|
def primary_key_name(self) -> str:
|
|
return "user_id"
|
|
|
|
@property
|
|
def key(self) -> Dict[str, str]:
|
|
return {self.primary_key_name: self.primary_key}
|
|
|
|
def _subscribed_alerts_as_trinity_list(self) -> List[str]:
|
|
return (
|
|
[str(alert) for alert in self.subscribed_alerts]
|
|
if self.subscribed_alerts
|
|
else []
|
|
)
|
|
|
|
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 put(self, table: Table) -> None:
|
|
await table.put_item(
|
|
item={
|
|
self.primary_key_name: self.primary_key,
|
|
"region": self.region,
|
|
"subscribed_alerts": self._subscribed_alerts_as_trinity_list(),
|
|
}
|
|
)
|
|
|
|
async def delete(self, table: Table) -> None:
|
|
if not self._loaded:
|
|
await self._lazy_load(table, consistent=True)
|
|
await table.delete_item(
|
|
key={self.primary_key_name: self.primary_key},
|
|
)
|
|
|
|
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.subscribed_alerts = []
|
|
for string_trinity in response["subscribed_alerts"]:
|
|
self.subscribed_alerts.append(pdb.Alert.from_str(string_trinity))
|
|
self.region = Region(response["region"])
|
|
return True
|
|
|
|
async def add_alert(
|
|
self, table: Table, alert: "pdb.Alert", consistent: bool = False
|
|
) -> None:
|
|
await self._lazy_load(table, consistent=consistent)
|
|
if alert not in self.subscribed_alerts:
|
|
self.subscribed_alerts.append(alert)
|
|
await self.put(table)
|
|
|
|
async def remove_alert(
|
|
self, table: Table, alert: "pdb.Alert", consistent: bool = True
|
|
) -> None:
|
|
await self._lazy_load(table, consistent=consistent)
|
|
if alert in self.subscribed_alerts:
|
|
self.subscribed_alerts.remove(alert)
|
|
await self.put(table)
|