From f4f84cea67e36f12d3502864764851e589c148f8 Mon Sep 17 00:00:00 2001 From: Choco <94597336+ChocoMeow@users.noreply.github.com> Date: Fri, 9 Feb 2024 21:02:07 +0800 Subject: [PATCH] Switched from pymongo to motor for asynchronous MongoDB operations --- cogs/settings.py | 102 +++++++++++++------------------------ cogs/task.py | 2 +- function.py | 116 +++++++++++++++++++----------------------- main.py | 32 ++++++++++-- requirements.txt | 2 +- views/embedBuilder.py | 4 +- voicelink/player.py | 20 ++++++-- 7 files changed, 134 insertions(+), 144 deletions(-) diff --git a/cogs/settings.py b/cogs/settings.py index 28c4776..f92fb54 100644 --- a/cogs/settings.py +++ b/cogs/settings.py @@ -34,7 +34,6 @@ from function import ( update_settings, get_settings, get_lang, - settings as sett, time as ctime, get_aliases, cooldown_check @@ -52,15 +51,6 @@ class Settings(commands.Cog, name="settings"): def __init__(self, bot) -> None: self.bot: commands.Bot = bot self.description = "This category is only available to admin permissions on the server." - - def get_settings(self, ctx: commands.Context) -> Tuple[voicelink.Player, dict]: - player: voicelink.Player = ctx.guild.voice_client - if not player: - settings = get_settings(ctx.guild.id) - else: - settings = player.settings - - return player, settings @commands.hybrid_group( name="settings", @@ -77,7 +67,7 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def prefix(self, ctx: commands.Context, prefix: str): "Change the default prefix for message commands." - update_settings(ctx.guild.id, {"prefix": prefix}) + await update_settings(ctx.guild.id, {"$set": {"prefix": prefix}}) await ctx.send(get_lang(ctx.guild.id, "setPrefix").format(ctx.prefix, prefix)) @settings.command(name="language", aliases=get_aliases("language")) @@ -89,7 +79,7 @@ class Settings(commands.Cog, name="settings"): if language not in LANGS: return await ctx.send(get_lang(ctx.guild.id, "languageNotFound")) - update_settings(ctx.guild.id, {'lang': language}) + await update_settings(ctx.guild.id, {"$set": {'lang': language}}) await ctx.send(get_lang(ctx.guild.id, 'changedLanguage').format(language)) @language.autocomplete('language') @@ -103,17 +93,7 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def dj(self, ctx: commands.Context, role: discord.Role = None): "Set a DJ role or remove DJ role." - player: voicelink.Player = ctx.guild.voice_client - - if not role: - if player: - player.settings.pop('dj', None) - update_settings(ctx.guild.id, {'dj': ''}, mode="unset") - else: - if player: - player.settings['dj'] = role.id - update_settings(ctx.guild.id, {'dj': role.id}) - + await update_settings(ctx.guild.id, {"$set": {'dj': role.id}} if role else {"$unset": {'dj': None}}) await ctx.send(get_lang(ctx.guild.id, 'setDJ').format(f"<@&{role.id}>" if role else "None"), allowed_mentions=discord.AllowedMentions.none()) @settings.command(name="queue", aliases=get_aliases("queue")) @@ -125,11 +105,8 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def queue(self, ctx: commands.Context, mode: str): "Change to another type of queue mode." - player, settings = self.get_settings(ctx) - mode = "FairQueue" if mode.lower() == "fairqueue" else "Queue" - settings["queueType"] = mode - update_settings(ctx.guild.id, {"queueType": mode}) + await update_settings(ctx.guild.id, {"$set": {"queueType": mode}}) await ctx.send(get_lang(ctx.guild.id, "setqueue").format(mode)) @settings.command(name="247", aliases=get_aliases("247")) @@ -137,33 +114,28 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def playforever(self, ctx: commands.Context): "Toggles 24/7 mode, which disables automatic inactivity-based disconnects." - player, settings = self.get_settings(ctx) + settings = await get_settings(ctx.guild.id) toggle = settings.get('24/7', False) - settings['24/7'] = not toggle - update_settings(ctx.guild.id, {'24/7': not toggle}) - toggle = get_lang(ctx.guild.id, "enabled" if not toggle else "disabled") - await ctx.send(get_lang(ctx.guild.id, '247').format(toggle)) + await update_settings(ctx.guild.id, {"$set": {'24/7': not toggle}}) + await ctx.send(get_lang(ctx.guild.id, '247').format(get_lang(ctx.guild.id, "enabled" if not toggle else "disabled"))) @settings.command(name="bypassvote", aliases=get_aliases("bypassvote")) @commands.has_permissions(manage_guild=True) @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def bypassvote(self, ctx: commands.Context): "Toggles voting system." - player, settings = self.get_settings(ctx) + settings = await get_settings(ctx.guild.id) toggle = settings.get('votedisable', True) - settings['votedisable'] = not toggle - update_settings(ctx.guild.id, {'votedisable': not toggle}) - toggle = get_lang(ctx.guild.id, - "enabled" if not toggle else "disabled") - await ctx.send(get_lang(ctx.guild.id, 'bypassVote').format(toggle)) + await update_settings(ctx.guild.id, {"$set": {'votedisable': not toggle}}) + await ctx.send(get_lang(ctx.guild.id, 'bypassVote').format(get_lang(ctx.guild.id, "enabled" if not toggle else "disabled"))) @settings.command(name="view", aliases=get_aliases("view")) @commands.has_permissions(manage_guild=True) @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def view(self, ctx: commands.Context): "Show all the bot settings in your server." - player, settings = self.get_settings(ctx) - embed = discord.Embed(color=sett.embed_color) + settings = await get_settings(ctx.guild.id) + embed = discord.Embed(color=func.settings.embed_color) embed.set_author(name=get_lang(ctx.guild.id, 'settingsMenu').format(ctx.guild.name), icon_url=self.bot.user.display_avatar.url) if ctx.guild.icon: embed.set_thumbnail(url=ctx.guild.icon.url) @@ -203,10 +175,9 @@ class Settings(commands.Cog, name="settings"): "Set the player's volume." player: voicelink.Player = ctx.guild.voice_client if player: - player.settings['volume'] = value await player.set_volume(value, ctx.author) - update_settings(ctx.guild.id, {'volume': value}) + await update_settings(ctx.guild.id, {"$set": {'volume': value}}) await ctx.send(get_lang(ctx.guild.id, 'setVolume').format(value)) @settings.command(name="togglecontroller", aliases=get_aliases("togglecontroller")) @@ -214,40 +185,39 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def togglecontroller(self, ctx: commands.Context): "Toggles the music controller." - player, settings = self.get_settings(ctx) - toggle = settings.get('controller', True) - settings['controller'] = not toggle - if player and settings['controller'] is False: - if player.controller: - try: - await player.controller.delete() - except: - discord.ui.View.from_message(player.controller).stop() + settings = await get_settings(ctx.guild.id) + toggle = not settings.get('controller', True) - update_settings(ctx.guild.id, {'controller': not toggle}) - toggle = get_lang(ctx.guild.id, "enabled" if not toggle else "disabled") - await ctx.send(get_lang(ctx.guild.id, 'togglecontroller').format(toggle)) + player: voicelink.Player = ctx.guild.voice_client + if player and toggle is False and player.controller: + try: + await player.controller.delete() + except: + discord.ui.View.from_message(player.controller).stop() + + await update_settings(ctx.guild.id, {"$set": {'controller': toggle}}) + await ctx.send(get_lang(ctx.guild.id, 'togglecontroller').format(get_lang(ctx.guild.id, "enabled" if toggle else "disabled"))) @settings.command(name="duplicatetrack", aliases=get_aliases("duplicatetrack")) @commands.has_permissions(manage_guild=True) @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def duplicatetrack(self, ctx: commands.Context): "Toggle Vocard to prevent duplicate songs from queuing." - player, settings = self.get_settings(ctx) - toggle = settings.get('duplicateTrack', False) + settings = await get_settings(ctx.guild.id) + toggle = not settings.get('duplicateTrack', False) + player: voicelink.Player = ctx.guild.voice_client if player: - player.queue._allow_duplicate = not toggle + player.queue._allow_duplicate = toggle - update_settings(ctx.guild.id, {'duplicateTrack': not toggle}) - toggle = get_lang(ctx.guild.id, "enabled" if toggle else "disabled") - return await ctx.send(get_lang(ctx.guild.id, "toggleDuplicateTrack").format(toggle)) + await update_settings(ctx.guild.id, {"$set": {'duplicateTrack': toggle}}) + return await ctx.send(get_lang(ctx.guild.id, "toggleDuplicateTrack").format(get_lang(ctx.guild.id, "disabled" if toggle else "enabled"))) @settings.command(name="customcontroller", aliases=get_aliases("customcontroller")) @commands.has_permissions(manage_guild=True) @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def customcontroller(self, ctx: commands.Context): "Customizes music controller embeds." - player, settings = self.get_settings(ctx) + settings = await get_settings(ctx.guild.id) controller_settings = settings.get("default_controller", func.settings.controller) view = EmbedBuilderView(ctx, controller_settings.get("embeds").copy()) @@ -258,13 +228,11 @@ class Settings(commands.Cog, name="settings"): @commands.dynamic_cooldown(cooldown_check, commands.BucketType.guild) async def controllermsg(self, ctx: commands.Context): "Toggles to send a message when clicking the button in the music controller." - player, settings = self.get_settings(ctx) - toggle = settings.get('controller_msg', True) + settings = await get_settings(ctx.guild.id) + toggle = not settings.get('controller_msg', True) - settings['controller_msg'] = not toggle - update_settings(ctx.guild.id, {'controller_msg': not toggle}) - toggle = get_lang(ctx.guild.id, "enabled" if not toggle else "disabled") - await ctx.send(get_lang(ctx.guild.id, 'toggleControllerMsg').format(toggle)) + await update_settings(ctx.guild.id, {"$set": {'controller_msg': toggle}}) + await ctx.send(get_lang(ctx.guild.id, 'toggleControllerMsg').format(get_lang(ctx.guild.id, "enabled" if toggle else "disabled"))) @app_commands.command(name="debug") async def debug(self, interaction: discord.Interaction): diff --git a/cogs/task.py b/cogs/task.py index 9620189..bd445fb 100644 --- a/cogs/task.py +++ b/cogs/task.py @@ -110,7 +110,7 @@ class Task(commands.Cog): @tasks.loop(hours=12.0) async def cache_cleaner(self): - func.GUILD_SETTINGS.clear() + func.SETTINGS_BUFFER.clear() func.PLAYLIST_NAME.clear() errorFile = func.gen_report() diff --git a/function.py b/function.py index 63115da..3e358ce 100644 --- a/function.py +++ b/function.py @@ -1,75 +1,37 @@ -import discord -import json -import aiohttp -import os +import discord, json, os from discord.ext import commands from datetime import datetime from time import strptime from io import BytesIO -from pymongo import MongoClient from typing import Optional, Union, Any from addons import Settings, TOKENS +from motor.motor_asyncio import ( + AsyncIOMotorClient, + AsyncIOMotorCollection, +) + ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) if not os.path.exists(os.path.join(ROOT_DIR, "settings.json")): raise Exception("Settings file not set!") -#-------------- API Clients -------------- -tokens: TOKENS = TOKENS() - -if not (tokens.mongodb_name and tokens.mongodb_url): - raise Exception("MONGODB_NAME and MONGODB_URL can't not be empty in .env") - -try: - mongodb = MongoClient(host=tokens.mongodb_url, serverSelectionTimeoutMS=5000) - mongodb.server_info() - if tokens.mongodb_name not in mongodb.list_database_names(): - raise Exception(f"{tokens.mongodb_name} does not exist in your mongoDB!") - print("Successfully connected to MongoDB!") - -except Exception as e: - raise Exception("Not able to connect MongoDB! Reason:", e) - -SETTINGS_DB = mongodb[tokens.mongodb_name]['Settings'] -PLAYLISTS_DB = mongodb[tokens.mongodb_name]['Playlist'] - #--------------- Cache Var --------------- +tokens: TOKENS = TOKENS() settings: Settings + +MONGO_DB: AsyncIOMotorClient +SETTINGS_DB: AsyncIOMotorCollection +PLAYLISTS_DB: AsyncIOMotorCollection + ERROR_LOGS: dict[int, dict[int, str]] = {} #Stores error that not a Voicelink Exception LANGS: dict[str, dict[str, str]] = {} #Stores all the languages in ./langs -GUILD_SETTINGS: dict[int, dict[str, Any]] = {} #Cache guild language +SETTINGS_BUFFER: dict[int, dict[str, Any]] = {} #Cache guild language LOCAL_LANGS: dict[str, dict[str, str]] = {} #Stores all the localization languages in ./local_langs PLAYLIST_NAME: dict[str, list[str]] = {} #Cache the user's playlist name #-------------- Vocard Functions -------------- -def get_settings(guild_id:int) -> dict: - settings = GUILD_SETTINGS.get(guild_id, None) - if not settings: - settings = SETTINGS_DB.find_one({"_id":guild_id}) - if not settings: - SETTINGS_DB.insert_one({"_id":guild_id}) - - GUILD_SETTINGS[guild_id] = settings or {} - return settings - -def update_settings(guild_id:int, data: dict, mode="set") -> bool: - settings = get_settings(guild_id) - - for key, value in data.items(): - if settings.get(key) != value: - match mode: - case "set": - GUILD_SETTINGS[guild_id][key] = value - case "unset": - GUILD_SETTINGS[guild_id].pop(key) - case _: - return False - - result = SETTINGS_DB.update_one({"_id":guild_id}, {f"${mode}":data}) - return result.modified_count > 0 - def open_json(path: str) -> dict: try: with open(os.path.join(ROOT_DIR, path), encoding="utf8") as json_file: @@ -88,19 +50,12 @@ def update_json(path: str, new_data: dict) -> None: json.dump(data, json_file, indent=4) def get_lang(guild_id:int, key:str) -> str: - lang = get_settings(guild_id).get("lang", "EN") + lang = SETTINGS_BUFFER.get(guild_id).get("lang", "EN") if lang in LANGS and not LANGS[lang]: LANGS[lang] = open_json(os.path.join("langs", f"{lang}.json")) return LANGS.get(lang, {}).get(key, "Language pack not found!") -def init() -> None: - global settings - - json = open_json("settings.json") - if json is not None: - settings = Settings(json) - def langs_setup() -> None: for language in os.listdir(os.path.join(ROOT_DIR, "langs")): if language.endswith('.json'): @@ -167,6 +122,41 @@ def get_aliases(name: str) -> list: def check_roles() -> tuple[str, int, int]: return 'Normal', 5, 500 +async def get_settings(guild_id:int) -> dict[str, Any]: + settings = SETTINGS_BUFFER.get(guild_id, None) + if not settings: + settings = await SETTINGS_DB.find_one({"_id": guild_id}) + if not settings: + await SETTINGS_DB.insert_one({"_id": guild_id}) + + settings = SETTINGS_BUFFER[guild_id] = settings or {} + return settings + +async def update_settings(guild_id: int, data: dict[str, dict[str, Any]]) -> bool: + settings = await get_settings(guild_id) + for mode, action in data.items(): + for key, value in action.items(): + cursors = key.split(".") + + nested_data = settings + for c in cursors[:-1]: + nested_data = nested_data.setdefault(c, {}) + + if mode == "$set": + try: + nested_data[cursors[-1]] = value + except TypeError: + nested_data[int(cursors[-1])] = value + + elif mode == "$unset": + nested_data.pop(cursors[-1], None) + + else: + return False + + result = await SETTINGS_DB.update_one({"_id": guild_id}, data) + return result.modified_count > 0 + async def create_account(ctx: Union[commands.Context, discord.Interaction]) -> None: author = ctx.author if isinstance(ctx, commands.Context) else ctx.user if not author: @@ -187,12 +177,12 @@ async def create_account(ctx: Union[commands.Context, discord.Interaction]) -> N await view.wait() if view.value: try: - PLAYLISTS_DB.insert_one({'_id':author.id, 'playlist': {'200':{'tracks':[],'perms':{ 'read': [], 'write':[], 'remove': []},'name':'Favourite', 'type':'playlist' }},'inbox':[] }) + await PLAYLISTS_DB.insert_one({'_id':author.id, 'playlist': {'200':{'tracks':[],'perms':{ 'read': [], 'write':[], 'remove': []},'name':'Favourite', 'type':'playlist' }},'inbox':[] }) except: pass async def get_playlist(user_id:int, dType:str=None, dId:str=None) -> bool: - user = PLAYLISTS_DB.find_one({"_id":user_id}, {"_id": 0}) + user = await PLAYLISTS_DB.find_one({"_id":user_id}, {"_id": 0}) if not user: return None if dType: @@ -204,9 +194,9 @@ async def get_playlist(user_id:int, dType:str=None, dId:str=None) -> bool: async def update_playlist(user_id:int, data:dict, *, mode:str="set", update_cache: bool=False) -> None: if update_cache: PLAYLIST_NAME.pop(str(user_id), None) - result = PLAYLISTS_DB.update_one({"_id":user_id}, {f"${mode}": data}) + result = await PLAYLISTS_DB.update_one({"_id":user_id}, {f"${mode}": data}) return result.modified_count > 0 async def update_inbox(user_id:int, data:dict) -> bool: - result = PLAYLISTS_DB.update_one({"_id":user_id}, {"$push":{'inbox':data}}) + result = await PLAYLISTS_DB.update_one({"_id":user_id}, {"$push":{'inbox':data}}) return result.modified_count > 0 \ No newline at end of file diff --git a/main.py b/main.py index a094e18..ad7bbac 100644 --- a/main.py +++ b/main.py @@ -8,10 +8,10 @@ import function as func from discord.ext import commands from web import IPCServer +from motor.motor_asyncio import AsyncIOMotorClient from datetime import datetime from voicelink import VoicelinkException - -func.init() +from addons import Settings class Translator(discord.app_commands.Translator): async def load(self): @@ -48,8 +48,29 @@ class Vocard(commands.Bot): await self.process_commands(message) - async def setup_hook(self): + async def connect_db(self) -> None: + if not ((db_name := func.tokens.mongodb_name) and (db_url := func.tokens.mongodb_url)): + raise Exception("MONGODB_NAME and MONGODB_URL can't not be empty in settings.json") + + try: + func.MONGO_DB = AsyncIOMotorClient(host=db_url, serverSelectionTimeoutMS=5000) + await func.MONGO_DB.server_info() + if db_name not in await func.MONGO_DB.list_database_names(): + raise Exception(f"{db_name} does not exist in your mongoDB!") + print("Successfully connected to MongoDB!") + + except Exception as e: + raise Exception("Not able to connect MongoDB! Reason:", e) + + func.SETTINGS_DB = func.MONGO_DB[db_name]["Settings"] + func.PLAYLISTS_DB = func.MONGO_DB[db_name]["Playlist"] + + async def setup_hook(self) -> None: func.langs_setup() + + await self.connect_db() + + # Loading all the module in `cogs` folder for module in os.listdir(func.ROOT_DIR + '/cogs'): if module.endswith('.py'): try: @@ -113,7 +134,6 @@ class Vocard(commands.Bot): pass class CommandCheck(discord.app_commands.CommandTree): - async def interaction_check(self, interaction: discord.Interaction, /) -> bool: if not interaction.guild: await interaction.response.send_message("This command can only be used in guilds!") @@ -122,9 +142,11 @@ class CommandCheck(discord.app_commands.CommandTree): return await super().interaction_check(interaction) async def get_prefix(bot, message: discord.Message): - settings = func.get_settings(message.guild.id) + settings = await func.get_settings(message.guild.id) return settings.get("prefix", func.settings.bot_prefix) + +# Setup the bot object intents = discord.Intents.default() intents.message_content = True if func.settings.bot_prefix else False member_cache = discord.MemberCacheFlags( diff --git a/requirements.txt b/requirements.txt index ce159b4..41d4156 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ discord.py==2.3.2 -pymongo==4.5.0 +motor==3.3.2 dnspython==2.2.1 tldextract==3.2.1 validators==0.18.2 diff --git a/views/embedBuilder.py b/views/embedBuilder.py index f2ed8a5..405dd0e 100644 --- a/views/embedBuilder.py +++ b/views/embedBuilder.py @@ -307,9 +307,9 @@ class EmbedBuilderView(discord.ui.View): @discord.ui.button(label="Apply", style=discord.ButtonStyle.green, row=1) async def apply(self, interaction: discord.Interaction, button: discord.ui.Button): - func.update_settings( + await func.update_settings( interaction.guild_id, - {"default_controller": {"embeds": self.data}}, + {"$set": {"default_controller.embeds": self.data}}, ) await self.on_timeout() diff --git a/voicelink/player.py b/voicelink/player.py index 125f773..7087eab 100644 --- a/voicelink/player.py +++ b/voicelink/player.py @@ -68,8 +68,11 @@ async def connect_channel(ctx: Union[commands.Context, Interaction], channel: Vo if check.connect == False or check.speak == False: raise VoicelinkException(func.get_lang(ctx.guild.id, 'noPermission')) - player: Player = await channel.connect(cls=Player( - ctx.bot if isinstance(ctx, commands.Context) else ctx.client, channel, ctx + settings = await func.get_settings(channel.guild.id) + player: Player = await channel.connect( + cls=Player( + ctx.bot if isinstance(ctx, commands.Context) else ctx.client, + channel, ctx, settings )) await player.send_ws({"op": "createPlayer", "members_id": [member.id for member in channel.members]}) @@ -95,6 +98,7 @@ class Player(VoiceProtocol): client: Optional[Client] = None, channel: Optional[VoiceChannel] = None, ctx: Union[commands.Context, Interaction] = None, + settings: dict[str, Any] = None ): self.client: Client = client self._bot: Client = client @@ -104,7 +108,7 @@ class Player(VoiceProtocol): self._guild = channel.guild if channel else None self._ipc_connection: bool = False - self.settings: dict = func.get_settings(ctx.guild.id) + self.settings: dict = settings self.joinTime: float = round(time.time()) self._volume: int = self.settings.get('volume', 100) self.queue: Queue = eval(self.settings.get("queueType", "Queue"))(self.settings.get("maxQueue", func.settings.max_queue), self.settings.get("duplicateTrack", True), self.get_msg) @@ -403,8 +407,14 @@ class Player(VoiceProtocol): return False async def teardown(self): - timeNow = round(time.time()) - func.update_settings(self.guild.id, {"lastActice": timeNow, "playTime": round(self.settings.get("playTime", 0) + ((timeNow - self.joinTime) / 60), 2)}) + await func.update_settings( + self.guild.id, + {"$set": { + "lastActice": (timeNow := round(time.time())), + "playTime": round(self.settings.get("playTime", 0) + ((timeNow - self.joinTime) / 60), 2) + }} + ) + if self.is_ipc_connected: await self.send_ws({"op": "playerClose"})