71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
import os
|
|
from typing import List
|
|
|
|
import aiodynamo.client
|
|
import aiohttp
|
|
|
|
from token_bot.persistant_database import database as pdb, Alert
|
|
from token_bot.persistant_database.user_schema import User
|
|
|
|
|
|
class UsersController:
|
|
def __init__(self, session: aiohttp.ClientSession):
|
|
self._pdb: pdb.Database = pdb.Database(session)
|
|
self.table: aiodynamo.client.Table = self._pdb.client.table(
|
|
os.getenv("USERS_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(self, user: int | User):
|
|
user = self._user_to_obj(user)
|
|
await user.put(self.table)
|
|
|
|
async def delete(self, user: int | User):
|
|
user = self._user_to_obj(user)
|
|
await user.delete(self.table)
|
|
|
|
async def exists(self, user: int | User) -> bool:
|
|
user = self._user_to_obj(user)
|
|
return await user.get(self.table)
|
|
|
|
async def get(self, user: int | User) -> User:
|
|
user = self._user_to_obj(user)
|
|
await user.get(self.table)
|
|
return user
|
|
|
|
async def list_alerts(self, user: int | User) -> List[Alert]:
|
|
user = self._user_to_obj(user)
|
|
await user.get(self.table)
|
|
return user.subscribed_alerts
|
|
|
|
async def is_subscribed(self, user: int | User, alert: str | Alert) -> bool:
|
|
user = self._user_to_obj(user)
|
|
alert = self._alert_to_obj(alert)
|
|
await user.get(self.table)
|
|
return alert in user.subscribed_alerts
|
|
|
|
async def add_alert(self, user: int | User, alert: str | Alert) -> None:
|
|
user = self._user_to_obj(user)
|
|
alert = self._alert_to_obj(alert)
|
|
await user.get(self.table)
|
|
user.subscribed_alerts.append(alert)
|
|
await user.put(self.table)
|
|
|
|
async def remove_alert(self, user: int | User, alert: str | Alert) -> None:
|
|
user = self._user_to_obj(user)
|
|
alert = self._alert_to_obj(alert)
|
|
await user.get(self.table)
|
|
user.subscribed_alerts.remove(alert)
|
|
await user.put(self.table)
|