users.py (2270B)
1 import os 2 from typing import List 3 4 import aiodynamo.client 5 import aiohttp 6 7 from token_bot.persistant_database import database as pdb, Alert 8 from token_bot.persistant_database.user_schema import User 9 10 11 class UsersController: 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("USERS_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 add(self, user: int | User): 31 user = self._user_to_obj(user) 32 await user.put(self.table) 33 34 async def delete(self, user: int | User): 35 user = self._user_to_obj(user) 36 await user.delete(self.table) 37 38 async def exists(self, user: int | User) -> bool: 39 user = self._user_to_obj(user) 40 return await user.get(self.table) 41 42 async def get(self, user: int | User) -> User: 43 user = self._user_to_obj(user) 44 await user.get(self.table) 45 return user 46 47 async def list_alerts(self, user: int | User) -> List[Alert]: 48 user = self._user_to_obj(user) 49 await user.get(self.table) 50 return user.subscribed_alerts 51 52 async def is_subscribed(self, user: int | User, alert: str | Alert) -> bool: 53 user = self._user_to_obj(user) 54 alert = self._alert_to_obj(alert) 55 await user.get(self.table) 56 return alert in user.subscribed_alerts 57 58 async def add_alert(self, user: int | User, alert: str | Alert) -> None: 59 user = self._user_to_obj(user) 60 alert = self._alert_to_obj(alert) 61 await user.get(self.table) 62 user.subscribed_alerts.append(alert) 63 await user.put(self.table) 64 65 async def remove_alert(self, user: int | User, alert: str | Alert) -> None: 66 user = self._user_to_obj(user) 67 alert = self._alert_to_obj(alert) 68 await user.get(self.table) 69 user.subscribed_alerts.remove(alert) 70 await user.put(self.table)