From 3383e84221fb8c462b6f2d090cbc0363c6251ebc Mon Sep 17 00:00:00 2001 From: Seaswimmer Date: Wed, 5 Jun 2024 00:39:56 -0400 Subject: [PATCH] fix(aurora): fixed some issues with aiosqlite --- aurora/aurora.py | 14 ++++++------- aurora/importers/aurora.py | 9 +++++---- aurora/models/moderation.py | 38 +++++++++++++++++++----------------- aurora/utilities/database.py | 11 ++--------- 4 files changed, 34 insertions(+), 38 deletions(-) diff --git a/aurora/aurora.py b/aurora/aurora.py index 2332d4a..3cac57d 100644 --- a/aurora/aurora.py +++ b/aurora/aurora.py @@ -128,13 +128,13 @@ 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): - async with connect() as database: - 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]:,})") + 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]:,})") @commands.Cog.listener("on_audit_log_entry_create") async def autologger(self, entry: discord.AuditLogEntry): diff --git a/aurora/importers/aurora.py b/aurora/importers/aurora.py index b2367a6..686797f 100644 --- a/aurora/importers/aurora.py +++ b/aurora/importers/aurora.py @@ -27,10 +27,11 @@ class ImportAuroraView(ui.View): "Deleting original table...", ephemeral=True ) - async with await connect() as database: - query = f"DROP TABLE IF EXISTS moderation_{self.ctx.guild.id};" - database.execute(query) - database.commit() + database = await connect() + query = f"DROP TABLE IF EXISTS moderation_{self.ctx.guild.id};" + database.execute(query) + database.commit() + database.close() await interaction.edit_original_response(content="Creating new table...") diff --git a/aurora/models/moderation.py b/aurora/models/moderation.py index 31912d1..412e195 100644 --- a/aurora/models/moderation.py +++ b/aurora/models/moderation.py @@ -122,24 +122,26 @@ class Moderation(AuroraGuildModel): from ..utilities.json import dumps query = f"UPDATE moderation_{self.guild_id} SET timestamp = ?, moderation_type = ?, target_type = ?, moderator_id = ?, role_id = ?, duration = ?, end_timestamp = ?, reason = ?, resolved = ?, resolved_by = ?, resolve_reason = ?, expired = ?, changes = ?, metadata = ? WHERE moderation_id = ?;" - async with connect() as database: - await database.execute(query, ( - self.timestamp.timestamp(), - self.moderation_type, - self.target_type, - self.moderator_id, - self.role_id, - str(self.duration) if self.duration else None, - self.end_timestamp.timestamp() if self.end_timestamp else None, - self.reason, - self.resolved, - self.resolved_by, - self.resolve_reason, - self.expired, - dumps(self.changes), - dumps(self.metadata), - self.moderation_id, - )) + database = await connect() + await database.execute(query, ( + self.timestamp.timestamp(), + self.moderation_type, + self.target_type, + self.moderator_id, + self.role_id, + str(self.duration) if self.duration else None, + self.end_timestamp.timestamp() if self.end_timestamp else None, + self.reason, + self.resolved, + self.resolved_by, + self.resolve_reason, + self.expired, + dumps(self.changes), + dumps(self.metadata), + self.moderation_id, + )) + await database.commit() + await database.close() logger.debug("Row updated in moderation_%s!\n%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s", self.moderation_id, diff --git a/aurora/utilities/database.py b/aurora/utilities/database.py index 7d50d77..3894bce 100644 --- a/aurora/utilities/database.py +++ b/aurora/utilities/database.py @@ -1,7 +1,5 @@ # pylint: disable=cyclic-import import json -from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator import aiosqlite from discord import Guild @@ -10,14 +8,13 @@ from redbot.core import data_manager from .logger import logger -@asynccontextmanager -async def connect() -> AsyncGenerator[aiosqlite.Connection, Any, None]: +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" ) - yield connection + return connection except aiosqlite.OperationalError as e: logger.error("Unable to access the SQLite database!\nError:\n%s", e.msg) @@ -25,10 +22,6 @@ async def connect() -> AsyncGenerator[aiosqlite.Connection, Any, None]: f"Unable to access the SQLite Database!\n{e.msg}" ) from e - finally: - if connection: - await connection.close() - async def create_guild_table(guild: Guild): database = await connect()