diff --git a/Dockerfile b/Dockerfile index 6009219..6e797bf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -24,4 +24,4 @@ USER appuser COPY . . -ENTRYPOINT ["python", "main.py"] \ No newline at end of file +ENTRYPOINT ["python", "main.py", "--migrate"] \ No newline at end of file diff --git a/README.md b/README.md index a0afa5e..defb708 100644 --- a/README.md +++ b/README.md @@ -30,9 +30,27 @@ 3. Clone the repository: `git clone https://git.end-play.xyz/HoloUA/Discord.git` 4. `cd Discord` -5. Install dependencies: +5. Create a virtual environment: + `python -m venv .venv` or `virtualenv .venv` +6. Activate the virtual environment: + Windows: `.venv\Scripts\activate.bat` + Linux/macOS: `.venv/bin/activate` +7. Install the dependencies: `python -m pip install -r requirements.txt` -6. Run the bot with `python main.py` after completing the [configuration](#Configuration) +8. Run the bot with `python main.py` after completing the [configuration](#Configuration) + +## Upgrading with Git + +1. Go to the bot's directory +2. `git pull` +3. Activate the virtual environment: + Windows: `.venv\Scripts\activate.bat` + Linux/macOS: `.venv/bin/activate` +4. Update the dependencies: + `python -m pip install -r requirements.txt` +5. First start after the upgrade must initiate the migration: + `python main.py --migrate` +6. Now the bot is up to date and the next run will not require `--migrate` anymore ## Configuration @@ -55,6 +73,34 @@ 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. + +## 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. 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..fe53e65 --- /dev/null +++ b/classes/cache/holo_cache_memcached.py @@ -0,0 +1,89 @@ +import logging +from logging import Logger +from typing import Dict, Any + +from pymemcache import Client + +from modules.cache_utils import string_to_json, json_to_string +from . import HoloCache + +logger: Logger = logging.getLogger(__name__) + + +class HoloCacheMemcached(HoloCache): + client: Client + + 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: + 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 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 + + # TODO Implement binary deserialization + def get_object(self, key: str) -> Any | None: + raise NotImplementedError() + + def set_json(self, key: str, value: Any) -> None: + 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: + 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: + 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/cache/holo_cache_redis.py b/classes/cache/holo_cache_redis.py new file mode 100644 index 0000000..7d6a125 --- /dev/null +++ b/classes/cache/holo_cache_redis.py @@ -0,0 +1,89 @@ +import logging +from logging import Logger +from typing import Dict, Any + +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: + if "uri" not in engine_config: + raise KeyError( + "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" + ) + + return cls(Redis.from_url(engine_config["uri"])) + + def get_json(self, key: str) -> Any | None: + 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: + 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: + 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: + 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: + 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 c2f4246..881e92d 100644 --- a/classes/holo_bot.py +++ b/classes/holo_bot.py @@ -1,6 +1,22 @@ +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_manager 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) + self._set_cache_engine() + + 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 508b8db..a4c2506 100644 --- a/classes/holo_user.py +++ b/classes/holo_user.py @@ -5,71 +5,114 @@ from typing import Any, Dict 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, sync_col_users, sync_col_warnings, col_users +from modules.database import col_warnings, col_users logger: Logger = logging.getLogger(__name__) class HoloUser: - def __init__(self, user: User | Member | int) -> None: + def __init__( + self, + _id: ObjectId, + id: int, + custom_role: int | None, + custom_channel: int | None, + ) -> None: + self._id: ObjectId = _id + + self.id: int = id + self.custom_role: int | None = custom_role + self.custom_channel: int | None = custom_channel + + @classmethod + async def from_user( + 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 | int): Object from which ID can be extracted + * `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 """ + if cache is not None: + cached_entry: Dict[str, Any] | None = cache.get_json(f"user_{user.id}") - self.id: int = user if not hasattr(user, "id") else user.id + if cached_entry is not None: + return cls(**cached_entry) - jav_user: Dict[str, Any] | None = sync_col_users.find_one({"user": self.id}) + db_entry: Dict[str, Any] | None = await col_users.find_one({"id": user.id}) - if jav_user is None: - raise UserNotFoundError(user=user, user_id=self.id) + if db_entry is None: + if not allow_creation: + raise UserNotFoundError(user=user, user_id=user.id) - self.db_id: ObjectId = jav_user["_id"] + db_entry = HoloUser.get_defaults(user.id) - self.customrole: int | None = jav_user["customrole"] - self.customchannel: int | None = jav_user["customchannel"] - self.warnings: int = self.warns() + insert_result: InsertOneResult = await col_users.insert_one(db_entry) - def warns(self) -> int: + 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) + + @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 = sync_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=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( - {"_id": self.db_id}, + {"_id": self._id}, {"$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: - """Set attribute data and save it into database + 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() @@ -77,10 +120,102 @@ class HoloUser: setattr(self, key, value) await col_users.update_one( - {"_id": self.db_id}, {"$set": {key: value}}, upsert=True + {"_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) + + 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() + + default_value: Any = HoloUser.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: HoloCache | None = 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: 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 328ca08..1adf48a 100644 --- a/cogs/custom_channels.py +++ b/cogs/custom_channels.py @@ -25,7 +25,7 @@ class CustomChannels(commands.Cog): @commands.Cog.listener() async def on_guild_channel_delete(self, channel: GuildChannel) -> None: await col_users.find_one_and_update( - {"customchannel": channel.id}, {"$set": {"customchannel": None}} + {"custom_channel": channel.id}, {"$set": {"custom_channel": None}} ) custom_channel_group: SlashCommandGroup = SlashCommandGroup( @@ -47,7 +47,9 @@ class CustomChannels(commands.Cog): Command to create a custom channel for a user. """ - holo_user_ctx: HoloUser = HoloUser(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"): @@ -62,7 +64,7 @@ class CustomChannels(commands.Cog): return # Return if the user already has a custom channel - if holo_user_ctx.customchannel is not None: + if holo_user_ctx.custom_channel is not None: await ctx.defer(ephemeral=True) await ctx.respond( embed=Embed( @@ -80,7 +82,7 @@ class CustomChannels(commands.Cog): reason=f"Користувач {guild_name(ctx.user)} отримав власний приватний канал", category=ds_utils.get( ctx.author.guild.categories, - id=await config_get("customchannels", "categories"), + id=await config_get("custom_channels", "categories"), ), ) @@ -100,7 +102,9 @@ class CustomChannels(commands.Cog): manage_channels=True, ) - await holo_user_ctx.set("customchannel", created_channel.id) + await holo_user_ctx.set_custom_channel( + created_channel.id, cache=self.client.cache + ) await ctx.respond( embed=Embed( @@ -136,10 +140,12 @@ class CustomChannels(commands.Cog): Command to change properties of a custom channel. """ - holo_user_ctx: HoloUser = HoloUser(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.customchannel + ctx.guild.channels, id=holo_user_ctx.custom_channel ) # Return if the channel was not found @@ -182,10 +188,12 @@ class CustomChannels(commands.Cog): """Command /customchannel remove [] Command to remove a custom channel. Requires additional confirmation.""" - holo_user_ctx: HoloUser = HoloUser(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.customchannel is None: + if holo_user_ctx.custom_channel is None: await ctx.defer(ephemeral=True) await ctx.respond( embed=Embed( @@ -199,7 +207,7 @@ class CustomChannels(commands.Cog): await ctx.defer() custom_channel: TextChannel | None = ds_utils.get( - ctx.guild.channels, id=holo_user_ctx.customchannel + ctx.guild.channels, id=holo_user_ctx.custom_channel ) # Return if the channel was not found @@ -211,7 +219,7 @@ class CustomChannels(commands.Cog): color=Color.FAIL, ) ) - await holo_user_ctx.set("customchannel", None) + 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.set("customchannel", None) + 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 66904a4..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,22 +25,13 @@ 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) and (message.channel is not None) and ( - await col_users.count_documents({"customchannel": message.channel.id}) + await col_users.count_documents({"custom_channel": message.channel.id}) > 0 ) ): @@ -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 a3cca7e..146e50f 100644 --- a/config_example.json +++ b/config_example.json @@ -22,18 +22,22 @@ "port": 27017, "name": "holo_discord" }, + "cache": { + "type": null, + "memcached": { + "uri": "127.0.0.1:11211" + }, + "redis": { + "uri": "redis://127.0.0.1:6379/0" + } + }, "logging": { "size": 512, "location": "logs" }, - "defaults": { - "user": { - "customrole": null, - "customchannel": null - } - }, + "defaults": {}, "categories": { - "customchannels": 0 + "custom_channels": 0 }, "channels": { "text": { diff --git a/main.py b/main.py index 6a1a038..7740609 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,7 @@ +import contextlib import logging import sys +from argparse import ArgumentParser from logging import Logger from pathlib import Path @@ -7,30 +9,49 @@ from discord import LoginFailure, Intents from libbot.utils import config_get 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]", ) logger: Logger = logging.getLogger(__name__) -try: - import uvloop # type: ignore +# Declare the parser that retrieves the command line arguments +parser = ArgumentParser( + prog="HoloUA Discord", + description="Discord bot for the HoloUA community.", +) + +# Add a switch argument --migrate to be parsed... +parser.add_argument("--migrate", action="store_true") + +# ...and parse the arguments we added +args = parser.parse_args() + +# Try to import the module that improves performance +# and ignore errors when module is not installed +with contextlib.suppress(ImportError): + import uvloop uvloop.install() -except ImportError: - pass 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...") + migrate_database() intents: Intents = Intents().all() client: HoloBot = HoloBot(intents=intents, scheduler=scheduler) diff --git a/migrations/202412272043.py b/migrations/202412272043.py new file mode 100644 index 0000000..07c3aed --- /dev/null +++ b/migrations/202412272043.py @@ -0,0 +1,79 @@ +import logging +from logging import Logger + +from libbot.utils import config_get, config_set, config_delete +from mongodb_migrations.base import BaseMigration + +logger: Logger = logging.getLogger(__name__) + + +class Migration(BaseMigration): + def upgrade(self): + try: + # Categories + config_set( + "custom_channels", + config_get("customchannels", "categories"), + "categories", + ) + config_delete("customchannels", "categories") + + # User defaults + config_delete( + "user", + "defaults", + ) + except Exception as exc: + logger.error( + "Could not upgrade the config during migration '%s' due to: %s", + __name__, + exc, + ) + + self.db.users.update_many( + {"customchannel": {"$exists": True}}, + {"$rename": {"customchannel": "custom_channel"}}, + ) + self.db.users.update_many( + {"customrole": {"$exists": True}}, + {"$rename": {"customrole": "custom_role"}}, + ) + + def downgrade(self): + try: + # Categories + config_set( + "customchannels", + config_get("custom_channels", "categories"), + "categories", + ) + config_delete("custom_channels", "categories") + + # User defaults + config_set( + "customrole", + None, + "defaults", + "user", + ) + config_set( + "customchannel", + None, + "defaults", + "user", + ) + except Exception as exc: + logger.error( + "Could not downgrade the config during migration '%s' due to: %s", + __name__, + exc, + ) + + self.db.users.update_many( + {"custom_channel": {"$exists": True}}, + {"$rename": {"custom_channel": "customchannel"}}, + ) + self.db.users.update_many( + {"custom_role": {"$exists": True}}, + {"$rename": {"custom_role": "customrole"}}, + ) diff --git a/migrations/202502092252.py b/migrations/202502092252.py new file mode 100644 index 0000000..b7430d6 --- /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": "redis://127.0.0.1:6379/0"}, + }, + *[], + ) + 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_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 new file mode 100644 index 0000000..2f95204 --- /dev/null +++ b/modules/cache_utils.py @@ -0,0 +1,25 @@ +from copy import deepcopy +from typing import Any + +from bson import ObjectId +from ujson import dumps, loads + + +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 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/modules/migrator.py b/modules/migrator.py new file mode 100644 index 0000000..2eedb71 --- /dev/null +++ b/modules/migrator.py @@ -0,0 +1,22 @@ +from typing import Any, Mapping + +from libbot.utils import config_get +from mongodb_migrations.cli import MigrationManager +from mongodb_migrations.config import Configuration + + +def migrate_database() -> None: + """Apply migrations from folder `migrations/` to the database""" + database_config: Mapping[str, Any] = config_get("database") + + manager_config = Configuration( + { + "mongo_host": database_config["host"], + "mongo_port": database_config["port"], + "mongo_database": database_config["name"], + "mongo_username": database_config["user"], + "mongo_password": database_config["password"], + } + ) + manager = MigrationManager(manager_config) + manager.run() diff --git a/requirements.txt b/requirements.txt index 05effea..f3937a3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,9 @@ requests>=2.32.2 aiofiles~=24.1.0 apscheduler>=3.10.0 async_pymongo==0.1.11 -libbot[speed,pycord]==4.0.0 -typing-extensions~=4.12.2 +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 d76120c..972d2e4 100644 --- a/validation/users.json +++ b/validation/users.json @@ -1,22 +1,27 @@ { "$jsonSchema": { "required": [ - "user", - "customrole", - "customchannel" - + "id", + "custom_role", + "custom_channel" ], "properties": { - "user": { + "id": { "bsonType": "long", "description": "Discord ID of user" }, - "customrole": { - "bsonType": ["null", "long"], + "custom_role": { + "bsonType": [ + "null", + "long" + ], "description": "Discord ID of custom role or 'null' if not set" }, - "customchannel": { - "bsonType": ["null", "long"], + "custom_channel": { + "bsonType": [ + "null", + "long" + ], "description": "Discord ID of custom channel or 'null' if not set" } } diff --git a/validation/warnings.json b/validation/warnings.json index e28a79a..23596ea 100644 --- a/validation/warnings.json +++ b/validation/warnings.json @@ -1,15 +1,15 @@ { "$jsonSchema": { "required": [ - "user", - "warns" + "user_id", + "warnings" ], "properties": { - "user": { + "user_id": { "bsonType": "long", "description": "Discord ID of user" }, - "warns": { + "warnings": { "bsonType": "int", "description": "Number of warnings on count" }