Merge pull request #287 from sarzz2/config

Config

bot/config.py

@@ -86,11 +86,6 @@ class ErrorHandling(BaseModel):
     webhook_url: str
 
 
-class CustomRoles(BaseModel):
-    log_channel_id: int
-    divider_role_id: int
-
-
 class YouTube(BaseModel):
     channel_id: str
     text_channel_id: int
@@ -109,7 +104,6 @@ class Settings(BaseSettings):
     tags: Tags
     hastebin: Hastebin
     errors: ErrorHandling
-    custom_roles: CustomRoles
     youtube: YouTube
 
     class Config:

bot/core.py

@@ -11,6 +11,7 @@
 from discord.ext import commands, tasks
 
 from bot.config import settings
+from bot.models import GuildConfig
 from bot.services import http, paste
 from bot.services.paste import Document
 from utils.errors import IgnorableException
@@ -59,12 +60,6 @@ async def setup_hook(self) -> None:
 
         self.error_webhook = discord.Webhook.from_url(url=settings.errors.webhook_url, session=http.session)
 
-    async def on_disconnect(self):
-        update_health("ready", False)
-
-    async def on_ready(self):
-        update_health("ready", True)
-
     async def when_online(self):
         log.info("Waiting until bot is ready to load extensions and app commands.")
         await self.wait_until_ready()
@@ -88,6 +83,34 @@ async def sync_commands(self) -> None:
         log.info("Commands synced.")
         log.info(f"Successfully logged in as {self.user}. In {len(self.guilds)} guilds")
 
+    @staticmethod
+    async def on_disconnect():
+        """Called when the"""
+        update_health("ready", False)
+
+    async def on_ready(self):
+        """Called when all guilds are cached."""
+        update_health("ready", True)
+
+        query = "SELECT COALESCE(array_agg(guild_id), '{}') FROM guild_configs"
+
+        stored_ids = await GuildConfig.fetchval(query)
+        missing_ids = [(guild.id,) for guild in self.guilds if guild.id not in stored_ids]
+
+        if missing_ids:
+            query = "INSERT INTO guild_configs (guild_id) VALUES ($1)"
+            await GuildConfig.pool.executemany(query, missing_ids)
+            log.info(f"Inserted new config for {len(missing_ids)} guilds.")
+
+    async def on_guild_join(self, guild: discord.Guild):
+        log.info(f"{self.user.name} has been added to a new guild: {guild.name}")
+
+        query = """INSERT INTO guild_configs (guild_id)
+                        VALUES ($1)
+                   ON CONFLICT (guild_id)
+                            DO NOTHING"""
+        await GuildConfig.execute(query, guild.id)
+
     async def on_message(self, message):
         await self.wait_until_ready()
 
@@ -127,7 +150,7 @@ def wrap(code: str) -> str:
         await self.error_webhook.send(embed=embed)
 
     async def on_error(self, event_method: str, *args: Any, **kwargs: Any) -> None:
-        content = "\n".join(traceback.format_exception(*sys.exc_info()))
+        content = "".join(traceback.format_exception(*sys.exc_info()))
         header = f"Ignored exception in event method **{event_method}**"
 
         await self.send_error(content, header)

bot/extensions/custom_roles/commands.py

@@ -1,15 +1,29 @@
 from __future__ import annotations
 
+import datetime
+
 import discord
 from discord import app_commands, utils
 from discord.ext import commands
 
 from bot import core
-from bot.config import settings
-from bot.models import CustomRole
+from bot.models import CustomRole, GuildConfig
 
 
 class CustomRoles(commands.Cog):
+    custom_roles = app_commands.Group(
+        name="custom_roles",
+        description="Custom Role commands",
+        default_permissions=discord.Permissions(administrator=True),
+    )
+
+    config = app_commands.Group(
+        parent=custom_roles,
+        name="config",
+        description="Set configuration for custom role",
+        default_permissions=discord.Permissions(administrator=True),
+    )
+
     def __init__(self, bot):
         self.bot = bot
 
