Switched from pymongo to motor for asynchronous MongoDB operations

2024-02-09 21:02:07 +08:00
parent abb53937d0
commit f4f84cea67
7 changed files with 134 additions and 144 deletions
--- 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):
--- 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()
--- 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
--- 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(
--- 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
--- 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()
--- 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"})