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)
        if self.subscribed_alerts:
            for alert in self.subscribed_alerts:
                await alert.remove_user(table, self)
        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