@@ -50,8 +64,15 @@ async def myrole(
             # Create and assign the role to user
             role = await interaction.guild.create_role(name=name, colour=color or discord.Color.random())
 
-            divider_role = interaction.guild.get_role(settings.custom_roles.divider_role_id)
-            await role.edit(position=divider_role.position + 1)
+            divider_role_query = """
+                            SELECT divider_role_id
+                             FROM guild_configs
+                            WHERE guild_id = $1"""
+            divider_role_id = await GuildConfig.fetchval(divider_role_query, interaction.guild.id)
+
+            if divider_role_id is not None:
+                divider_role = interaction.guild.get_role(divider_role_id)
+                await role.edit(position=divider_role.position + 1)
 
             record = await CustomRole.ensure_exists(
                 guild_id=interaction.guild.id,
@@ -105,6 +126,62 @@ async def myrole(
             ephemeral=True,
         )
 
+    @config.command(name="log-channel")
+    @app_commands.describe(channel="New channel")
+    async def log_channel(self, interaction: discord.Interaction, channel: discord.TextChannel = None):
+        """Set custom role log channel"""
+        query = """
+            UPDATE guild_configs
+               SET custom_role_log_channel_id = $1
+             WHERE guild_id = $2
+        """
+
+        if channel is None:
+            await GuildConfig.execute(query, None, interaction.guild.id)
+            return await interaction.response.send_message("Cleared log channel selection", ephemeral=True)
+
+        await GuildConfig.execute(query, channel.id, interaction.guild.id)
+        return await interaction.response.send_message(f"Log channel set to {channel.mention}", ephemeral=True)
+
+    @config.command(name="divider-role-id")
+    @app_commands.describe(role="divider role")
+    async def divider_role(self, interaction: discord.Interaction, role: discord.Role = None):
+        """Set Divider role"""
+        query = """
+            UPDATE guild_configs
+               SET divider_role_id = $1
+             WHERE guild_id = $2
+        """
+
+        if role is None:
+            await GuildConfig.execute(query, None, interaction.guild.id)
+            return await interaction.response.send_message("Cleared divider role selection", ephemeral=True)
+
+        await GuildConfig.execute(query, role.id, interaction.guild.id)
+        return await interaction.response.send_message(f"Divider role set to {role.mention}")
+
+    @config.command(name="show")
+    async def show(self, interaction: discord.Interaction):
+        """Return the current custom role configuration"""
+        query = """
+            SELECT *
+              FROM guild_configs
+             WHERE guild_id = $1
+        """
+        data = await GuildConfig.fetchrow(query, interaction.guild.id)
+
+        embed = discord.Embed(
+            title=f"Server Information - {interaction.guild.name}",
+            description="Custom Role Configuration",
+            color=discord.Color.blue(),
+            timestamp=datetime.datetime.utcnow(),
+        )
+
+        embed.set_thumbnail(url=interaction.guild.icon.url)
+        embed.add_field(name="Log channel", value=f"<#{data.custom_role_log_channel_id}>")
+
+        return await interaction.response.send_message(embed=embed)
+
 
 async def setup(bot: commands.Bot):
     await bot.add_cog(CustomRoles(bot=bot))

bot/extensions/custom_roles/events.py

@@ -5,8 +5,7 @@
 from discord.ext import commands
 
 from bot import core
-from bot.config import settings
-from bot.models.custom_roles import CustomRole
+from bot.models import CustomRole, GuildConfig
 
 
 class CustomRoleEvents(commands.Cog):
@@ -16,9 +15,11 @@ def __init__(self, bot: core.DiscordBot):
         self.bot = bot
         self.updated_at: dict[int, datetime] = {}
 
-    @property
-    def custom_roles_logs_channel(self) -> discord.TextChannel | None:
-        return self.bot.guild.get_channel(settings.custom_roles.log_channel_id)
+    async def get_custom_roles_logs_channel(self, guild_id: int):
+        query = """SELECT custom_role_log_channel_id FROM guild_configs
+                    WHERE guild_id = $1"""
+        data = await GuildConfig.fetchval(query, guild_id)
+        return self.bot.guild.get_channel(data)
 
     @commands.Cog.listener()
     async def on_guild_role_update(self, before: discord.Role, after: discord.Role):
@@ -54,7 +55,12 @@ async def on_custom_role_create(self, custom_role: CustomRole):
         embed.set_thumbnail(url=user.avatar)
         embed.set_footer(text=f"user_id: {custom_role.user_id}")
 
-        return await self.custom_roles_logs_channel.send(embed=embed)
+        channel = await self.get_custom_roles_logs_channel(custom_role.guild_id)
+        # Return if no log channel has been set
+        if channel is None:
+            return
+
+        return await channel.send(embed=embed)
 
     @commands.Cog.listener()
     async def on_custom_role_update(self, before: CustomRole, after: CustomRole):
@@ -85,4 +91,9 @@ async def on_custom_role_update(self, before: CustomRole, after: CustomRole):
         embed.set_author(name=user.name, icon_url=user.avatar)
         embed.set_footer(text=f"user_id: {after.user_id}")
 
-        return await self.custom_roles_logs_channel.send(embed=embed)
+        channel = await self.get_custom_roles_logs_channel(before.guild_id)
+        # Return if no log channel has been set
+        if channel is None:
+            return
+
+        return await channel.send(embed=embed)