New Command:

Added /removebirthday command to manually remove birthday roles from users
Includes permission checks and proper error handling
Cleans up any scheduled tasks for the user
Enhanced Error Handling:

Added comprehensive try/except blocks throughout the cog
Added detailed logging for all operations and errors
Better handling of Discord API errors
Improved error messages for users
Improved Role Removal:

Added hourly cleanup task to ensure role removals are processed
Better handling of timezone issues
Proper cleanup of tasks when cog is unloaded
Improved task management to prevent memory leaks
Better Logging:

Added detailed logging for all operations
Logs include guild IDs, member IDs, and error details
Helps track issues and debug problems
This commit is contained in:
pacnpal
2024-11-15 13:40:28 +00:00
parent b4479c951b
commit 0194142c56

View File

@@ -5,6 +5,11 @@ from redbot.core.config import Config
from datetime import datetime, time, timedelta from datetime import datetime, time, timedelta
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
import random import random
import logging
import asyncio
# Set up logging
logger = logging.getLogger("red.birthday")
# Define context menu command outside the class # Define context menu command outside the class
@app_commands.context_menu(name="Give Birthday Role") @app_commands.context_menu(name="Give Birthday Role")
@@ -15,27 +20,36 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor
cog = interaction.client.get_cog("Birthday") cog = interaction.client.get_cog("Birthday")
if not cog: if not cog:
logger.error("Birthday cog not loaded during context menu execution")
await interaction.followup.send("Birthday cog is not loaded.", ephemeral=True) await interaction.followup.send("Birthday cog is not loaded.", ephemeral=True)
return return
# Check if the user has permission to use this command # Check if the user has permission to use this command
allowed_roles = await cog.config.guild(interaction.guild).allowed_roles() allowed_roles = await cog.config.guild(interaction.guild).allowed_roles()
if not any(role.id in allowed_roles for role in interaction.user.roles): if not any(role.id in allowed_roles for role in interaction.user.roles):
logger.warning(f"User {interaction.user.id} attempted to use birthday context menu without permission")
return await interaction.followup.send("You don't have permission to use this command.", ephemeral=True) return await interaction.followup.send("You don't have permission to use this command.", ephemeral=True)
birthday_role_id = await cog.config.guild(interaction.guild).birthday_role() birthday_role_id = await cog.config.guild(interaction.guild).birthday_role()
if not birthday_role_id: if not birthday_role_id:
logger.error(f"Birthday role not set for guild {interaction.guild.id}")
return await interaction.followup.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True) return await interaction.followup.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True)
birthday_role = interaction.guild.get_role(birthday_role_id) birthday_role = interaction.guild.get_role(birthday_role_id)
if not birthday_role: if not birthday_role:
logger.error(f"Birthday role {birthday_role_id} not found in guild {interaction.guild.id}")
return await interaction.followup.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True) return await interaction.followup.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True)
# Assign the role, ignoring hierarchy # Assign the role, ignoring hierarchy
try: try:
await member.add_roles(birthday_role, reason="Birthday role") await member.add_roles(birthday_role, reason="Birthday role")
logger.info(f"Birthday role assigned to {member.id} in guild {interaction.guild.id}")
except discord.Forbidden: except discord.Forbidden:
logger.error(f"Failed to assign birthday role to {member.id} in guild {interaction.guild.id}: Insufficient permissions")
return await interaction.followup.send("I don't have permission to assign that role.", ephemeral=True) return await interaction.followup.send("I don't have permission to assign that role.", ephemeral=True)
except discord.HTTPException as e:
logger.error(f"Failed to assign birthday role to {member.id} in guild {interaction.guild.id}: {str(e)}")
return await interaction.followup.send("Failed to assign the birthday role due to a Discord error.", ephemeral=True)
# Generate birthday message with random cakes (or pie) # Generate birthday message with random cakes (or pie)
cakes = random.randint(0, 5) cakes = random.randint(0, 5)
@@ -48,7 +62,8 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor
birthday_channel_id = await cog.config.guild(interaction.guild).birthday_channel() birthday_channel_id = await cog.config.guild(interaction.guild).birthday_channel()
if birthday_channel_id: if birthday_channel_id:
channel = interaction.client.get_channel(birthday_channel_id) channel = interaction.client.get_channel(birthday_channel_id)
if not channel: # If the set channel doesn't exist anymore if not channel:
logger.warning(f"Birthday channel {birthday_channel_id} not found in guild {interaction.guild.id}")
channel = interaction.channel channel = interaction.channel
else: else:
channel = interaction.channel channel = interaction.channel
@@ -61,6 +76,7 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor
try: try:
tz = ZoneInfo(timezone) tz = ZoneInfo(timezone)
except ZoneInfoNotFoundError: except ZoneInfoNotFoundError:
logger.warning(f"Invalid timezone {timezone} for guild {interaction.guild.id}, defaulting to UTC")
await interaction.followup.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True) await interaction.followup.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True)
tz = ZoneInfo("UTC") tz = ZoneInfo("UTC")
@@ -69,6 +85,7 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor
await cog.schedule_birthday_role_removal(interaction.guild, member, birthday_role, midnight) await cog.schedule_birthday_role_removal(interaction.guild, member, birthday_role, midnight)
except Exception as e: except Exception as e:
logger.error(f"Unexpected error in birthday context menu: {str(e)}", exc_info=True)
try: try:
await interaction.followup.send(f"An error occurred: {str(e)}", ephemeral=True) await interaction.followup.send(f"An error occurred: {str(e)}", ephemeral=True)
except: except:
@@ -89,7 +106,124 @@ class Birthday(commands.Cog):
} }
self.config.register_guild(**default_guild) self.config.register_guild(**default_guild)
self.birthday_tasks = {} self.birthday_tasks = {}
self.cleanup_task = None
self.bot.loop.create_task(self.initialize())
async def initialize(self):
"""Initialize the cog and start the cleanup task."""
await self.bot.wait_until_ready()
await self.reload_scheduled_tasks()
self.cleanup_task = self.bot.loop.create_task(self.daily_cleanup())
async def cog_unload(self):
"""Clean up tasks when the cog is unloaded."""
if self.cleanup_task:
self.cleanup_task.cancel()
for task in self.birthday_tasks.values():
task.cancel()
@commands.hybrid_command(name="removebirthday")
@app_commands.guild_only()
@app_commands.describe(member="The member to remove the birthday role from")
async def remove_birthday(self, ctx: commands.Context, member: discord.Member):
"""Remove the birthday role from a user."""
try:
# Check if the user has permission to use this command
allowed_roles = await self.config.guild(ctx.guild).allowed_roles()
if not any(role.id in allowed_roles for role in ctx.author.roles):
logger.warning(f"User {ctx.author.id} attempted to remove birthday role without permission")
return await ctx.send("You don't have permission to use this command.", ephemeral=True)
birthday_role_id = await self.config.guild(ctx.guild).birthday_role()
if not birthday_role_id:
logger.error(f"Birthday role not set for guild {ctx.guild.id}")
return await ctx.send("The birthday role hasn't been set.", ephemeral=True)
birthday_role = ctx.guild.get_role(birthday_role_id)
if not birthday_role:
logger.error(f"Birthday role {birthday_role_id} not found in guild {ctx.guild.id}")
return await ctx.send("The birthday role doesn't exist anymore.", ephemeral=True)
if birthday_role not in member.roles:
return await ctx.send(f"{member.display_name} doesn't have the birthday role.", ephemeral=True)
try:
await member.remove_roles(birthday_role, reason="Birthday role manually removed")
logger.info(f"Birthday role manually removed from {member.id} in guild {ctx.guild.id}")
except discord.Forbidden:
logger.error(f"Failed to remove birthday role from {member.id} in guild {ctx.guild.id}: Insufficient permissions")
return await ctx.send("I don't have permission to remove that role.", ephemeral=True)
except discord.HTTPException as e:
logger.error(f"Failed to remove birthday role from {member.id} in guild {ctx.guild.id}: {str(e)}")
return await ctx.send("Failed to remove the birthday role due to a Discord error.", ephemeral=True)
# Remove scheduled task if it exists
if str(member.id) in (await self.config.guild(ctx.guild).scheduled_tasks()):
await self.config.guild(ctx.guild).scheduled_tasks.clear_raw(str(member.id))
if ctx.guild.id in self.birthday_tasks:
self.birthday_tasks[ctx.guild.id].cancel()
del self.birthday_tasks[ctx.guild.id]
await ctx.send(f"Birthday role removed from {member.display_name}!", ephemeral=True)
except Exception as e:
logger.error(f"Unexpected error in remove_birthday command: {str(e)}", exc_info=True)
await ctx.send(f"An error occurred while removing the birthday role: {str(e)}", ephemeral=True)
async def daily_cleanup(self):
"""Daily task to ensure all birthday roles are properly removed."""
while True:
try:
await asyncio.sleep(3600) # Check every hour
logger.info("Running daily birthday role cleanup")
for guild in self.bot.guilds:
try:
scheduled_tasks = await self.config.guild(guild).scheduled_tasks()
timezone = await self.config.guild(guild).timezone()
try:
tz = ZoneInfo(timezone)
except ZoneInfoNotFoundError:
logger.warning(f"Invalid timezone {timezone} for guild {guild.id}, defaulting to UTC")
tz = ZoneInfo("UTC")
now = datetime.now(tz)
for member_id, task_info in scheduled_tasks.items():
try:
member = guild.get_member(int(member_id))
if not member:
logger.warning(f"Member {member_id} not found in guild {guild.id}, removing task")
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
continue
role = guild.get_role(task_info["role_id"])
if not role:
logger.warning(f"Role {task_info['role_id']} not found in guild {guild.id}, removing task")
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
continue
remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=tz)
if now >= remove_at:
try:
await member.remove_roles(role, reason="Birthday role duration expired (cleanup)")
logger.info(f"Removed expired birthday role from {member_id} in guild {guild.id}")
except discord.Forbidden:
logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: Insufficient permissions")
except discord.HTTPException as e:
logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: {str(e)}")
finally:
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
except Exception as e:
logger.error(f"Error processing task for member {member_id} in guild {guild.id}: {str(e)}", exc_info=True)
except Exception as e:
logger.error(f"Error processing guild {guild.id} in cleanup: {str(e)}", exc_info=True)
except Exception as e:
logger.error(f"Error in daily cleanup task: {str(e)}", exc_info=True)
finally:
await asyncio.sleep(3600) # Wait an hour before next check
# [Previous commands remain unchanged...]
@commands.hybrid_command(name="setrole") @commands.hybrid_command(name="setrole")
@app_commands.guild_only() @app_commands.guild_only()
@app_commands.describe(role="The role to set as the birthday role") @app_commands.describe(role="The role to set as the birthday role")
@@ -146,24 +280,33 @@ class Birthday(commands.Cog):
@app_commands.describe(member="The member to give the birthday role to") @app_commands.describe(member="The member to give the birthday role to")
async def birthday(self, ctx: commands.Context, member: discord.Member): async def birthday(self, ctx: commands.Context, member: discord.Member):
"""Assign the birthday role to a user until midnight in the set timezone.""" """Assign the birthday role to a user until midnight in the set timezone."""
try:
# Check if the user has permission to use this command # Check if the user has permission to use this command
allowed_roles = await self.config.guild(ctx.guild).allowed_roles() allowed_roles = await self.config.guild(ctx.guild).allowed_roles()
if not any(role.id in allowed_roles for role in ctx.author.roles): if not any(role.id in allowed_roles for role in ctx.author.roles):
logger.warning(f"User {ctx.author.id} attempted to use birthday command without permission")
return await ctx.send("You don't have permission to use this command.", ephemeral=True) return await ctx.send("You don't have permission to use this command.", ephemeral=True)
birthday_role_id = await self.config.guild(ctx.guild).birthday_role() birthday_role_id = await self.config.guild(ctx.guild).birthday_role()
if not birthday_role_id: if not birthday_role_id:
logger.error(f"Birthday role not set for guild {ctx.guild.id}")
return await ctx.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True) return await ctx.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True)
birthday_role = ctx.guild.get_role(birthday_role_id) birthday_role = ctx.guild.get_role(birthday_role_id)
if not birthday_role: if not birthday_role:
logger.error(f"Birthday role {birthday_role_id} not found in guild {ctx.guild.id}")
return await ctx.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True) return await ctx.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True)
# Assign the role, ignoring hierarchy # Assign the role, ignoring hierarchy
try: try:
await member.add_roles(birthday_role, reason="Birthday role") await member.add_roles(birthday_role, reason="Birthday role")
logger.info(f"Birthday role assigned to {member.id} in guild {ctx.guild.id}")
except discord.Forbidden: except discord.Forbidden:
logger.error(f"Failed to assign birthday role to {member.id} in guild {ctx.guild.id}: Insufficient permissions")
return await ctx.send("I don't have permission to assign that role.", ephemeral=True) return await ctx.send("I don't have permission to assign that role.", ephemeral=True)
except discord.HTTPException as e:
logger.error(f"Failed to assign birthday role to {member.id} in guild {ctx.guild.id}: {str(e)}")
return await ctx.send("Failed to assign the birthday role due to a Discord error.", ephemeral=True)
# Generate birthday message with random cakes (or pie) # Generate birthday message with random cakes (or pie)
cakes = random.randint(0, 5) cakes = random.randint(0, 5)
@@ -177,6 +320,7 @@ class Birthday(commands.Cog):
if birthday_channel_id: if birthday_channel_id:
channel = self.bot.get_channel(birthday_channel_id) channel = self.bot.get_channel(birthday_channel_id)
if not channel: # If the set channel doesn't exist anymore if not channel: # If the set channel doesn't exist anymore
logger.warning(f"Birthday channel {birthday_channel_id} not found in guild {ctx.guild.id}")
channel = ctx.channel channel = ctx.channel
else: else:
channel = ctx.channel channel = ctx.channel
@@ -189,6 +333,7 @@ class Birthday(commands.Cog):
try: try:
tz = ZoneInfo(timezone) tz = ZoneInfo(timezone)
except ZoneInfoNotFoundError: except ZoneInfoNotFoundError:
logger.warning(f"Invalid timezone {timezone} for guild {ctx.guild.id}, defaulting to UTC")
await ctx.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True) await ctx.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True)
tz = ZoneInfo("UTC") tz = ZoneInfo("UTC")
@@ -196,14 +341,19 @@ class Birthday(commands.Cog):
midnight = datetime.combine(now.date() + timedelta(days=1), time.min).replace(tzinfo=tz) midnight = datetime.combine(now.date() + timedelta(days=1), time.min).replace(tzinfo=tz)
await self.schedule_birthday_role_removal(ctx.guild, member, birthday_role, midnight) await self.schedule_birthday_role_removal(ctx.guild, member, birthday_role, midnight)
except Exception as e:
logger.error(f"Unexpected error in birthday command: {str(e)}", exc_info=True)
await ctx.send(f"An error occurred: {str(e)}", ephemeral=True)
@commands.hybrid_command(name="bdaycheck") @commands.hybrid_command(name="bdaycheck")
@app_commands.guild_only() @app_commands.guild_only()
async def bdaycheck(self, ctx: commands.Context): async def bdaycheck(self, ctx: commands.Context):
"""Check the upcoming birthday role removal tasks.""" """Check the upcoming birthday role removal tasks."""
try:
# Check if the user has permission to use this command # Check if the user has permission to use this command
allowed_roles = await self.config.guild(ctx.guild).allowed_roles() allowed_roles = await self.config.guild(ctx.guild).allowed_roles()
if not any(role.id in allowed_roles for role in ctx.author.roles): if not any(role.id in allowed_roles for role in ctx.author.roles):
logger.warning(f"User {ctx.author.id} attempted to use bdaycheck command without permission")
return await ctx.send("You don't have permission to use this command.", ephemeral=True) return await ctx.send("You don't have permission to use this command.", ephemeral=True)
scheduled_tasks = await self.config.guild(ctx.guild).scheduled_tasks() scheduled_tasks = await self.config.guild(ctx.guild).scheduled_tasks()
@@ -222,36 +372,87 @@ class Birthday(commands.Cog):
message += f"- {member.display_name} ({member.id}): {role.name} will be removed at {remove_at}\n" message += f"- {member.display_name} ({member.id}): {role.name} will be removed at {remove_at}\n"
await ctx.send(message, ephemeral=True) await ctx.send(message, ephemeral=True)
except Exception as e:
logger.error(f"Unexpected error in bdaycheck command: {str(e)}", exc_info=True)
await ctx.send(f"An error occurred while checking birthday tasks: {str(e)}", ephemeral=True)
async def schedule_birthday_role_removal(self, guild, member, role, when): async def schedule_birthday_role_removal(self, guild, member, role, when):
"""Schedule the removal of the birthday role.""" """Schedule the removal of the birthday role."""
try:
await self.config.guild(guild).scheduled_tasks.set_raw(str(member.id), value={ await self.config.guild(guild).scheduled_tasks.set_raw(str(member.id), value={
"role_id": role.id, "role_id": role.id,
"remove_at": when.isoformat() "remove_at": when.isoformat()
}) })
if guild.id in self.birthday_tasks:
self.birthday_tasks[guild.id].cancel()
self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, when)) self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, when))
logger.info(f"Scheduled birthday role removal for {member.id} in guild {guild.id} at {when}")
except Exception as e:
logger.error(f"Failed to schedule birthday role removal for {member.id} in guild {guild.id}: {str(e)}", exc_info=True)
raise
async def remove_birthday_role(self, guild, member, role, when): async def remove_birthday_role(self, guild, member, role, when):
"""Remove the birthday role at the specified time.""" """Remove the birthday role at the specified time."""
try:
await discord.utils.sleep_until(when) await discord.utils.sleep_until(when)
try: try:
await member.remove_roles(role, reason="Birthday role duration expired") await member.remove_roles(role, reason="Birthday role duration expired")
except (discord.Forbidden, discord.HTTPException): logger.info(f"Birthday role removed from {member.id} in guild {guild.id}")
pass # If we can't remove the role, we'll just let it be except discord.Forbidden:
logger.error(f"Failed to remove birthday role from {member.id} in guild {guild.id}: Insufficient permissions")
except discord.HTTPException as e:
logger.error(f"Failed to remove birthday role from {member.id} in guild {guild.id}: {str(e)}")
except asyncio.CancelledError:
logger.info(f"Birthday role removal task cancelled for {member.id} in guild {guild.id}")
raise
except Exception as e:
logger.error(f"Error removing birthday role from {member.id} in guild {guild.id}: {str(e)}", exc_info=True)
finally: finally:
if guild.id in self.birthday_tasks:
del self.birthday_tasks[guild.id] del self.birthday_tasks[guild.id]
await self.config.guild(guild).scheduled_tasks.clear_raw(str(member.id)) await self.config.guild(guild).scheduled_tasks.clear_raw(str(member.id))
async def reload_scheduled_tasks(self): async def reload_scheduled_tasks(self):
"""Reload and reschedule tasks from the configuration.""" """Reload and reschedule tasks from the configuration."""
try:
logger.info("Reloading scheduled birthday tasks")
for guild in self.bot.guilds: for guild in self.bot.guilds:
try:
scheduled_tasks = await self.config.guild(guild).scheduled_tasks() scheduled_tasks = await self.config.guild(guild).scheduled_tasks()
for member_id, task_info in scheduled_tasks.items(): for member_id, task_info in scheduled_tasks.items():
try:
member = guild.get_member(int(member_id)) member = guild.get_member(int(member_id))
if not member: if not member:
logger.warning(f"Member {member_id} not found in guild {guild.id}, removing task")
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
continue continue
role = guild.get_role(task_info["role_id"]) role = guild.get_role(task_info["role_id"])
if not role: if not role:
logger.warning(f"Role {task_info['role_id']} not found in guild {guild.id}, removing task")
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
continue continue
remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=ZoneInfo(await self.config.guild(guild).timezone()))
self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, remove_at)) remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(
tzinfo=ZoneInfo(await self.config.guild(guild).timezone()))
if datetime.now(remove_at.tzinfo) >= remove_at:
try:
await member.remove_roles(role, reason="Birthday role duration expired (reload)")
logger.info(f"Removed expired birthday role from {member_id} in guild {guild.id}")
except discord.Forbidden:
logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: Insufficient permissions")
except discord.HTTPException as e:
logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: {str(e)}")
finally:
await self.config.guild(guild).scheduled_tasks.clear_raw(member_id)
else:
self.birthday_tasks[guild.id] = self.bot.loop.create_task(
self.remove_birthday_role(guild, member, role, remove_at))
logger.info(f"Rescheduled birthday role removal for {member_id} in guild {guild.id}")
except Exception as e:
logger.error(f"Error processing task for member {member_id} in guild {guild.id}: {str(e)}", exc_info=True)
except Exception as e:
logger.error(f"Error processing guild {guild.id} during reload: {str(e)}", exc_info=True)
except Exception as e:
logger.error(f"Error reloading scheduled tasks: {str(e)}", exc_info=True)