import logging from dataclasses import dataclass from typing import Dict, Any, Optional from bson import ObjectId from pymongo.results import InsertOneResult from classes.cache import Cache from classes.errors.pycord_user import UserNotFoundError from modules.database import col_users logger = logging.getLogger(__name__) @dataclass class PycordUser: """Dataclass of DB entry of a user""" __slots__ = ("_id", "id") _id: ObjectId id: int @classmethod async def from_id( cls, user_id: int, allow_creation: bool = True, cache: Optional[Cache] = None ) -> "PycordUser": """Find user in database and create new record if user does not exist. Args: user_id (int): User's Discord ID allow_creation (:obj:`bool`, optional): Create new user record if none found in the database cache (:obj:`Cache`, optional): Cache engine to get the cache from Returns: PycordUser: User object Raises: UserNotFoundError: User was not found and creation was not allowed """ if cache is not None: cached_entry: Dict[str, Any] | None = cache.get_json(f"user_{user_id}") if cached_entry is not None: return cls(**cached_entry) db_entry = await col_users.find_one({"id": user_id}) if db_entry is None: if not allow_creation: raise UserNotFoundError(user_id) db_entry = PycordUser.get_defaults(user_id) insert_result: InsertOneResult = await col_users.insert_one(db_entry) db_entry["_id"] = insert_result.inserted_id if cache is not None: cache.set_json(f"user_{user_id}", db_entry) return cls(**db_entry) def _to_dict(self) -> Dict[str, Any]: return { "_id": self._id, "id": self.id, } 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 col_users.update_one( {"_id": self._id}, {"$set": {key: value}}, upsert=True ) self._update_cache(cache) logger.info("Set attribute '%s' of user %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 = PycordUser.get_default_value(key) setattr(self, key, default_value) await col_users.update_one( {"_id": self._id}, {"$set": {key: default_value}}, upsert=True ) self._update_cache(cache) logger.info("Removed attribute '%s' of user %s", key, self.id) def _get_cache_key(self) -> str: return f"user_{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()) @staticmethod def get_defaults(user_id: int | None = None) -> Dict[str, Any]: return { "id": user_id, } @staticmethod def get_default_value(key: str) -> Any: if key not in PycordUser.get_defaults(): raise KeyError(f"There's no default value for key '{key}' in PycordUser") return PycordUser.get_defaults()[key] async def purge(self, cache: Optional[Cache] = None) -> None: """Completely remove user data from database. Currently only removes the user record from users collection. Args: cache (:obj:`Cache`, optional): Cache engine to write the update into """ await col_users.delete_one({"_id": self._id}) self._delete_cache(cache)