From 0bcbcd6c0cd45cf40fd1027b643b11612f7d0fb1 Mon Sep 17 00:00:00 2001 From: Seaswimmer Date: Sat, 8 Jun 2024 20:12:22 -0400 Subject: [PATCH] feat(aurora): bunch of changes --- aurora/aurora.py | 52 ++++++++--------- aurora/importers/aurora.py | 11 +--- aurora/importers/galacticbot.py | 5 +- aurora/models/moderation.py | 35 ++++++++---- aurora/utilities/database.py | 99 --------------------------------- aurora/utilities/utils.py | 69 +++++++++++++++++++++++ 6 files changed, 122 insertions(+), 149 deletions(-) delete mode 100644 aurora/utilities/database.py diff --git a/aurora/aurora.py b/aurora/aurora.py index 521cf5b..85d81d6 100644 --- a/aurora/aurora.py +++ b/aurora/aurora.py @@ -19,7 +19,7 @@ from redbot.core import app_commands, commands, data_manager from redbot.core.app_commands import Choice from redbot.core.bot import Red from redbot.core.commands.converter import parse_relativedelta, parse_timedelta -from redbot.core.utils.chat_formatting import box, error, humanize_list, humanize_timedelta, warning +from redbot.core.utils.chat_formatting import bold, box, error, humanize_list, humanize_timedelta, warning from .importers.aurora import ImportAuroraView from .importers.galacticbot import ImportGalacticBotView @@ -30,11 +30,10 @@ from .menus.overrides import Overrides from .models.change import Change from .models.moderation import Moderation from .utilities.config import config, register_config -from .utilities.database import connect, create_guild_table from .utilities.factory import addrole_embed, case_factory, changes_factory, evidenceformat_factory, guild_embed, immune_embed, message_factory, overrides_embed from .utilities.json import dump from .utilities.logger import logger -from .utilities.utils import check_moddable, check_permissions, get_footer_image, log, send_evidenceformat, timedelta_from_relativedelta +from .utilities.utils import check_moddable, check_permissions, create_guild_table, get_footer_image, log, send_evidenceformat, timedelta_from_relativedelta class Aurora(commands.Cog): @@ -43,28 +42,22 @@ class Aurora(commands.Cog): This cog stores all of its data in an SQLite database.""" __author__ = ["SeaswimmerTheFsh"] - __version__ = "2.2.0" + __version__ = "2.3.0" __documentation__ = "https://seacogs.coastalcommits.com/aurora/" async def red_delete_data_for_user(self, *, requester, user_id: int): if requester == "discord_deleted_user": await config.user_from_id(user_id).clear() - database = await connect() - cursor = await database.cursor() - - await cursor.execute("SHOW TABLES;") - tables = [table[0] for table in cursor.fetchall()] + results = await Moderation.execute(query="SHOW TABLES;", return_obj=False) + tables = [table[0] for table in results] condition = "target_id = %s OR moderator_id = %s;" for table in tables: delete_query = f"DELETE FROM {table[0]} WHERE {condition}" - await cursor.execute(delete_query, (user_id, user_id)) + await Moderation.execute(query=delete_query, parameters=(user_id, user_id), return_obj=False) - await database.commit() - await cursor.close() - await database.close() if requester == "owner": await config.user_from_id(user_id).clear() if requester == "user": @@ -86,16 +79,16 @@ class Aurora(commands.Cog): # and the information that aiosqlite logs is not useful to the bot owner. # This is a bad solution though as it overrides it for any other cogs that are using aiosqlite too. # If there's a better solution that you're aware of, please let me know in Discord or in a CoastalCommits issue. - py_logging.getLogger('aiosqlite').setLevel(py_logging.WARNING) + py_logging.getLogger('aiosqlite').setLevel(py_logging.INFO) def format_help_for_context(self, ctx: commands.Context) -> str: pre_processed = super().format_help_for_context(ctx) or "" n = "\n" if "\n\n" not in pre_processed else "" text = [ f"{pre_processed}{n}", - f"Cog Version: **{self.__version__}**", - f"Author: {humanize_list(self.__author__)}", - f"Documentation: {self.__documentation__}", + f"{bold('Cog Version:')} {self.__version__}", + f"{bold('Author:')} {humanize_list(self.__author__)}", + f"{bold('Documentation:')} {self.__documentation__}", ] return "\n".join(text) @@ -127,13 +120,12 @@ class Aurora(commands.Cog): async def addrole_on_member_join(self, member: discord.Member): """This method automatically adds roles to users when they join the server.""" if not await self.bot.cog_disabled_in_guild(self, member.guild): - database = await connect() query = f"""SELECT moderation_id, role_id, reason FROM moderation_{member.guild.id} WHERE target_id = ? AND moderation_type = 'ADDROLE' AND expired = 0 AND resolved = 0;""" - async with database.execute(query, (member.id,)) as cursor: - async for row in cursor: - role = member.guild.get_role(row[1]) - reason = row[2] - await member.add_roles(role, reason=f"Role automatically added on member rejoin for: {reason} (Case #{row[0]:,})") + results = Moderation.execute(query, (member.id,)) + for row in results: + role = member.guild.get_role(row[1]) + reason = row[2] + await member.add_roles(role, reason=f"Role automatically added on member rejoin for: {reason} (Case #{row[0]:,})") @commands.Cog.listener("on_audit_log_entry_create") async def autologger(self, entry: discord.AuditLogEntry): @@ -1093,7 +1085,7 @@ class Aurora(commands.Cog): ) return - database = await connect() + database = await Moderation.connect() if export: try: @@ -1268,7 +1260,7 @@ class Aurora(commands.Cog): ephemeral: bool | None = None, evidenceformat: bool = False, changes: bool = False, - export: Choice[str] | None = None, + raw: Choice[str] | None = None, ): """Check the details of a specific case. @@ -1278,9 +1270,11 @@ class Aurora(commands.Cog): What case are you looking up? ephemeral: bool Hide the command response + evidenceformat: bool + Display the evidence format of the case changes: bool List the changes made to the case - export: bool + raw: bool Export the case to a JSON file or codeblock""" permissions = check_permissions( interaction.client.user, ["embed_links"], interaction @@ -1309,8 +1303,8 @@ class Aurora(commands.Cog): ) return - if export: - if export.value == "file" or len(mod.to_json(2)) > 1800: + if raw: + if raw.value == "file" or len(mod.to_json(2)) > 1800: filename = ( str(data_manager.cog_data_path(cog_instance=self)) + str(os.sep) @@ -1319,7 +1313,7 @@ class Aurora(commands.Cog): with open(filename, "w", encoding="utf-8") as f: mod.to_json(2, f) - if export.value == "codeblock": + if raw.value == "codeblock": content = f"Case #{case:,} exported.\n" + warning( "Case was too large to export as codeblock, so it has been uploaded as a `.json` file." ) diff --git a/aurora/importers/aurora.py b/aurora/importers/aurora.py index 9bc98ae..2500993 100644 --- a/aurora/importers/aurora.py +++ b/aurora/importers/aurora.py @@ -9,9 +9,8 @@ from redbot.core import commands, data_manager from redbot.core.utils.chat_formatting import warning from ..models.moderation import Moderation -from ..utilities.database import connect, create_guild_table from ..utilities.json import dump -from ..utilities.utils import timedelta_from_string +from ..utilities.utils import create_guild_table, timedelta_from_string class ImportAuroraView(ui.View): @@ -29,10 +28,8 @@ class ImportAuroraView(ui.View): "Deleting original table...", ephemeral=True ) - database = await connect() query = f"DROP TABLE IF EXISTS moderation_{self.ctx.guild.id};" - await database.execute(query) - await database.commit() + await Moderation.execute(query=query, return_obj=False) await interaction.edit_original_response(content="Creating new table...") @@ -111,15 +108,11 @@ class ImportAuroraView(ui.View): expired=case["expired"], changes=changes, metadata=metadata, - database=database, return_obj=False ) except Exception as e: # pylint: disable=broad-exception-caught failed_cases.append(str(case["moderation_id"]) + f": {e}") - await database.commit() - await database.close() - await interaction.edit_original_response(content="Import complete.") if failed_cases: filename = ( diff --git a/aurora/importers/galacticbot.py b/aurora/importers/galacticbot.py index 453cac0..2fbc62d 100644 --- a/aurora/importers/galacticbot.py +++ b/aurora/importers/galacticbot.py @@ -8,13 +8,14 @@ from redbot.core import commands from redbot.core.utils.chat_formatting import box, warning from ..models.moderation import Change, Moderation -from ..utilities.database import connect, create_guild_table +from ..utilities.database import create_guild_table class ImportGalacticBotView(ui.View): def __init__(self, timeout, ctx, message): super().__init__() self.ctx: commands.Context = ctx + self.timeout = timeout self.message: Message = message @ui.button(label="Yes", style=ButtonStyle.success) @@ -26,7 +27,7 @@ class ImportGalacticBotView(ui.View): "Deleting original table...", ephemeral=True ) - database = await connect() + database = await Moderation.connect() cursor = await database.cursor() query = f"DROP TABLE IF EXISTS moderation_{self.ctx.guild.id};" diff --git a/aurora/models/moderation.py b/aurora/models/moderation.py index 9cc534c..afd9c8e 100644 --- a/aurora/models/moderation.py +++ b/aurora/models/moderation.py @@ -5,8 +5,10 @@ from time import time from typing import Dict, Iterable, List, Optional, Tuple, Union import discord -from aiosqlite import Cursor +from aiosqlite import Connection, Cursor, OperationalError, Row +from aiosqlite import connect as aiosqlite_connect from discord import NotFound +from redbot.core import data_manager from redbot.core.bot import Red from ..utilities.logger import logger @@ -211,34 +213,48 @@ class Moderation(AuroraGuildModel): } return cls.from_dict(bot=bot, data=case) + @staticmethod + async def connect() -> Connection: + """Connects to the SQLite database, and returns a connection object.""" + try: + connection = await aiosqlite_connect( + database=data_manager.cog_data_path(raw_name="Aurora") / "aurora.db" + ) + return connection + + except OperationalError as e: + logger.error("Unable to access the SQLite database!\nError:\n%s", e.msg) + raise ConnectionRefusedError( + f"Unable to access the SQLite Database!\n{e.msg}" + ) from e + @classmethod - async def execute(cls, bot: Red, guild_id: int, query: str, parameters: tuple | None = None, cursor: Cursor | None = None, return_obj: bool = True) -> Tuple["Moderation"]: - from ..utilities.database import connect - logger.trace("Executing query: %s", query) - logger.trace("With parameters: %s", parameters) + async def execute(cls, query: str, parameters: tuple | None = None, bot: Red | None = None, guild_id: int | None = None, cursor: Cursor | None = None, return_obj: bool = True) -> Union[Tuple["Moderation"], Iterable[Row]]: + logger.trace("Executing query: \"%s\" with parameters \"%s\"", query, parameters) if not parameters: parameters = () if not cursor: no_cursor = True - database = await connect() + database = await cls.connect() cursor = await database.cursor() else: no_cursor = False await cursor.execute(query, parameters) results = await cursor.fetchall() + await database.commit() if no_cursor: await cursor.close() await database.close() - if results and return_obj: + if results and return_obj and bot and guild_id: cases = [] for result in results: case = cls.from_result(bot=bot, result=result, guild_id=guild_id) if case.moderation_id != 0: cases.append(case) return tuple(cases) - return () + return results @classmethod async def get_latest(cls, bot: Red, guild_id: int, limit: int | None = None, offset: int = 0, types: Iterable | None = None, cursor: Cursor | None = None) -> Tuple["Moderation"]: @@ -306,7 +322,6 @@ class Moderation(AuroraGuildModel): metadata: dict | None = None, return_obj: bool = True, ) -> Union["Moderation", int]: - from ..utilities.database import connect from ..utilities.json import dumps if not timestamp: timestamp = datetime.fromtimestamp(time()) @@ -341,7 +356,7 @@ class Moderation(AuroraGuildModel): role_id = None if not database: - database = await connect() + database = await cls.connect() close_db = True else: close_db = False diff --git a/aurora/utilities/database.py b/aurora/utilities/database.py deleted file mode 100644 index a1d1808..0000000 --- a/aurora/utilities/database.py +++ /dev/null @@ -1,99 +0,0 @@ -# pylint: disable=cyclic-import -import json - -import aiosqlite -from discord import Guild -from redbot.core import data_manager - -from .logger import logger - - -async def connect() -> aiosqlite.Connection: - """Connects to the SQLite database, and returns a connection object.""" - try: - connection = await aiosqlite.connect( - database=data_manager.cog_data_path(raw_name="Aurora") / "aurora.db" - ) - return connection - - except aiosqlite.OperationalError as e: - logger.error("Unable to access the SQLite database!\nError:\n%s", e.msg) - raise ConnectionRefusedError( - f"Unable to access the SQLite Database!\n{e.msg}" - ) from e - - -async def create_guild_table(guild: Guild): - database = await connect() - - try: - await database.execute(f"SELECT * FROM `moderation_{guild.id}`") - logger.verbose("SQLite Table exists for server %s (%s)", guild.name, guild.id) - - except aiosqlite.OperationalError: - query = f""" - CREATE TABLE `moderation_{guild.id}` ( - moderation_id INTEGER PRIMARY KEY, - timestamp INTEGER NOT NULL, - moderation_type TEXT NOT NULL, - target_type TEXT NOT NULL, - target_id INTEGER NOT NULL, - moderator_id INTEGER NOT NULL, - role_id INTEGER, - duration TEXT, - end_timestamp INTEGER, - reason TEXT, - resolved INTEGER NOT NULL, - resolved_by TEXT, - resolve_reason TEXT, - expired INTEGER NOT NULL, - changes JSON NOT NULL, - metadata JSON NOT NULL - ) - """ - await database.execute(query) - - index_query_1 = f"CREATE INDEX IF NOT EXISTS idx_target_id ON moderation_{guild.id}(target_id);" - await database.execute(index_query_1) - - index_query_2 = f"CREATE INDEX IF NOT EXISTS idx_moderator_id ON moderation_{guild.id}(moderator_id);" - await database.execute(index_query_2) - - index_query_3 = f"CREATE INDEX IF NOT EXISTS idx_moderation_id ON moderation_{guild.id}(moderation_id);" - await database.execute(index_query_3) - - insert_query = f""" - INSERT INTO `moderation_{guild.id}` - (moderation_id, timestamp, moderation_type, target_type, target_id, moderator_id, role_id, duration, end_timestamp, reason, resolved, resolved_by, resolve_reason, expired, changes, metadata) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """ - insert_values = ( - 0, - 0, - "NULL", - "NULL", - 0, - 0, - None, - None, - None, - None, - 0, - None, - None, - 0, - json.dumps([]), - json.dumps({}), - ) - await database.execute(insert_query, insert_values) - - await database.commit() - - logger.debug( - "SQLite Table (moderation_%s) created for %s (%s)", - guild.id, - guild.name, - guild.id, - ) - - await database.close() diff --git a/aurora/utilities/utils.py b/aurora/utilities/utils.py index 49346dc..c898e60 100644 --- a/aurora/utilities/utils.py +++ b/aurora/utilities/utils.py @@ -2,6 +2,7 @@ from datetime import datetime, timedelta from typing import Optional, Tuple, Union +import aiosqlite from dateutil.relativedelta import relativedelta as rd from discord import File, Guild, Interaction, Member, SelectOption, TextChannel, User from discord.errors import Forbidden @@ -9,6 +10,8 @@ from redbot.core import commands, data_manager from redbot.core.utils.chat_formatting import error from ..utilities.config import config +from ..utilities.json import dumps +from ..utilities.logger import logger def check_permissions( @@ -210,3 +213,69 @@ def get_footer_image(coginstance: commands.Cog) -> File: """Returns the footer image for the embeds.""" image_path = data_manager.bundled_data_path(coginstance) / "arrow.png" return File(image_path, filename="arrow.png", description="arrow") + +async def create_guild_table(guild: Guild) -> None: + from ..models.moderation import Moderation + + try: + await Moderation.execute(f"SELECT * FROM `moderation_{guild.id}`", return_obj=False) + logger.trace("SQLite Table exists for server %s (%s)", guild.name, guild.id) + + except aiosqlite.OperationalError: + query = f""" + CREATE TABLE `moderation_{guild.id}` ( + moderation_id INTEGER PRIMARY KEY, + timestamp INTEGER NOT NULL, + moderation_type TEXT NOT NULL, + target_type TEXT NOT NULL, + target_id INTEGER NOT NULL, + moderator_id INTEGER NOT NULL, + role_id INTEGER, + duration TEXT, + end_timestamp INTEGER, + reason TEXT, + resolved INTEGER NOT NULL, + resolved_by TEXT, + resolve_reason TEXT, + expired INTEGER NOT NULL, + changes JSON NOT NULL, + metadata JSON NOT NULL + ) + """ + await Moderation.execute(query=query, return_obj=False) + + index_query_1 = f"CREATE INDEX IF NOT EXISTS idx_target_id ON moderation_{guild.id}(target_id);" + await Moderation.execute(query=index_query_1, return_obj=False) + + index_query_2 = f"CREATE INDEX IF NOT EXISTS idx_moderator_id ON moderation_{guild.id}(moderator_id);" + await Moderation.execute(query=index_query_2, return_obj=False) + + index_query_3 = f"CREATE INDEX IF NOT EXISTS idx_moderation_id ON moderation_{guild.id}(moderation_id);" + await Moderation.execute(query=index_query_3, return_obj=False) + + insert_query = f""" + INSERT INTO `moderation_{guild.id}` + (moderation_id, timestamp, moderation_type, target_type, target_id, moderator_id, role_id, duration, end_timestamp, reason, resolved, resolved_by, resolve_reason, expired, changes, metadata) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + insert_values = ( + 0, + 0, + "NULL", + "NULL", + 0, + 0, + None, + None, + None, + None, + 0, + None, + None, + 0, + dumps([]), + dumps({}), + ) + await Moderation.execute(query=insert_query, parameters=insert_values, return_obj=False) + + logger.trace("SQLite Table created for server %s (%s)", guild.name, guild.id)