Better database getter and cleanup

Author: gccody

launcher.py

@@ -73,9 +73,14 @@ async def on_command_error(ctx: Interaction, exc) -> None:
             else:
                 raise exc
 
-    async def exit_handler():
+    def exit_handler():
         bot.db.commit()
-        await bot.close()
+        # try:
+        #     loop = asyncio.get_running_loop()
+        # except RuntimeError:
+        #     loop = asyncio.new_event_loop()
+        # loop.run_until_complete(bot.close())
+        # asyncio.run(bot.close())
 
     atexit.register(exit_handler)
     await bot.start(bot.config.token)

lib/bot/__init__.py

@@ -2,27 +2,22 @@
 import json
 import os
 from datetime import datetime
-from http.client import HTTPException
 import random
 from typing import Any
 
 import discord
 import psutil
 import tzlocal
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
-from discord import Guild, app_commands, Interaction, VoiceChannel, TextChannel, Message
-from discord.app_commands import errors, Command
+from discord import Guild, Interaction, VoiceChannel
+from discord.app_commands import Command
 from discord.embeds import Embed
 from discord.errors import NotFound
-from discord.ext.commands import AutoShardedBot, Context
-from discord import AutoShardedClient
-from discord.ext.commands._types import BotT
-from discord.ext.commands.errors import CommandNotFound, BadArgument, MissingRequiredArgument, CommandOnCooldown, \
-    MissingPermissions, BotMissingPermissions
+from discord.ext.commands import AutoShardedBot
+from discord.ext.commands.errors import CommandNotFound, BadArgument
 
 from lib.config import Config
 from lib.db import DB
-from lib.errors import BotNotReady
 from lib.progress import Progress, Timer
 from lib.urban import UrbanDictionary
 
@@ -56,6 +51,7 @@ def __init__(self, shards: list[int], version: str):
     async def on_connect(self):
         await self.tree.sync()
         while not self.shards_ready:
+            print('Shards Not Ready')
             await asyncio.sleep(.5)
         self.register_guilds()
         print(f"Signed into {self.user.display_name}#{self.user.discriminator}")
@@ -73,7 +69,7 @@ async def on_interaction(self, ctx: Interaction):
         if isinstance(ctx.command, Command):
             name_list, options = self.get_name(ctx.data, [])
             name = " ".join(name_list)
