import os from typing import List import aiodynamo.client import aiohttp from token_bot.persistant_database import Alert, User from token_bot.persistant_database import database as pdb class AlertsController: def __init__(self, session: aiohttp.ClientSession): self._pdb: pdb.Database = pdb.Database(session) self.table = aiodynamo.client.Table = self._pdb.client.table(os.getenv('ALERTS_TABLE')) @staticmethod def _user_to_obj(user: int | User) -> User: if isinstance(user, int): return User(user) return user @staticmethod def _alert_to_obj(alert: str | Alert) -> Alert: if isinstance(alert, str): return Alert.from_str(alert) return alert async def add_user(self, user: int | User, alert: str | Alert) -> None: user = self._user_to_obj(user) alert = self._alert_to_obj(alert) await alert.add_user(self.table, user) async def delete_user(self, user: int | User, alert: str | Alert): user = self._user_to_obj(user) alert = self._alert_to_obj(alert) await alert.remove_user(self.table, user) async def get_users(self, alert: str | Alert, consistent: bool = False) -> List[User]: alert = self._alert_to_obj(alert) return await alert.get_users(self.table, consistent=consistent )