59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
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)
|
|
|
|
async def get_all_by_type(self, alert_type: int) -> List[Alert]:
|
|
"""Query all alerts of a specific type from the database."""
|
|
from aiodynamo.expressions import F
|
|
alerts = []
|
|
async for item in self.table.query(
|
|
key_condition=F("alert").equals(alert_type)
|
|
):
|
|
alert = Alert.from_item(
|
|
primary_key=item["alert"],
|
|
sort_key=item["flavor-region-price"],
|
|
users=item.get("users", [])
|
|
)
|
|
alerts.append(alert)
|
|
return alerts
|