wow-token-app-bot/token_bot/persistant_database/user_schema.py

from typing import List, Dict

from aiodynamo.client import Table
from aiodynamo.errors import ItemNotFound

import token_bot.persistant_database as pdb
from token_bot.token_database.region import Region


class User:
    def __init__(
        self,
        user_id: int,
        region: Region = None,
        subscribed_alerts: List["pdb.Alert"] = None,
    ) -> None:
        self.user_id: int = user_id
        self._loaded: bool = False
        self.region: Region = region
        self.subscribed_alerts: List[pdb.Alert] = subscribed_alerts

    def __eq__(self, other):
        return self.user_id == other.user_id

    def __hash__(self):
        return hash(self.user_id)

    @classmethod
    def from_item(
        cls, primary_key: int, region: Region, subscribed_alerts: List[str]
    ) -> "User":
        alerts = [pdb.Alert.from_str(alert_str) for alert_str in subscribed_alerts]
        return cls(primary_key, region, alerts)

    @property
    def primary_key(self) -> str:
        return str(self.user_id)

    @property
    def primary_key_name(self) -> str:
        return "user_id"

    @property
    def key(self) -> Dict[str, str]:
        return {self.primary_key_name: self.primary_key}

    def _subscribed_alerts_as_trinity_list(self) -> List[str]:
        return (
            [str(alert) for alert in self.subscribed_alerts]
            if self.subscribed_alerts
            else []
        )

    async def _lazy_load(self, table: Table, consistent: bool = False) -> None:
        if consistent or not self._loaded:
            await self.get(table, consistent=consistent)

    async def put(self, table: Table) -> None:
        await table.put_item(
            item={
                self.primary_key_name: self.primary_key,
                "region": self.region,
                "subscribed_alerts": self._subscribed_alerts_as_trinity_list(),
            }
        )

    async def delete(self, table: Table) -> None:
        if not self._loaded:
            await self._lazy_load(table, consistent=True)
        await table.delete_item(
            key={self.primary_key_name: self.primary_key},
        )

    async def get(self, table: Table, consistent: bool = False) -> bool:
        try:
            response = await table.get_item(key=self.key, consistent_read=consistent)
        except ItemNotFound:
            return False

        self.subscribed_alerts = []
        for string_trinity in response["subscribed_alerts"]:
            self.subscribed_alerts.append(pdb.Alert.from_str(string_trinity))
        self.region = Region(response["region"])
        return True

    async def add_alert(
        self, table: Table, alert: "pdb.Alert", consistent: bool = False
    ) -> None:
        await self._lazy_load(table, consistent=consistent)
        if alert not in self.subscribed_alerts:
            self.subscribed_alerts.append(alert)
            await self.put(table)

    async def remove_alert(
        self, table: Table, alert: "pdb.Alert", consistent: bool = True
    ) -> None:
        await self._lazy_load(table, consistent=consistent)
        if alert in self.subscribed_alerts:
            self.subscribed_alerts.remove(alert)
            await self.put(table)