alerts.py (1910B)
1 import os 2 from typing import List 3 4 import aiodynamo.client 5 import aiohttp 6 7 from token_bot.persistant_database import Alert, User 8 from token_bot.persistant_database import database as pdb 9 10 11 class AlertsController: 12 def __init__(self, session: aiohttp.ClientSession): 13 self._pdb: pdb.Database = pdb.Database(session) 14 self.table: aiodynamo.client.Table = self._pdb.client.table( 15 os.getenv("ALERTS_TABLE") 16 ) 17 18 @staticmethod 19 def _user_to_obj(user: int | User) -> User: 20 if isinstance(user, int): 21 return User(user) 22 return user 23 24 @staticmethod 25 def _alert_to_obj(alert: str | Alert) -> Alert: 26 if isinstance(alert, str): 27 return Alert.from_str(alert) 28 return alert 29 30 async def get_users(self, alert: str | Alert) -> List[User]: 31 alert = self._alert_to_obj(alert) 32 await alert.get(self.table, consistent=True) 33 return alert.users 34 35 async def add_user(self, alert: str | Alert, user: int | User): 36 alert = self._alert_to_obj(alert) 37 user = self._user_to_obj(user) 38 await alert.add_user(self.table, user) 39 40 async def remove_user(self, alert: str | Alert, user: int | User): 41 alert = self._alert_to_obj(alert) 42 user = self._user_to_obj(user) 43 await alert.remove_user(self.table, user) 44 45 async def get_all_by_type(self, alert_type: int) -> List[Alert]: 46 """Query all alerts of a specific type from the database.""" 47 from aiodynamo.expressions import F 48 alerts = [] 49 async for item in self.table.query( 50 key_condition=F("alert").equals(alert_type) 51 ): 52 alert = Alert.from_item( 53 primary_key=item["alert"], 54 sort_key=item["flavor-region-price"], 55 users=item.get("users", []) 56 ) 57 alerts.append(alert) 58 return alerts