From b3a7e3623a979b6a4059820d52584c89200c417f Mon Sep 17 00:00:00 2001 From: profitroll Date: Sun, 9 Feb 2025 23:00:18 +0100 Subject: [PATCH 1/7] Implemented memcached caching --- classes/__init__.py | 0 classes/cache/__init__.py | 3 + classes/cache/holo_cache.py | 44 +++++++++ classes/cache/holo_cache_memcached.py | 106 ++++++++++++++++++++ classes/cache/holo_cache_redis.py | 34 +++++++ classes/holo_bot.py | 15 +++ classes/holo_user.py | 133 ++++++++++++++++++++------ cogs/analytics.py | 2 +- cogs/custom_channels.py | 20 ++-- cogs/data.py | 19 +--- cogs/logger.py | 38 ++------ config_example.json | 16 ++-- main.py | 15 +-- migrations/202502092252.py | 63 ++++++++++++ modules/cache_utils.py | 29 ++++++ modules/database.py | 3 + requirements.txt | 2 + validation/analytics.json | 11 ++- validation/users.json | 4 +- validation/warnings.json | 4 +- 20 files changed, 460 insertions(+), 101 deletions(-) create mode 100644 classes/__init__.py create mode 100644 classes/cache/__init__.py create mode 100644 classes/cache/holo_cache.py create mode 100644 classes/cache/holo_cache_memcached.py create mode 100644 classes/cache/holo_cache_redis.py create mode 100644 migrations/202502092252.py create mode 100644 modules/cache_utils.py diff --git a/classes/__init__.py b/classes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/classes/cache/__init__.py b/classes/cache/__init__.py new file mode 100644 index 0000000..05d7db8 --- /dev/null +++ b/classes/cache/__init__.py @@ -0,0 +1,3 @@ +from .holo_cache import HoloCache +from .holo_cache_memcached import HoloCacheMemcached +from .holo_cache_redis import HoloCacheRedis diff --git a/classes/cache/holo_cache.py b/classes/cache/holo_cache.py new file mode 100644 index 0000000..5969365 --- /dev/null +++ b/classes/cache/holo_cache.py @@ -0,0 +1,44 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict + +import pymemcache +import redis + + +class HoloCache(ABC): + client: pymemcache.Client | redis.Redis + + @classmethod + @abstractmethod + def from_config(cls, engine_config: Dict[str, Any]) -> Any: + pass + + @abstractmethod + def get_json(self, key: str) -> Any | None: + # TODO This method must also carry out ObjectId conversion! + pass + + @abstractmethod + def get_string(self, key: str) -> str | None: + pass + + @abstractmethod + def get_object(self, key: str) -> Any | None: + pass + + @abstractmethod + def set_json(self, key: str, value: Any) -> None: + # TODO This method must also carry out ObjectId conversion! + pass + + @abstractmethod + def set_string(self, key: str, value: str) -> None: + pass + + @abstractmethod + def set_object(self, key: str, value: Any) -> None: + pass + + @abstractmethod + def delete(self, key: str) -> None: + pass diff --git a/classes/cache/holo_cache_memcached.py b/classes/cache/holo_cache_memcached.py new file mode 100644 index 0000000..a764d24 --- /dev/null +++ b/classes/cache/holo_cache_memcached.py @@ -0,0 +1,106 @@ +import logging +from logging import Logger +from typing import Dict, Any + +from bson import ObjectId +from pymemcache import Client +from ujson import dumps, loads + +from . import HoloCache + +logger: Logger = logging.getLogger(__name__) + + +class HoloCacheMemcached(HoloCache): + client: Client + + def __init__(self, client: Client): + self.client = client + + @classmethod + def from_config(cls, engine_config: Dict[str, Any]) -> "HoloCacheMemcached": + if "uri" not in engine_config: + raise KeyError( + "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" + ) + + return cls(Client(engine_config["uri"], default_noreply=True)) + + def get_json(self, key: str) -> Any | None: + try: + result: Any | None = self.client.get(key, None) + + logger.debug( + "Got json cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + except Exception as exc: + logger.error("Could not get json cache key '%s' due to: %s", key, exc) + return None + + return None if result is None else self._string_to_json(result) + + def get_string(self, key: str) -> str | None: + try: + result: str | None = self.client.get(key, None) + + logger.debug( + "Got string cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + + return result + except Exception as exc: + logger.error("Could not get string cache key '%s' due to: %s", key, exc) + return None + + def get_object(self, key: str) -> Any | None: + # TODO Implement binary deserialization + raise NotImplementedError() + + def set_json(self, key: str, value: Any) -> None: + try: + self.client.set(key, self._json_to_string(value)) + logger.debug("Set json cache key '%s'", key) + except Exception as exc: + logger.error("Could not set json cache key '%s' due to: %s", key, exc) + return None + + def set_string(self, key: str, value: str) -> None: + try: + self.client.set(key, value) + logger.debug("Set string cache key '%s'", key) + except Exception as exc: + logger.error("Could not set string cache key '%s' due to: %s", key, exc) + return None + + def set_object(self, key: str, value: Any) -> None: + # TODO Implement binary serialization + raise NotImplementedError() + + def delete(self, key: str) -> None: + try: + self.client.delete(key) + logger.debug("Deleted cache key '%s'", key) + except Exception as exc: + logger.error("Could not delete cache key '%s' due to: %s", key, exc) + + @staticmethod + def _json_to_string(json_object: Any) -> str: + if isinstance(json_object, dict) and "_id" in json_object: + json_object["_id"] = str(json_object["_id"]) + + return dumps( + json_object, ensure_ascii=False, indent=0, escape_forward_slashes=False + ) + + @staticmethod + def _string_to_json(json_string: str) -> Any: + json_object: Any = loads(json_string) + + if isinstance(json_object, dict) and "_id" in json_object: + json_object["_id"] = ObjectId(json_object["_id"]) + + return json_object diff --git a/classes/cache/holo_cache_redis.py b/classes/cache/holo_cache_redis.py new file mode 100644 index 0000000..a22a9aa --- /dev/null +++ b/classes/cache/holo_cache_redis.py @@ -0,0 +1,34 @@ +from typing import Dict, Any + +from redis import Redis + +from classes.cache import HoloCache + + +class HoloCacheRedis(HoloCache): + client: Redis + + @classmethod + def from_config(cls, engine_config: Dict[str, Any]) -> Any: + raise NotImplementedError() + + def get_json(self, key: str) -> Any | None: + raise NotImplementedError() + + def get_string(self, key: str) -> str | None: + raise NotImplementedError() + + def get_object(self, key: str) -> Any | None: + raise NotImplementedError() + + def set_json(self, key: str, value: Any) -> None: + raise NotImplementedError() + + def set_string(self, key: str, value: str) -> None: + raise NotImplementedError() + + def set_object(self, key: str, value: Any) -> None: + raise NotImplementedError() + + def delete(self, key: str) -> None: + raise NotImplementedError() diff --git a/classes/holo_bot.py b/classes/holo_bot.py index c2f4246..3b8aa4b 100644 --- a/classes/holo_bot.py +++ b/classes/holo_bot.py @@ -1,6 +1,21 @@ +import logging +from logging import Logger + from libbot.pycord.classes import PycordBot +from classes.cache.holo_cache_memcached import HoloCacheMemcached +from classes.cache.holo_cache_redis import HoloCacheRedis +from modules.cache_utils import create_cache_client + +logger: Logger = logging.getLogger(__name__) + class HoloBot(PycordBot): + cache: HoloCacheMemcached | HoloCacheRedis | None = None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + + def _set_cache_engine(self) -> None: + if "cache" in self.config and self.config["cache"]["type"] is not None: + self.cache = create_cache_client(self.config, self.config["cache"]["type"]) diff --git a/classes/holo_user.py b/classes/holo_user.py index d78b720..a4c2506 100644 --- a/classes/holo_user.py +++ b/classes/holo_user.py @@ -6,7 +6,9 @@ from bson import ObjectId from discord import User, Member from libbot.utils import config_get from pymongo.results import InsertOneResult +from typing_extensions import deprecated +from classes.cache import HoloCache from errors import UserNotFoundError from modules.database import col_warnings, col_users @@ -29,33 +31,41 @@ class HoloUser: @classmethod async def from_user( - cls, user: User | Member, allow_creation: bool = True + cls, + user: User | Member, + allow_creation: bool = True, + cache: HoloCache | None = None, ) -> "HoloUser": """Get an object that has a proper binding between Discord ID and database ### Args: * `user` (User | Member): Object from which an ID can be extracted + * `allow_creation` (bool, optional): Whether to allow creation of a new user record if none found. Defaults to True. + * `cache` (HoloCache | None, optional): Cache engine to get the cache from ### Raises: * `UserNotFoundError`: User with such ID does not seem to exist in database """ - db_entry: Dict[str, Any] | None = await col_users.find_one({"user": user.id}) + 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: Dict[str, Any] | None = await col_users.find_one({"id": user.id}) if db_entry is None: if not allow_creation: raise UserNotFoundError(user=user, user_id=user.id) - db_entry = { - "user": user.id, - "custom_role": None, - "custom_channel": None, - } + db_entry = HoloUser.get_defaults(user.id) insert_result: InsertOneResult = await col_users.insert_one(db_entry) - db_entry["_id"] = insert_result.inserted_id() + db_entry["_id"] = insert_result.inserted_id - db_entry["id"] = db_entry.pop("user") + if cache is not None: + cache.set_json(f"user_{user.id}", db_entry) return cls(**db_entry) @@ -63,23 +73,28 @@ class HoloUser: async def from_id(cls, user_id: int) -> "HoloUser": return NotImplemented + # TODO Deprecate and remove warnings + @deprecated("Warnings are deprecated") async def get_warnings(self) -> int: """Get number of warnings user has ### Returns: * `int`: Number of warnings """ - warns: Dict[str, Any] | None = await col_warnings.find_one({"user": self.id}) + warns: Dict[str, Any] | None = await col_warnings.find_one({"id": self.id}) return 0 if warns is None else warns["warns"] - async def warn(self, count: int = 1, reason: str = "Not provided") -> None: + # TODO Deprecate and remove warnings + @deprecated("Warnings are deprecated") + async def warn(self, count: int = 1, reason: str = "Reason not provided") -> None: """Warn and add count to warns number ### Args: * `count` (int, optional): Count of warnings to be added. Defaults to 1. + * `reason` (int, optional): Count of warnings to be added. Defaults to 1. """ - warns: Dict[str, Any] | None = await col_warnings.find_one({"user": self.id}) + warns: Dict[str, Any] | None = await col_warnings.find_one({"id": self.id}) if warns is not None: await col_warnings.update_one( @@ -87,16 +102,17 @@ class HoloUser: {"$set": {"warns": warns["warns"] + count}}, ) else: - await col_warnings.insert_one(document={"user": self.id, "warns": count}) + await col_warnings.insert_one(document={"id": self.id, "warns": count}) logger.info("User %s was warned %s times due to: %s", self.id, count, reason) - async def _set(self, key: str, value: Any) -> None: + async def _set(self, key: str, value: Any, cache: HoloCache | None = None) -> None: """Set attribute data and save it into the database ### Args: * `key` (str): Attribute to be changed * `value` (Any): Value to set + * `cache` (HoloCache | None, optional): Cache engine to write the update into """ if not hasattr(self, key): raise AttributeError() @@ -107,40 +123,99 @@ class HoloUser: {"_id": self._id}, {"$set": {key: value}}, upsert=True ) - logger.info("Set attribute %s of user %s to %s", key, self.id, value) + self._update_cache(cache) - async def _remove(self, key: str) -> None: + logger.info("Set attribute '%s' of user %s to '%s'", key, self.id, value) + + async def _remove(self, key: str, cache: HoloCache | None = None) -> None: """Remove attribute data and save it into the database ### Args: * `key` (str): Attribute to be removed + * `cache` (HoloCache | None, optional): Cache engine to write the update into """ if not hasattr(self, key): raise AttributeError() - setattr(self, key, None) + default_value: Any = HoloUser.get_default_value(key) + + setattr(self, key, default_value) await col_users.update_one( - {"_id": self._id}, {"$unset": {key: None}}, upsert=True + {"_id": self._id}, {"$set": {key: default_value}}, upsert=True ) - logger.info("Removed attribute %s of user %s", key, self.id) + self._update_cache(cache) - async def set_custom_channel(self, channel_id: int) -> None: - await self._set("custom_channel", channel_id) + logger.info("Removed attribute '%s' of user %s", key, self.id) - async def set_custom_role(self, role_id: int) -> None: - await self._set("custom_role", role_id) + def _get_cache_key(self) -> str: + return f"user_{self.id}" - async def remove_custom_channel(self) -> None: - await self._remove("custom_channel") + def _update_cache(self, cache: HoloCache | None = None) -> None: + if cache is None: + return - async def remove_custom_role(self) -> None: - await self._remove("custom_role") + user_dict: Dict[str, Any] = self._to_dict() - async def purge(self) -> None: - """Completely remove user data from database. Will not remove transactions logs and warnings.""" + 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: HoloCache | None = 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, + "custom_role": None, + "custom_channel": None, + } + + @staticmethod + def get_default_value(key: str) -> Any: + if key not in HoloUser.get_defaults(): + raise KeyError(f"There's no default value for key '{key}' in HoloUser") + + return HoloUser.get_defaults()[key] + + def _to_dict(self) -> Dict[str, Any]: + return { + "_id": self._id, + "id": self.id, + "custom_role": self.custom_role, + "custom_channel": self.custom_channel, + } + + async def set_custom_channel( + self, channel_id: int, cache: HoloCache | None = None + ) -> None: + await self._set("custom_channel", channel_id, cache=cache) + + async def set_custom_role( + self, role_id: int, cache: HoloCache | None = None + ) -> None: + await self._set("custom_role", role_id, cache=cache) + + async def remove_custom_channel(self, cache: HoloCache | None = None) -> None: + await self._remove("custom_channel", cache=cache) + + async def remove_custom_role(self, cache: HoloCache | None = None) -> None: + await self._remove("custom_role", cache=cache) + + async def purge(self, cache: HoloCache | None = None) -> None: + """Completely remove user data from database. Will not remove transactions logs and warnings. + + ### Args: + * `cache` (HoloCache | None, optional): Cache engine to write the update into + """ await col_users.delete_one({"_id": self._id}) + self._delete_cache(cache) @staticmethod async def is_moderator(member: User | Member) -> bool: diff --git a/cogs/analytics.py b/cogs/analytics.py index af5f951..5a452b9 100644 --- a/cogs/analytics.py +++ b/cogs/analytics.py @@ -56,7 +56,7 @@ class Analytics(commands.Cog): # Insert entry into the database await col_analytics.insert_one( { - "user": message.author.id, + "user_id": message.author.id, "channel": message.channel.id, "content": message.content, "stickers": stickers, diff --git a/cogs/custom_channels.py b/cogs/custom_channels.py index 6074c12..1adf48a 100644 --- a/cogs/custom_channels.py +++ b/cogs/custom_channels.py @@ -47,7 +47,9 @@ class CustomChannels(commands.Cog): Command to create a custom channel for a user. """ - holo_user_ctx: HoloUser = await HoloUser.from_user(ctx.user) + holo_user_ctx: HoloUser = await HoloUser.from_user( + ctx.user, cache=self.client.cache + ) # Return if the user is using the command outside of a guild if not hasattr(ctx.author, "guild"): @@ -100,7 +102,9 @@ class CustomChannels(commands.Cog): manage_channels=True, ) - await holo_user_ctx.set_custom_channel(created_channel.id) + await holo_user_ctx.set_custom_channel( + created_channel.id, cache=self.client.cache + ) await ctx.respond( embed=Embed( @@ -136,7 +140,9 @@ class CustomChannels(commands.Cog): Command to change properties of a custom channel. """ - holo_user_ctx: HoloUser = await HoloUser.from_user(ctx.user) + holo_user_ctx: HoloUser = await HoloUser.from_user( + ctx.user, cache=self.client.cache + ) custom_channel: TextChannel | None = ds_utils.get( ctx.guild.channels, id=holo_user_ctx.custom_channel @@ -182,7 +188,9 @@ class CustomChannels(commands.Cog): """Command /customchannel remove [] Command to remove a custom channel. Requires additional confirmation.""" - holo_user_ctx: HoloUser = await HoloUser.from_user(ctx.user) + holo_user_ctx: HoloUser = await HoloUser.from_user( + ctx.user, cache=self.client.cache + ) # Return if the user does not have a custom channel if holo_user_ctx.custom_channel is None: @@ -211,7 +219,7 @@ class CustomChannels(commands.Cog): color=Color.FAIL, ) ) - await holo_user_ctx.remove_custom_channel() + await holo_user_ctx.remove_custom_channel(cache=self.client.cache) return # Return if the confirmation is missing @@ -227,7 +235,7 @@ class CustomChannels(commands.Cog): await custom_channel.delete(reason="Власник запросив видалення") - await holo_user_ctx.remove_custom_channel() + await holo_user_ctx.remove_custom_channel(cache=self.client.cache) try: await ctx.respond( diff --git a/cogs/data.py b/cogs/data.py index b5af890..b2b0e96 100644 --- a/cogs/data.py +++ b/cogs/data.py @@ -14,7 +14,6 @@ from libbot.utils import config_get, json_write from classes.holo_bot import HoloBot from classes.holo_user import HoloUser from enums import Color -from modules.database import col_users from modules.utils_sync import guild_name logger: Logger = logging.getLogger(__name__) @@ -90,7 +89,7 @@ class Data(commands.Cog): { "id": member.id, "nick": member.nick, - "username": f"{member.name}#{member.discriminator}", + "username": member.name, "bot": member.bot, } ) @@ -101,6 +100,7 @@ class Data(commands.Cog): await ctx.respond(file=File(Path(f"tmp/{uuid}"), filename="users.json")) + # TODO Deprecate this command @data.command( name="migrate", description="Мігрувати всіх користувачів до бази", @@ -164,20 +164,7 @@ class Data(commands.Cog): if member.bot: continue - if (await col_users.find_one({"user": member.id})) is None: - user: Dict[str, Any] = {} - defaults: Dict[str, Any] = await config_get("user", "defaults") - - user["user"] = member.id - - for key in defaults: - user[key] = defaults[key] - - await col_users.insert_one(document=user) - - logging.info( - "Added DB record for user %s during migration", member.id - ) + await HoloUser.from_user(member, cache=self.client.cache) await ctx.respond( embed=Embed( diff --git a/cogs/logger.py b/cogs/logger.py index 88eb8f2..4ba34ad 100644 --- a/cogs/logger.py +++ b/cogs/logger.py @@ -1,6 +1,5 @@ import logging from logging import Logger -from typing import Dict, Any from discord import Member, Message, TextChannel, MessageType from discord import utils as ds_utils @@ -8,6 +7,7 @@ from discord.ext import commands from libbot.utils import config_get from classes.holo_bot import HoloBot +from classes.holo_user import HoloUser from modules.database import col_users logger: Logger = logging.getLogger(__name__) @@ -25,16 +25,7 @@ class Logger(commands.Cog): and (message.author.bot is False) and (message.author.system is False) ): - if (await col_users.find_one({"user": message.author.id})) is None: - user: Dict[str, Any] = {} - defaults: Dict[str, Any] = await config_get("user", "defaults") - - user["user"] = message.author.id - - for key in defaults: - user[key] = defaults[key] - - await col_users.insert_one(document=user) + await HoloUser.from_user(message.author, cache=self.client.cache) if ( (message.type == MessageType.thread_created) @@ -69,30 +60,21 @@ class Logger(commands.Cog): id=await config_get("rules", "channels", "text"), ) - if welcome_chan is None: - logger.warning("Could not find a welcome channel by its id") - if ( (member != self.client.user) and (member.bot is False) and (member.system is False) ): - await welcome_chan.send( - content=(await config_get("welcome", "messages")).format( - mention=member.mention, rules=rules_chan.mention + if welcome_chan is not None and rules_chan is not None: + await welcome_chan.send( + content=(await config_get("welcome", "messages")).format( + mention=member.mention, rules=rules_chan.mention + ) ) - ) + else: + logger.warning("Could not find a welcome and/or rules channel by id") - if (await col_users.find_one({"user": member.id})) is None: - user: Dict[str, Any] = {} - defaults: Dict[str, Any] = await config_get("user", "defaults") - - user["user"] = member.id - - for key in defaults: - user[key] = defaults[key] - - await col_users.insert_one(document=user) + await HoloUser.from_user(member, cache=self.client.cache) def setup(client: HoloBot) -> None: diff --git a/config_example.json b/config_example.json index b04f2dd..4e0008a 100644 --- a/config_example.json +++ b/config_example.json @@ -22,16 +22,20 @@ "port": 27017, "name": "holo_discord" }, + "cache": { + "type": null, + "memcached": { + "uri": "127.0.0.1:11211" + }, + "redis": { + "uri": "127.0.0.1:6379" + } + }, "logging": { "size": 512, "location": "logs" }, - "defaults": { - "user": { - "custom_role": null, - "custom_channel": null - } - }, + "defaults": {}, "categories": { "custom_channels": 0 }, diff --git a/main.py b/main.py index d186f18..7740609 100644 --- a/main.py +++ b/main.py @@ -12,8 +12,15 @@ from classes.holo_bot import HoloBot from modules.migrator import migrate_database from modules.scheduler import scheduler +if not Path("config.json").exists(): + print( + "Config file is missing: Make sure the configuration file 'config.json' is in place.", + flush=True, + ) + sys.exit() + logging.basicConfig( - level=logging.INFO, + level=logging.INFO if not config_get("debug") else logging.DEBUG, format="%(name)s.%(funcName)s | %(levelname)s | %(message)s", datefmt="[%X]", ) @@ -41,12 +48,6 @@ with contextlib.suppress(ImportError): def main() -> None: - if not Path("config.json").exists(): - logger.error( - "Config file is missing: Make sure the configuration file 'config.json' is in place." - ) - sys.exit() - # Perform migration if command line argument was provided if args.migrate: logger.info("Performing migrations...") diff --git a/migrations/202502092252.py b/migrations/202502092252.py new file mode 100644 index 0000000..9e5603a --- /dev/null +++ b/migrations/202502092252.py @@ -0,0 +1,63 @@ +import logging +from logging import Logger + +from libbot.utils import config_set, config_delete +from mongodb_migrations.base import BaseMigration + +logger: Logger = logging.getLogger(__name__) + + +class Migration(BaseMigration): + def upgrade(self): + try: + config_set( + "cache", + { + "type": None, + "memcached": {"uri": "127.0.0.1:11211"}, + "redis": {"uri": "127.0.0.1:6379"}, + }, + *[], + ) + except Exception as exc: + logger.error( + "Could not upgrade the config during migration '%s' due to: %s", + __name__, + exc, + ) + + self.db.users.update_many( + {"user": {"$exists": True}}, + {"$rename": {"user": "id"}}, + ) + self.db.analytics.update_many( + {"user": {"$exists": True}}, + {"$rename": {"user": "user_id"}}, + ) + self.db.warnings.update_many( + {"user": {"$exists": True}}, + {"$rename": {"user": "user_id"}}, + ) + + def downgrade(self): + try: + config_delete("cache", *[]) + except Exception as exc: + logger.error( + "Could not downgrade the config during migration '%s' due to: %s", + __name__, + exc, + ) + + self.db.users.update_many( + {"id": {"$exists": True}}, + {"$rename": {"id": "user"}}, + ) + self.db.analytics.update_many( + {"user": {"$exists": True}}, + {"$rename": {"user_id": "user"}}, + ) + self.db.warnings.update_many( + {"user": {"$exists": True}}, + {"$rename": {"user_id": "user"}}, + ) diff --git a/modules/cache_utils.py b/modules/cache_utils.py new file mode 100644 index 0000000..02d01fb --- /dev/null +++ b/modules/cache_utils.py @@ -0,0 +1,29 @@ +from typing import Dict, Any, Literal + +from classes.cache.holo_cache_memcached import HoloCacheMemcached +from classes.cache.holo_cache_redis import HoloCacheRedis + + +def create_cache_client( + config: Dict[str, Any], + engine: Literal["memcached", "redis"] | None = None, +) -> HoloCacheMemcached | HoloCacheRedis: + if engine not in ["memcached", "redis"] or engine is None: + raise KeyError( + f"Incorrect cache engine provided. Expected 'memcached' or 'redis', got '{engine}'" + ) + + if "cache" not in config or engine not in config["cache"]: + raise KeyError( + f"Cache configuration is invalid. Please check if all keys are set (engine: '{engine}')" + ) + + match engine: + case "memcached": + return HoloCacheMemcached.from_config(config["cache"][engine]) + case "redis": + return HoloCacheRedis.from_config(config["cache"][engine]) + case _: + raise KeyError( + f"Cache implementation for the engine '{engine}' is not present." + ) diff --git a/modules/database.py b/modules/database.py index f78d61a..fb62da4 100644 --- a/modules/database.py +++ b/modules/database.py @@ -38,3 +38,6 @@ sync_db: Database = db_client_sync.get_database(name=db_config["name"]) sync_col_users: Collection = sync_db.get_collection("users") sync_col_warnings: Collection = sync_db.get_collection("warnings") sync_col_analytics: Collection = sync_db.get_collection("analytics") + +# Update indexes +sync_col_users.create_index(["id"], unique=True) diff --git a/requirements.txt b/requirements.txt index 5388ca9..f3937a3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,5 +7,7 @@ apscheduler>=3.10.0 async_pymongo==0.1.11 libbot[speed,pycord]==4.0.2 mongodb-migrations==1.3.1 +pymemcache~=4.0.0 +redis~=5.2.1 ujson~=5.10.0 WaifuPicsPython==0.2.0 \ No newline at end of file diff --git a/validation/analytics.json b/validation/analytics.json index d07f175..3f13cc4 100644 --- a/validation/analytics.json +++ b/validation/analytics.json @@ -1,14 +1,14 @@ { "$jsonSchema": { "required": [ - "user", + "user_id", "channel", "content", "stickers", "attachments" ], "properties": { - "user": { + "user_id": { "bsonType": "long", "description": "Discord ID of user" }, @@ -17,7 +17,10 @@ "description": "Discord ID of a channel" }, "content": { - "bsonType": ["null", "string"], + "bsonType": [ + "null", + "string" + ], "description": "Text of the message" }, "stickers": { @@ -40,7 +43,7 @@ "format": { "bsonType": "array" }, - "user": { + "user_id": { "bsonType": "string" } } diff --git a/validation/users.json b/validation/users.json index 34e0a8b..972d2e4 100644 --- a/validation/users.json +++ b/validation/users.json @@ -1,12 +1,12 @@ { "$jsonSchema": { "required": [ - "user", + "id", "custom_role", "custom_channel" ], "properties": { - "user": { + "id": { "bsonType": "long", "description": "Discord ID of user" }, diff --git a/validation/warnings.json b/validation/warnings.json index aa80293..23596ea 100644 --- a/validation/warnings.json +++ b/validation/warnings.json @@ -1,11 +1,11 @@ { "$jsonSchema": { "required": [ - "user", + "user_id", "warnings" ], "properties": { - "user": { + "user_id": { "bsonType": "long", "description": "Discord ID of user" }, From a54ef39e53cbdd8196a3e81934ff6356e9aed6c2 Mon Sep 17 00:00:00 2001 From: profitroll Date: Sun, 9 Feb 2025 23:37:44 +0100 Subject: [PATCH 2/7] Added basic Redis support --- classes/cache/holo_cache_memcached.py | 31 +++--------- classes/cache/holo_cache_redis.py | 71 ++++++++++++++++++++++++--- classes/holo_bot.py | 1 + modules/cache_utils.py | 24 +++++++++ 4 files changed, 96 insertions(+), 31 deletions(-) diff --git a/classes/cache/holo_cache_memcached.py b/classes/cache/holo_cache_memcached.py index a764d24..fe53e65 100644 --- a/classes/cache/holo_cache_memcached.py +++ b/classes/cache/holo_cache_memcached.py @@ -2,10 +2,9 @@ import logging from logging import Logger from typing import Dict, Any -from bson import ObjectId from pymemcache import Client -from ujson import dumps, loads +from modules.cache_utils import string_to_json, json_to_string from . import HoloCache logger: Logger = logging.getLogger(__name__) @@ -17,6 +16,8 @@ class HoloCacheMemcached(HoloCache): def __init__(self, client: Client): self.client = client + logger.info("Initialized Memcached for caching") + @classmethod def from_config(cls, engine_config: Dict[str, Any]) -> "HoloCacheMemcached": if "uri" not in engine_config: @@ -39,7 +40,7 @@ class HoloCacheMemcached(HoloCache): logger.error("Could not get json cache key '%s' due to: %s", key, exc) return None - return None if result is None else self._string_to_json(result) + return None if result is None else string_to_json(result) def get_string(self, key: str) -> str | None: try: @@ -56,13 +57,13 @@ class HoloCacheMemcached(HoloCache): logger.error("Could not get string cache key '%s' due to: %s", key, exc) return None + # TODO Implement binary deserialization def get_object(self, key: str) -> Any | None: - # TODO Implement binary deserialization raise NotImplementedError() def set_json(self, key: str, value: Any) -> None: try: - self.client.set(key, self._json_to_string(value)) + self.client.set(key, json_to_string(value)) logger.debug("Set json cache key '%s'", key) except Exception as exc: logger.error("Could not set json cache key '%s' due to: %s", key, exc) @@ -76,8 +77,8 @@ class HoloCacheMemcached(HoloCache): logger.error("Could not set string cache key '%s' due to: %s", key, exc) return None + # TODO Implement binary serialization def set_object(self, key: str, value: Any) -> None: - # TODO Implement binary serialization raise NotImplementedError() def delete(self, key: str) -> None: @@ -86,21 +87,3 @@ class HoloCacheMemcached(HoloCache): logger.debug("Deleted cache key '%s'", key) except Exception as exc: logger.error("Could not delete cache key '%s' due to: %s", key, exc) - - @staticmethod - def _json_to_string(json_object: Any) -> str: - if isinstance(json_object, dict) and "_id" in json_object: - json_object["_id"] = str(json_object["_id"]) - - return dumps( - json_object, ensure_ascii=False, indent=0, escape_forward_slashes=False - ) - - @staticmethod - def _string_to_json(json_string: str) -> Any: - json_object: Any = loads(json_string) - - if isinstance(json_object, dict) and "_id" in json_object: - json_object["_id"] = ObjectId(json_object["_id"]) - - return json_object diff --git a/classes/cache/holo_cache_redis.py b/classes/cache/holo_cache_redis.py index a22a9aa..5b30bb5 100644 --- a/classes/cache/holo_cache_redis.py +++ b/classes/cache/holo_cache_redis.py @@ -1,34 +1,91 @@ -from typing import Dict, Any +import logging +from logging import Logger +from typing import Dict, Any, List from redis import Redis from classes.cache import HoloCache +from modules.cache_utils import string_to_json, json_to_string + +logger: Logger = logging.getLogger(__name__) class HoloCacheRedis(HoloCache): client: Redis + def __init__(self, client: Redis): + self.client = client + + logger.info("Initialized Redis for caching") + @classmethod def from_config(cls, engine_config: Dict[str, Any]) -> Any: - raise NotImplementedError() + if "uri" not in engine_config: + raise KeyError( + "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" + ) + + uri_split: List[str] = engine_config["uri"].split(":") + + return cls(Redis(host=uri_split[0], port=int(uri_split[1]))) def get_json(self, key: str) -> Any | None: - raise NotImplementedError() + try: + result: Any | None = self.client.get(key) + + logger.debug( + "Got json cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + except Exception as exc: + logger.error("Could not get json cache key '%s' due to: %s", key, exc) + return None + + return None if result is None else string_to_json(result) def get_string(self, key: str) -> str | None: - raise NotImplementedError() + try: + result: str | None = self.client.get(key) + logger.debug( + "Got string cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + + return result + except Exception as exc: + logger.error("Could not get string cache key '%s' due to: %s", key, exc) + return None + + # TODO Implement binary deserialization def get_object(self, key: str) -> Any | None: raise NotImplementedError() def set_json(self, key: str, value: Any) -> None: - raise NotImplementedError() + try: + self.client.set(key, json_to_string(value)) + logger.debug("Set json cache key '%s'", key) + except Exception as exc: + logger.error("Could not set json cache key '%s' due to: %s", key, exc) + return None def set_string(self, key: str, value: str) -> None: - raise NotImplementedError() + try: + self.client.set(key, value) + logger.debug("Set string cache key '%s'", key) + except Exception as exc: + logger.error("Could not set string cache key '%s' due to: %s", key, exc) + return None + # TODO Implement binary serialization def set_object(self, key: str, value: Any) -> None: raise NotImplementedError() def delete(self, key: str) -> None: - raise NotImplementedError() + try: + self.client.delete(key) + logger.debug("Deleted cache key '%s'", key) + except Exception as exc: + logger.error("Could not delete cache key '%s' due to: %s", key, exc) diff --git a/classes/holo_bot.py b/classes/holo_bot.py index 3b8aa4b..26310dc 100644 --- a/classes/holo_bot.py +++ b/classes/holo_bot.py @@ -15,6 +15,7 @@ class HoloBot(PycordBot): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self._set_cache_engine() def _set_cache_engine(self) -> None: if "cache" in self.config and self.config["cache"]["type"] is not None: diff --git a/modules/cache_utils.py b/modules/cache_utils.py index 02d01fb..85aa629 100644 --- a/modules/cache_utils.py +++ b/modules/cache_utils.py @@ -1,5 +1,9 @@ +from copy import deepcopy from typing import Dict, Any, Literal +from bson import ObjectId +from ujson import dumps, loads + from classes.cache.holo_cache_memcached import HoloCacheMemcached from classes.cache.holo_cache_redis import HoloCacheRedis @@ -27,3 +31,23 @@ def create_cache_client( raise KeyError( f"Cache implementation for the engine '{engine}' is not present." ) + + +def json_to_string(json_object: Any) -> str: + json_object_copy: Any = deepcopy(json_object) + + if isinstance(json_object_copy, dict) and "_id" in json_object_copy: + json_object_copy["_id"] = str(json_object_copy["_id"]) + + return dumps( + json_object_copy, ensure_ascii=False, indent=0, escape_forward_slashes=False + ) + + +def string_to_json(json_string: str) -> Any: + json_object: Any = loads(json_string) + + if "_id" in json_object: + json_object["_id"] = ObjectId(json_object["_id"]) + + return json_object From cda570eb374e0d4e3c8247d1a7d92895769e7634 Mon Sep 17 00:00:00 2001 From: profitroll Date: Sun, 9 Feb 2025 23:40:09 +0100 Subject: [PATCH 3/7] Changed default URI for Redis --- classes/cache/holo_cache_redis.py | 6 ++---- config_example.json | 2 +- migrations/202502092252.py | 2 +- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/classes/cache/holo_cache_redis.py b/classes/cache/holo_cache_redis.py index 5b30bb5..7d6a125 100644 --- a/classes/cache/holo_cache_redis.py +++ b/classes/cache/holo_cache_redis.py @@ -1,6 +1,6 @@ import logging from logging import Logger -from typing import Dict, Any, List +from typing import Dict, Any from redis import Redis @@ -25,9 +25,7 @@ class HoloCacheRedis(HoloCache): "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" ) - uri_split: List[str] = engine_config["uri"].split(":") - - return cls(Redis(host=uri_split[0], port=int(uri_split[1]))) + return cls(Redis.from_url(engine_config["uri"])) def get_json(self, key: str) -> Any | None: try: diff --git a/config_example.json b/config_example.json index 4e0008a..146e50f 100644 --- a/config_example.json +++ b/config_example.json @@ -28,7 +28,7 @@ "uri": "127.0.0.1:11211" }, "redis": { - "uri": "127.0.0.1:6379" + "uri": "redis://127.0.0.1:6379/0" } }, "logging": { diff --git a/migrations/202502092252.py b/migrations/202502092252.py index 9e5603a..b7430d6 100644 --- a/migrations/202502092252.py +++ b/migrations/202502092252.py @@ -15,7 +15,7 @@ class Migration(BaseMigration): { "type": None, "memcached": {"uri": "127.0.0.1:11211"}, - "redis": {"uri": "127.0.0.1:6379"}, + "redis": {"uri": "redis://127.0.0.1:6379/0"}, }, *[], ) From 9e9b6bc7dc84267f0f7444a2b251ee381b847673 Mon Sep 17 00:00:00 2001 From: profitroll Date: Sun, 9 Feb 2025 23:45:18 +0100 Subject: [PATCH 4/7] Moved create_cache_client() to cache_manager --- classes/holo_bot.py | 2 +- modules/cache_manager.py | 29 +++++++++++++++++++++++++++++ modules/cache_utils.py | 30 +----------------------------- 3 files changed, 31 insertions(+), 30 deletions(-) create mode 100644 modules/cache_manager.py diff --git a/classes/holo_bot.py b/classes/holo_bot.py index 26310dc..881e92d 100644 --- a/classes/holo_bot.py +++ b/classes/holo_bot.py @@ -5,7 +5,7 @@ from libbot.pycord.classes import PycordBot from classes.cache.holo_cache_memcached import HoloCacheMemcached from classes.cache.holo_cache_redis import HoloCacheRedis -from modules.cache_utils import create_cache_client +from modules.cache_manager import create_cache_client logger: Logger = logging.getLogger(__name__) diff --git a/modules/cache_manager.py b/modules/cache_manager.py new file mode 100644 index 0000000..02d01fb --- /dev/null +++ b/modules/cache_manager.py @@ -0,0 +1,29 @@ +from typing import Dict, Any, Literal + +from classes.cache.holo_cache_memcached import HoloCacheMemcached +from classes.cache.holo_cache_redis import HoloCacheRedis + + +def create_cache_client( + config: Dict[str, Any], + engine: Literal["memcached", "redis"] | None = None, +) -> HoloCacheMemcached | HoloCacheRedis: + if engine not in ["memcached", "redis"] or engine is None: + raise KeyError( + f"Incorrect cache engine provided. Expected 'memcached' or 'redis', got '{engine}'" + ) + + if "cache" not in config or engine not in config["cache"]: + raise KeyError( + f"Cache configuration is invalid. Please check if all keys are set (engine: '{engine}')" + ) + + match engine: + case "memcached": + return HoloCacheMemcached.from_config(config["cache"][engine]) + case "redis": + return HoloCacheRedis.from_config(config["cache"][engine]) + case _: + raise KeyError( + f"Cache implementation for the engine '{engine}' is not present." + ) diff --git a/modules/cache_utils.py b/modules/cache_utils.py index 85aa629..2f95204 100644 --- a/modules/cache_utils.py +++ b/modules/cache_utils.py @@ -1,37 +1,9 @@ from copy import deepcopy -from typing import Dict, Any, Literal +from typing import Any from bson import ObjectId from ujson import dumps, loads -from classes.cache.holo_cache_memcached import HoloCacheMemcached -from classes.cache.holo_cache_redis import HoloCacheRedis - - -def create_cache_client( - config: Dict[str, Any], - engine: Literal["memcached", "redis"] | None = None, -) -> HoloCacheMemcached | HoloCacheRedis: - if engine not in ["memcached", "redis"] or engine is None: - raise KeyError( - f"Incorrect cache engine provided. Expected 'memcached' or 'redis', got '{engine}'" - ) - - if "cache" not in config or engine not in config["cache"]: - raise KeyError( - f"Cache configuration is invalid. Please check if all keys are set (engine: '{engine}')" - ) - - match engine: - case "memcached": - return HoloCacheMemcached.from_config(config["cache"][engine]) - case "redis": - return HoloCacheRedis.from_config(config["cache"][engine]) - case _: - raise KeyError( - f"Cache implementation for the engine '{engine}' is not present." - ) - def json_to_string(json_object: Any) -> str: json_object_copy: Any = deepcopy(json_object) From 0228983d52870279dcde60dadd5002447f372f40 Mon Sep 17 00:00:00 2001 From: kku Date: Mon, 10 Feb 2025 10:15:28 +0100 Subject: [PATCH 5/7] Updated README for caching --- README.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/README.md b/README.md index 2c47899..f0e3165 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,27 @@ Mandatory keys to modify: After all of that you're good to go! Happy using :) +## Caching + +Although general database access speed is fast, caching might become helpful for +bigger servers with many bot interactions. Currently, Redis and Memcached are supported. + +Configuration happens through the config key `caching`. + +Set `caching.type` to the service of you choice ("redis" or "memcached") and then update +the URI to access the service. It's Redis' default URI format for Redis and "address:port" +for Memcached. + +Which one should I choose? + +| Service | Read/write speed | Config flexibility | +|-----------|------------------|--------------------| +| Redis | High | Very flexible | +| Memcached | Very high | Basic | + +> Performance difference between Redis and Memcached is generally quite low, so your setup +> should normally depend more on the configuration flexibility than on raw speed. + ## Docker [Experimental] As an experiment, Docker deployment option has been added. From 25f2595cf7bdd268990b9cd2d9d7a9d985c9335c Mon Sep 17 00:00:00 2001 From: kku Date: Mon, 10 Feb 2025 10:21:16 +0100 Subject: [PATCH 6/7] Added a note about a debug setting --- README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/README.md b/README.md index f0e3165..defb708 100644 --- a/README.md +++ b/README.md @@ -94,6 +94,13 @@ Which one should I choose? > Performance difference between Redis and Memcached is generally quite low, so your setup > should normally depend more on the configuration flexibility than on raw speed. +## Debugging + +There's a config key `debug` that can be set to `true` for debugging purposes. + +It should be set to `false` in production, otherwise log becomes very verbose and might +contain data that shouldn't normally be logged. + ## Docker [Experimental] As an experiment, Docker deployment option has been added. From dae89fd1ac12562a5b0ec3a3ad3229b0141e7869 Mon Sep 17 00:00:00 2001 From: kku Date: Mon, 10 Feb 2025 12:50:56 +0100 Subject: [PATCH 7/7] Removed warnings functionality --- classes/holo_user.py | 40 +++------------------------------------- modules/database.py | 2 -- validation/warnings.json | 18 ------------------ 3 files changed, 3 insertions(+), 57 deletions(-) delete mode 100644 validation/warnings.json diff --git a/classes/holo_user.py b/classes/holo_user.py index a4c2506..2d08564 100644 --- a/classes/holo_user.py +++ b/classes/holo_user.py @@ -6,11 +6,10 @@ from bson import ObjectId from discord import User, Member from libbot.utils import config_get from pymongo.results import InsertOneResult -from typing_extensions import deprecated from classes.cache import HoloCache from errors import UserNotFoundError -from modules.database import col_warnings, col_users +from modules.database import col_users logger: Logger = logging.getLogger(__name__) @@ -71,40 +70,7 @@ class HoloUser: @classmethod async def from_id(cls, user_id: int) -> "HoloUser": - return NotImplemented - - # TODO Deprecate and remove warnings - @deprecated("Warnings are deprecated") - async def get_warnings(self) -> int: - """Get number of warnings user has - - ### Returns: - * `int`: Number of warnings - """ - warns: Dict[str, Any] | None = await col_warnings.find_one({"id": self.id}) - - return 0 if warns is None else warns["warns"] - - # TODO Deprecate and remove warnings - @deprecated("Warnings are deprecated") - async def warn(self, count: int = 1, reason: str = "Reason not provided") -> None: - """Warn and add count to warns number - - ### Args: - * `count` (int, optional): Count of warnings to be added. Defaults to 1. - * `reason` (int, optional): Count of warnings to be added. Defaults to 1. - """ - warns: Dict[str, Any] | None = await col_warnings.find_one({"id": self.id}) - - if warns is not None: - await col_warnings.update_one( - {"_id": self._id}, - {"$set": {"warns": warns["warns"] + count}}, - ) - else: - await col_warnings.insert_one(document={"id": self.id, "warns": count}) - - logger.info("User %s was warned %s times due to: %s", self.id, count, reason) + raise NotImplementedError() async def _set(self, key: str, value: Any, cache: HoloCache | None = None) -> None: """Set attribute data and save it into the database @@ -209,7 +175,7 @@ class HoloUser: await self._remove("custom_role", cache=cache) async def purge(self, cache: HoloCache | None = None) -> None: - """Completely remove user data from database. Will not remove transactions logs and warnings. + """Completely remove user data from database. Only removes the user record from users collection. ### Args: * `cache` (HoloCache | None, optional): Cache engine to write the update into diff --git a/modules/database.py b/modules/database.py index fb62da4..97bf49a 100644 --- a/modules/database.py +++ b/modules/database.py @@ -29,14 +29,12 @@ db_client_sync: MongoClient = MongoClient(con_string) db: AsyncDatabase = db_client.get_database(name=db_config["name"]) col_users: AsyncCollection = db.get_collection("users") -col_warnings: AsyncCollection = db.get_collection("warnings") col_analytics: AsyncCollection = db.get_collection("analytics") # Sync declarations as a fallback sync_db: Database = db_client_sync.get_database(name=db_config["name"]) sync_col_users: Collection = sync_db.get_collection("users") -sync_col_warnings: Collection = sync_db.get_collection("warnings") sync_col_analytics: Collection = sync_db.get_collection("analytics") # Update indexes diff --git a/validation/warnings.json b/validation/warnings.json deleted file mode 100644 index 23596ea..0000000 --- a/validation/warnings.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "$jsonSchema": { - "required": [ - "user_id", - "warnings" - ], - "properties": { - "user_id": { - "bsonType": "long", - "description": "Discord ID of user" - }, - "warnings": { - "bsonType": "int", - "description": "Number of warnings on count" - } - } - } -} \ No newline at end of file