-            self.db.execute(f"""INSERT INTO commands VALUES (?,?,?,?,?)""", name, ctx.guild_id,
+            self.db.run(f"""INSERT INTO commands VALUES (?,?,?,?,?)""", name, ctx.guild_id,
                             ctx.user.id, json.dumps(options), datetime.now().timestamp())
 
     @staticmethod
@@ -100,21 +96,22 @@ async def send(ctx: Interaction, *args, **kwargs):
     def register_guilds(self):
         progress = Progress('Registering guilds', len(self.guilds))
         for guild in self.guilds:
-            self.db.execute("""INSERT OR IGNORE INTO guilds VALUES (?,?,?)""", guild.id, None, None)
+            self.db.run("""INSERT OR IGNORE INTO guilds VALUES (?,?,?)""", guild.id, None, None)
             progress.next()
 
-        for guild in self.db.records("""SELECT * FROM guilds"""):
+        for guild in self.db.get.guilds():
             if not self.get_guild(guild.id):
-                self.db.execute("""DELETE FROM guilds WHERE id=?""", guild.id)
+                self.db.run("""DELETE FROM guilds WHERE id=?""", guild.id)
         self.db.commit()
         self.ready = True
+        print('End')
         self.tasks.start()
 
     async def on_guild_join(self, guild: Guild):
-        self.db.execute("""INSERT OR IGNORE INTO guilds VALUES (?,?,?)""", guild.id, None, None)
+        self.db.run("""INSERT OR IGNORE INTO guilds VALUES (?,?,?)""", guild.id, None, None)
 
     async def on_guild_remove(self, guild: Guild):
-        self.db.execute("""DELETE FROM guilds WHERE id=?""", guild.id)
+        self.db.run("""DELETE FROM guilds WHERE id=?""", guild.id)
 
     def get_name(self, data: Any, groups: list[str]):
         if isinstance(data, dict):

lib/cache/__init__.py

@@ -0,0 +1,4 @@
+from lib.levels import Player
+
+class Cache:
+    levels: dict[str, Player]

lib/cogs/Announcment.py

@@ -33,15 +33,15 @@ async def announce_enable(self, ctx: Interaction, channel: Optional[discord.Text
         if not channel.permissions_for(await ctx.guild.fetch_member(self.bot.user.id)).send_messages:
             return await ctx.response.send_message(
                 Embed(title=f">>> I can't send messages in channel: {channel.name} ({channel.id})", colour=0xff0000))
-        self.bot.db.execute("""UPDATE guilds SET a_id=? WHERE id=?""", channel.id, ctx.guild.id)
+        self.bot.db.run("""UPDATE guilds SET a_id=? WHERE id=?""", channel.id, ctx.guild.id)
         print('here at send')
         await self.bot.send(ctx, embed=Embed(
             title=f">>> Successfully Enabled Announcements in {channel.mention} ({channel.id})"))
 
     @announce.command(name='disable', description='Disable bot announcements')
     @has_guild_permissions(manage_guild=True)
     async def announce_disable(self, ctx: Interaction):
-        self.bot.db.execute("""UPDATE guilds SET a_id=?""")
+        self.bot.db.run("""UPDATE guilds SET a_id=?""")
         await self.bot.send(ctx,
                             embed=Embed(title=f">>> Successfully Disabled Announcements"))
 

lib/cogs/Handlers.py

@@ -20,7 +20,7 @@ def __init__(self, bot):
         # if isinstance(ctx.command, discord.app_commands.Command):
         #     name_list, options = self.bot.get_name(ctx.data, [])
         #     name = " ".join(name_list)
-            # self.bot.db.execute(f"""INSERT INTO commands VALUES (?,?,?,?,?)""", name, ctx.guild_id,
+            # self.bot.db.run(f"""INSERT INTO commands VALUES (?,?,?,?,?)""", name, ctx.guild_id,
             #                     ctx.user.id, json.dumps(options), datetime.now().timestamp())
 
 

lib/cogs/Mod.py

@@ -81,13 +81,11 @@ async def ban_add(self, ctx: Interaction, member: discord.Member, delete_message
         if isinstance(member, Member):
             await member.ban(delete_message_days=self.days[delete_messages] if delete_messages else 0,
                              reason=f"{reason} | {ctx.user.display_name}#{ctx.user.discriminator}")
-            print('Banned Member')
         # Member is not a part of the guild
         elif isinstance(member, User):
             await ctx.guild.ban(user=member, delete_message_days=self.days[delete_messages] if delete_messages else 0,
                                 reason=f"{reason} | {ctx.user.display_name}#{ctx.user.discriminator}")
-            print('Banned User')
-        self.bot.db.execute("""INSERT INTO bans VALUES (?,?,?,?,?)""", ctx.guild_id, member.id, ctx.user.id, reason,
+        self.bot.db.run("""INSERT INTO bans VALUES (?,?,?,?,?)""", ctx.guild_id, member.id, ctx.user.id, reason,
                             datetime.now().timestamp())
         embed = Embed(title=f"✅ | {member.display_name}#{member.discriminator} ({member.id}) has been banned")
         await ctx.response.send_message(embed=embed)
@@ -109,7 +107,7 @@ async def ban_remove(self, ctx: Interaction, member: discord.Member, reason: Opt
         # Member is not part of the guild
         elif isinstance(member, User):
             await ctx.guild.unban(user=member, reason=reason)
-        self.bot.db.execute("""DELETE FROM bans WHERE guild_id=? AND user_id=? """, ctx.guild_id, member.id)
+        self.bot.db.run("""DELETE FROM bans WHERE guild_id=? AND user_id=? """, ctx.guild_id, member.id)
         await ctx.response.send_message(
             embed=Embed(title=f"✅ | {member.display_name}#{member.discriminator} ({member.id}) has been unbanned"))
 
@@ -143,7 +141,7 @@ async def massban_add(self, ctx: Interaction, str_members: str, delete_messages:
                 continue
             await member.ban(delete_message_days=self.days[delete_messages],
                              reason=f"{reason} | {ctx.user.display_name}#{ctx.user.discriminator}")
-            self.bot.db.execute("""INSERT INTO bans VALUES (?,?,?,?,?)""", ctx.guild_id, member.id, ctx.user.id, reason,
+            self.bot.db.run("""INSERT INTO bans VALUES (?,?,?,?,?)""", ctx.guild_id, member.id, ctx.user.id, reason,
                                 datetime.now().timestamp())
             success.append(f"{member.display_name}#{member.discriminator} ({id})")
         embed = Embed()
@@ -179,7 +177,7 @@ async def massban_remove(self, ctx: Interaction, str_members: str, reason: Optio
                 continue
 
             await member.unban(reason=f"{reason} | {ctx.user.display_name}#{ctx.user.discriminator}")
-            self.bot.db.execute("""DELETE FROM bans WHERE guild_id=? AND user_id=? """, ctx.guild_id, member.id)
+            self.bot.db.run("""DELETE FROM bans WHERE guild_id=? AND user_id=? """, ctx.guild_id, member.id)
             success.append(f"{member.display_name}#{member.discriminator} ({id})")
         embed = Embed()
         s_users = "\n".join(success)

lib/cogs/Youtube.py

@@ -36,8 +36,7 @@ async def youtube_notify(self, ctx: Interaction, handle: str, message: Optional[
             elif res.status not in range(200, 300):
                 return await ctx.response.send_message(embed=Embed(title=':x: | Something went wrong', colour=0xff0000))
 
-            user = self.bot.db.record("""SELECT * FROM youtube WHERE handle = ? AND guild_id = ?""", handle,
-                                      ctx.guild_id)
+            user = self.bot.db.get.youtube(handle=handle, guild_id=ctx.guild_id)
             if user:
                 return await ctx.response.send_message(embed=Embed(title=':x: | Already watching for new videos', colour=0xff0000))
 
@@ -48,7 +47,7 @@ async def youtube_notify(self, ctx: Interaction, handle: str, message: Optional[
             latest_video_url = "https://www.youtube.com/watch?v=" + re.search('(?<="videoId":").*?(?=")', html).group()
         except AttributeError:
             latest_video_url = ""
-        self.bot.db.execute("""INSERT INTO youtube VALUES (?,?,?,?)""", handle, ctx.guild_id,
+        self.bot.db.run("""INSERT INTO youtube VALUES (?,?,?,?)""", handle, ctx.guild_id,
                             message if message else f"@{handle} has uploaded a video", latest_video_url)
         await ctx.response.send_message(embed=Embed(title=f'✅ | Waiting for @{handle} to post more videos'))
         self.bot.tasks.add_job(id=f"{handle}|{ctx.guild_id}", args=(handle, message, ctx.guild_id,), trigger='interval',

lib/db/__init__.py

@@ -1,14 +1,52 @@
 import re
-import sys
-import time
 from sqlite3 import Connection, Cursor, connect, Error
+from typing import List, Dict, Optional
 
 from lib.progress import Progress
 
+
+class Get:
+    def __init__(self, parent):
+        self.parent = parent
+        self._table_names = [list(table["table"].keys())[0] for table in parent.table_info]
+
+    def __getattribute__(self, name):
+        table_names = super().__getattribute__('_table_names')
+        if name in table_names:
+            _get_rows = super().__getattribute__('_get_rows')
+            return lambda **kwargs: _get_rows(name, **kwargs)
+        return super().__getattribute__(name)
+
+    def _get_rows(self, table_name: str, count: Optional[int] = None, **kwargs: Dict) -> List[object]:
+        if not kwargs:
+            query = f"SELECT * FROM {table_name}"
+        else:
+            query = f"SELECT * FROM {table_name} WHERE "
+            query += " AND ".join([f"{key}='{value}'" for key, value in kwargs.items()])
+        if count:
+            query += f" LIMIT {count}"
+            self.parent.cur.execute(query)
+        rows = self.parent.cur.fetchall()
+        for table in self.parent.table_info:
+            if list(table["table"].keys())[0] == table_name:
+                keys = list(table["table"][table_name].keys())
+                class_name = table_name
+                row_class = type(class_name, (object,), {})
+                return [self._create_row(row_class, keys, row) for row in rows]
+
+    def _create_row(self, row_class, keys, row):
+        new_row = row_class()
+        for i in range(len(keys)):
+            setattr(new_row, keys[i], row[i])
+        return new_row
+
+
 class DB:
     DB_PATH = "./data/database.db"
     cxn: Connection
     cur: Cursor
+    table_info: List[Dict[str, Dict[str, Dict[str, str | List[str]]]]]
+    table_names: List[str]
 
     table_info: list[dict[str, dict[str, dict[str, str | list[str]]]]] = [
         {
@@ -81,9 +119,12 @@ class DB:
         }
     ]
 
+    table_names: list[str] = []
+
     def __init__(self):
         self.cxn = connect(self.DB_PATH, check_same_thread=False)
         self.cur = self.cxn.cursor()
+        self.get: Get = Get(self)
 
     def update_table(self, table: str, data: str):
         try:
@@ -102,6 +143,7 @@ def build(self) -> None:
         for table in self.table_info:
             table_data = table['table']
             for table_name, columns in table_data.items():
+                self.table_names.append(table_name)
                 create_table_query = f"CREATE TABLE IF NOT EXISTS {table_name}("
                 constraints = table.get('constraints')
                 for column_name, data_type in columns.items():
@@ -145,24 +187,5 @@ def commit(self) -> None:
     def close(self) -> None:
         self.cxn.close()
 
-    def field(self, command, *values):
-        self.cur.execute(command, tuple(values))
-
-        if (fetch := self.cur.fetchone()) is not None:
-            return fetch[0]
-
-    def record(self, command, *values):
-        self.cur.execute(command, tuple(values))
-
-        return self.cur.fetchone()
-
-    def records(self, command, *values) -> list:
-        self.cur.execute(command, tuple(values))
-
-        return self.cur.fetchall()
-
-    def execute(self, command, *values) -> None:
-        self.cur.execute(command, tuple(values))
-
-    def multiexec(self, command, valueset) -> None:
-        self.cur.execute(command, valueset)
+    def run(self, command, *values) -> None:
+        self.cur.execute(command, tuple(values))

lib/levels/__init__.py

@@ -0,0 +1,11 @@
+class Player:
+    xp: int = 0
+    guild_id: str
+    user_id: str
+    level: int = 0
+
+    def __init__(self, guild_id: str, user_id: str, xp: int = 0):
+        self.xp = xp
+        self.guild_id = guild_id
+        self.user_id = user_id
+        self.level = 0

template.txt

@@ -16,3 +16,18 @@ class Class(Cog):
 
 async def setup(bot):
     await bot.add_cog(Class(bot))
+
+
+        commands = self.all().get('commands')
+        print(commands)
+
+    def all(self):
+        def bans():
+            self.cur.execute("""SELECT * FROM bans""")
+            return self.cur.fetchall()
+
+        def commands():
+            query = """SELECT * FROM commands"""
+            self.cur.execute(query)
+            return self.cur.fetchall()
+        return {'bans': bans(), 'commands': commands()}