wow-token-app-bot/token_bot/controller/alerts.py

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 get_users(self, alert: str | Alert) -> List[User]:
        alert = self._alert_to_obj(alert)
        await alert.get(self.table, consistent=True)
        return alert.users

    async def add_user(self, alert: str | Alert, user: int | User):
        alert = self._alert_to_obj(alert)
        user = self._user_to_obj(user)
        await alert.add_user(self.table, user)

    async def remove_user(self, alert: str | Alert, user: int | User):
        alert = self._alert_to_obj(alert)
        user = self._user_to_obj(user)
        await alert.remove_user(self.table, user)