Rosie-Val/SeaCogs
1
0
Fork 0
Code Pull requests Activity

feat(aurora): migrated Aurora.handle_expiry() to use Moderation.execute() instead of opening its own connections

Browse source
This commit is contained in:
Seaswimmer 2024-06-05 01:29:47 -04:00
parent d07e5ed804
commit 42f7f9f69b
Signed by untrusted user: cswimr
GPG key ID: 5D671B5D03D65A7F
1 changed files with 33 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
aurora/aurora.py
View file

@ -1490,8 +1490,6 @@ class Aurora(commands.Cog):
async def handle_expiry(self): async def handle_expiry(self):
await self.bot.wait_until_red_ready() await self.bot.wait_until_red_ready()
current_time = time.time() current_time = time.time()
database = await connect()
cursor = await database.cursor()
global_unban_num = 0 global_unban_num = 0
global_addrole_num = 0 global_addrole_num = 0
global_removerole_num = 0 global_removerole_num = 0
@ -1501,20 +1499,21 @@ class Aurora(commands.Cog):
if not await self.bot.cog_disabled_in_guild(self, guild): if not await self.bot.cog_disabled_in_guild(self, guild):
time_per_guild = time.time() time_per_guild = time.time()
tempban_query = f"SELECT target_id, moderation_id FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL 0 AND end_timestamp <= ? AND moderation_type = 'TEMPBAN' AND expired = 0" tempban_query = f"SELECT * FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL 0 AND end_timestamp <= ? AND moderation_type = 'TEMPBAN' AND expired = 0"
try: try:
await cursor.execute(tempban_query, (time.time(),)) tempbans = await Moderation.execute(bot=self.bot, guild_id=guild.id, query=tempban_query, parameters=(time.time(),))
result = cursor.fetchall()
except aiosqlite.OperationalError: except aiosqlite.OperationalError:
continue continue
target_ids = [row[0] for row in result]
moderation_ids = [row[1] for row in result]
unban_num = 0 unban_num = 0
for target_id, moderation_id in zip(target_ids, moderation_ids): for moderation in tempbans:
user: discord.User = await self.bot.fetch_user(target_id) user = self.bot.get_user(moderation.target_id)
if user is None:
try:
user = self.bot.fetch_user(moderation.target_id)
except discord.errors.NotFound:
continue
name = ( name = (
f"{user.name}#{user.discriminator}" f"{user.name}#{user.discriminator}"
if user.discriminator != "0" if user.discriminator != "0"
@ -1522,14 +1521,14 @@ class Aurora(commands.Cog):
) )
try: try:
await guild.unban( await guild.unban(
user, reason=f"Automatic unban from case #{moderation_id}" user, reason=f"Automatic unban from case #{moderation.id}"
) )
embed = await message_factory( embed = await message_factory(
bot=self.bot, bot=self.bot,
color=await self.bot.get_embed_color(guild.channels[0]), color=await self.bot.get_embed_color(guild.channels[0]),
guild=guild, guild=guild,
reason=f"Automatic unban from case #{moderation_id}", reason=f"Automatic unban from case #{moderation.id}",
moderation_type="unbanned", moderation_type="unbanned",
) )
@ -1538,14 +1537,14 @@ class Aurora(commands.Cog):
except discord.errors.HTTPException: except discord.errors.HTTPException:
pass pass
logger.debug( logger.trace(
"Unbanned %s (%s) from %s (%s)", "Unbanned %s (%s) from %s (%s)",
name, name,
user.id, user.id,
guild.name, guild.name,
guild.id, guild.id,
) )
unban_num = unban_num + 1 unban_num += 1
except ( except (
discord.errors.NotFound, discord.errors.NotFound,
discord.errors.Forbidden, discord.errors.Forbidden,
@ -1561,26 +1560,23 @@ class Aurora(commands.Cog):
) )
removerole_num = 0 removerole_num = 0
addrole_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL AND end_timestamp <= ? AND moderation_type = 'ADDROLE' AND expired = 0" addrole_query = f"SELECT * FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL AND end_timestamp <= ? AND moderation_type = 'ADDROLE' AND expired = 0"
try: addroles = await Moderation.execute(bot=self.bot, guild_id=guild.id, query=addrole_query, parameters=(time.time(),))
await cursor.execute(addrole_query, (time.time(),))
result = await cursor.fetchall()
except aiosqlite.OperationalError:
continue
target_ids = [row[0] for row in result]
moderation_ids = [row[1] for row in result]
role_ids = [row[2] for row in result]
for target_id, moderation_id, role_id in zip( for moderation in addroles:
target_ids, moderation_ids, role_ids
):
try: try:
member = await guild.fetch_member(target_id) member = await guild.fetch_member(moderation.target_id)
await member.remove_roles( await member.remove_roles(
Object(role_id), reason=f"Automatic role removal from case #{moderation_id}" Object(moderation.role_id), reason=f"Automatic role removal from case #{moderation.id}"
) )
logger.trace(
"Removed role %s from %s (%s)",
moderation.role_id,
member.name,
member.id,
)
removerole_num = removerole_num + 1 removerole_num = removerole_num + 1
except ( except (
discord.errors.NotFound, discord.errors.NotFound,
@ -1589,44 +1585,36 @@ class Aurora(commands.Cog):
) as e: ) as e:
logger.error( logger.error(
"Removing the role %s from user %s failed due to: \n%s", "Removing the role %s from user %s failed due to: \n%s",
role_id, moderation.role_id,
target_id, moderation.target_id,
e, e,
) )
continue continue
addrole_num = 0 addrole_num = 0
removerole_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL 0 AND end_timestamp <= ? AND moderation_type = 'REMOVEROLE' AND expired = 0" removerole_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp IS NOT NULL 0 AND end_timestamp <= ? AND moderation_type = 'REMOVEROLE' AND expired = 0"
try: removeroles = await Moderation.execute(bot=self.bot, guild_id=guild.id, query=removerole_query, parameters=(time.time(),))
await cursor.execute(removerole_query, (time.time(),))
result = await cursor.fetchall()
except aiosqlite.OperationalError:
continue
target_ids = [row[0] for row in result]
moderation_ids = [row[1] for row in result]
role_ids = [row[2] for row in result]
for target_id, moderation_id, role_id in zip( for moderation in removeroles:
target_ids, moderation_ids, role_ids
):
try: try:
member = await guild.fetch_member(target_id) member = await guild.fetch_member(moderation.target_id)
await member.add_roles( await member.add_roles(
Object(role_id), reason=f"Automatic role addition from case #{moderation_id}" Object(moderation.role_id), reason=f"Automatic role addition from case #{moderation.id}"
) )
logger.trace("Added role %s to %s (%s)", moderation.role_id, member.name, member.id)
addrole_num = addrole_num + 1 addrole_num = addrole_num + 1
except ( except (
discord.errors.NotFound, discord.errors.NotFound,
discord.errors.Forbidden, discord.errors.Forbidden,
discord.errors.HTTPException, discord.errors.HTTPException,
) as e: ) as e:
logger.error("Adding the role %s to user %s failed due to: \n%s", role_id, target_id, e) logger.error("Adding the role %s to user %s failed due to: \n%s", moderation.role_id, moderation.target_id, e)
continue continue
expiry_query = f"UPDATE `moderation_{guild.id}` SET expired = 1 WHERE (end_timestamp IS NOT NULL AND end_timestamp <= ? AND expired = 0) OR (expired = 0 AND resolved = 1);" expiry_query = f"UPDATE `moderation_{guild.id}` SET expired = 1 WHERE (end_timestamp IS NOT NULL AND end_timestamp <= ? AND expired = 0) OR (expired = 0 AND resolved = 1);"
await cursor.execute(expiry_query, (time.time(),)) await Moderation.execute(bot=self.bot, guild_id=guild.id, query=expiry_query, parameters=(time.time(),))
per_guild_completion_time = (time.time() - time_per_guild) * 1000 per_guild_completion_time = (time.time() - time_per_guild) * 1000
logger.debug( logger.debug(
@ -1642,9 +1630,6 @@ class Aurora(commands.Cog):
global_addrole_num = global_addrole_num + addrole_num global_addrole_num = global_addrole_num + addrole_num
global_removerole_num = global_removerole_num + removerole_num global_removerole_num = global_removerole_num + removerole_num
await database.commit()
await cursor.close()
await database.close()
completion_time = (time.time() - current_time) * 1000 completion_time = (time.time() - current_time) * 1000
logger.debug( logger.debug(