from dataclasses import dataclass from logging import Logger from typing import Any, Dict, Optional from bson import ObjectId from libbot.cache.classes import Cache from pymongo.results import InsertOneResult from classes.errors import GuildNotFoundError from modules.database import col_guilds from modules.logging_utils import get_logger logger: Logger = get_logger(__name__) @dataclass class PycordGuild: """Dataclass of DB entry of a guild""" __slots__ = ("_id", "id", "channel_id", "category_id", "timezone") __short_name__ = "guild" __collection__ = col_guilds _id: ObjectId id: int channel_id: Optional[int] category_id: Optional[int] timezone: str @classmethod async def from_id( cls, guild_id: int, allow_creation: bool = True, cache: Optional[Cache] = None ) -> "PycordGuild": """Find guild in database and create new record if guild does not exist. Args: guild_id (int): User's Discord ID allow_creation (:obj:`bool`, optional): Create new guild record if none found in the database cache (:obj:`Cache`, optional): Cache engine to get the cache from Returns: PycordGuild: User object Raises: GuildNotFoundError: User was not found and creation was not allowed """ if cache is not None: cached_entry: Dict[str, Any] | None = cache.get_json(f"{cls.__short_name__}_{guild_id}") if cached_entry is not None: return cls(**cached_entry) db_entry = await cls.__collection__.find_one({"id": guild_id}) if db_entry is None: if not allow_creation: raise GuildNotFoundError(guild_id) db_entry = PycordGuild.get_defaults(guild_id) insert_result: InsertOneResult = await cls.__collection__.insert_one(db_entry) db_entry["_id"] = insert_result.inserted_id if cache is not None: cache.set_json(f"{cls.__short_name__}_{guild_id}", db_entry) return cls(**db_entry) async def _set(self, key: str, value: Any, cache: Optional[Cache] = None) -> None: """Set attribute data and save it into the database. Args: key (str): Attribute to change value (Any): Value to set cache (:obj:`Cache`, optional): Cache engine to write the update into """ if not hasattr(self, key): raise AttributeError() setattr(self, key, value) await self.__collection__.update_one({"_id": self._id}, {"$set": {key: value}}, upsert=True) self._update_cache(cache) logger.info("Set attribute '%s' of guild %s to '%s'", key, self.id, value) async def _remove(self, key: str, cache: Optional[Cache] = None) -> None: """Remove attribute data and save it into the database. Args: key (str): Attribute to remove cache (:obj:`Cache`, optional): Cache engine to write the update into """ if not hasattr(self, key): raise AttributeError() default_value: Any = PycordGuild.get_default_value(key) setattr(self, key, default_value) await self.__collection__.update_one({"_id": self._id}, {"$set": {key: default_value}}, upsert=True) self._update_cache(cache) logger.info("Removed attribute '%s' of guild %s", key, self.id) def _get_cache_key(self) -> str: return f"{self.__short_name__}_{self.id}" def _update_cache(self, cache: Optional[Cache] = None) -> None: if cache is None: return user_dict: Dict[str, Any] = self.to_dict() if user_dict is not None: cache.set_json(self._get_cache_key(), user_dict) else: self._delete_cache(cache) def _delete_cache(self, cache: Optional[Cache] = None) -> None: if cache is None: return cache.delete(self._get_cache_key()) def to_dict(self, json_compatible: bool = False) -> Dict[str, Any]: """Convert PycordGuild object to a JSON representation. Args: json_compatible (bool): Whether the JSON-incompatible objects like ObjectId need to be converted Returns: Dict[str, Any]: JSON representation of PycordGuild """ return { "_id": self._id if not json_compatible else str(self._id), "id": self.id, "channel_id": self.channel_id, "category_id": self.category_id, "timezone": self.timezone, } @staticmethod def get_defaults(guild_id: Optional[int] = None) -> Dict[str, Any]: return {"id": guild_id, "channel_id": None, "category_id": None, "timezone": "UTC"} @staticmethod def get_default_value(key: str) -> Any: if key not in PycordGuild.get_defaults(): raise KeyError(f"There's no default value for key '{key}' in PycordGuild") return PycordGuild.get_defaults()[key] async def purge(self, cache: Optional[Cache] = None) -> None: """Completely remove guild data from database. Currently only removes the guild record from guilds collection. Args: cache (:obj:`Cache`, optional): Cache engine to write the update into """ await self.__collection__.delete_one({"_id": self._id}) self._delete_cache(cache) # TODO Add documentation def is_configured(self) -> bool: return ( (self.id is not None) and (self.channel_id is not None) and (self.category_id is not None) and (self.timezone is not None) ) # TODO Add documentation async def set_channel(self, channel_id: Optional[int] = None, cache: Optional[Cache] = None) -> None: await self._set("channel_id", channel_id, cache) # TODO Add documentation async def reset_channel(self, cache: Optional[Cache] = None) -> None: await self._remove("channel_id", cache) # TODO Add documentation async def set_category(self, category_id: Optional[int] = None, cache: Optional[Cache] = None) -> None: await self._set("category_id", category_id, cache) # TODO Add documentation async def reset_category(self, cache: Optional[Cache] = None) -> None: await self._remove("category_id", cache) # TODO Add documentation async def set_timezone(self, timezone: str, cache: Optional[Cache] = None) -> None: await self._set("timezone", timezone, cache) # TODO Add documentation async def reset_timezone(self, cache: Optional[Cache] = None) -> None: await self._remove("timezone", cache)