[CustomRoles] guild setup for privileged role, fix deletion loop

cogs/CustomRoles.py

@@ -7,27 +7,30 @@
 
 import db
 from cogs import AinitMixin, CustomCog
-
 from models import CustomRole
-from views import RoleCreatorView, RoleCreatorResult
+from views import RoleCreatorView, RoleCreatorResult, CustomRoleSetup
 
 logger = logging.getLogger(__name__)
 
 
-class NotBoosting(commands.CheckFailure):
+class MissingRole(commands.CheckFailure):
     def __init__(self):
-        super().__init__("You are not a server booster.")
-
+        super().__init__("You are missing a role to do this.")
 
-def is_booster():
-    def predicate(ctx: commands.Context):
-        if (
-            ctx.author.premium_since is None
-            and not ctx.author.guild_permissions.administrator
-        ):
-            raise NotBoosting
 
-        return True
+def has_guild_role():
+    async def predicate(ctx: commands.Context):
+        async with ctx.bot.Session() as session:
+            custom_role_settings = await db.get_custom_role_settings(
+                session, ctx.guild.id
+            )
+            if ctx.author.guild_permissions.administrator or (
+                custom_role_settings is not None
+                and custom_role_settings._role in [role.id for role in ctx.author.roles]
+            ):
+                return True
+            else:
+                raise MissingRole
 
     return commands.check(predicate)
 
@@ -42,13 +45,32 @@ def __init__(self, bot):
 
         CustomRole.inject_bot(bot)
 
-    @commands.group(aliases=["customrole", "cr"], brief="Custom roles")
+        self._role_removal_loop.start()
+
+    def cog_unload(self):
+        self._role_removal_loop.stop()
+
+    @commands.group(
+        name="customroles", aliases=["customrole", "cr"], brief="Custom roles"
+    )
     async def custom_roles(self, ctx: commands.Context):
         if not ctx.invoked_subcommand:
             await ctx.send_help(self.custom_roles)
 
+    @custom_roles.command(brief="Set up custom roles")
+    @commands.has_permissions(administrator=True)
+    async def setup(self, ctx: commands.Context):
+        embed = Embed(
+            title="Set up custom roles",
+            description="Choose the role below that members need to have in order to be allowed to create"
+            " custom roles in the server, then hit confirm.",
+        )
+        embed.set_author(name=ctx.guild.name, icon_url=ctx.guild.icon.url)
+
+        await ctx.reply(view=CustomRoleSetup(), embed=embed, ephemeral=True)
+
     @custom_roles.command(brief="Removes your custom role")
-    @is_booster()
+    @has_guild_role()
     async def delete(self, ctx: commands.Context):
         async with self.bot.Session() as session:
             custom_role = await db.get_user_custom_role_in_guild(
@@ -68,7 +90,7 @@ async def delete(self, ctx: commands.Context):
             await ctx.reply(f"Your role '{role_name}' has been deleted.")
 
     @custom_roles.command(brief="Creates a custom role for you")
-    @is_booster()
+    @has_guild_role()
     async def create(self, ctx: commands.Context):
         embed = Embed(
             title="Create custom role",
@@ -110,13 +132,14 @@ async def _delete_custom_role(self, session, custom_role: CustomRole) -> None:
         try:
             logger.info(
                 "removing custom role from %s (%d) in %s (%d)",
+                str(custom_role.member),
+                custom_role._user,
                 str(custom_role.guild),
                 custom_role._guild,
-                str(custom_role.member),
-                custom_role._member,
             )
             await custom_role.role.delete(
-                reason=f"Member {custom_role.member} ({custom_role.member.id}) removed their boost"
+                reason=f"Member {custom_role.member} ({custom_role._user}) no longer has required role,"
+                f" removing custom role"
             )
             await db.delete_custom_role(session, custom_role._guild, custom_role._user)
         except discord.Forbidden:
@@ -125,10 +148,10 @@ async def _delete_custom_role(self, session, custom_role: CustomRole) -> None:
                 str(custom_role.guild),
                 custom_role._guild,
                 str(custom_role.member),
-                custom_role._member,
+                custom_role._user,
             )
 
-    @tasks.loop(hours=24.0)
+    @tasks.loop(minutes=1)
     async def _role_removal_loop(self) -> None:
         logger.info("running role removal task")
         async with self.bot.Session() as session:
@@ -137,10 +160,27 @@ async def _role_removal_loop(self) -> None:
             role_deletion_tasks = []
 
             for custom_role in custom_roles:
-                if custom_role.member.premium_since is None:
-                    # member is no longer boosting, add to deletion list
+                # TODO fix this N+1 issue
+                custom_role_settings = await db.get_custom_role_settings(
+                    session, custom_role._guild
+                )
+
+                has_guild_role = custom_role_settings._role in (
+                    member_role.id for member_role in custom_role.member.roles
+                )
+
+                if (
+                    not custom_role.member.guild_permissions.administrator
+                    and not has_guild_role
+                ):
+                    # member no longer has the guild's required role, add to deletion list
                     role_deletion_tasks.append(
                         self._delete_custom_role(session, custom_role)
                     )
 
-            await asyncio.gather(*custom_roles)
+            await asyncio.gather(*role_deletion_tasks)
+            await session.commit()
+
+    @_role_removal_loop.before_loop
+    async def loop_before(self):
+        await self.bot.wait_until_ready()

db.py

@@ -27,6 +27,7 @@
     RoleSettings,
     GuildCog,
     CustomRole,
+    CustomRoleSettings,
 )
 
 
@@ -381,6 +382,15 @@ async def get_custom_roles(session: AsyncSession) -> list[CustomRole]:
     return [r for (r,) in result]
 
 
+async def get_custom_role_settings(
+    session: AsyncSession, guild_id: int
+) -> typing.Optional[CustomRoleSettings]:
+    statement = select(CustomRoleSettings).where(CustomRoleSettings._guild == guild_id)
+    result = (await session.execute(statement)).first()
+
+    return result[0] if result else None
+
+
 async def get_greeter(session, guild_id, greeter_type):
     statement = select(Greeter).where(
         (Greeter._guild == guild_id) & (Greeter.type == greeter_type)

models/__init__.py

@@ -1,6 +1,6 @@
 from .botw import BotwState, BotwWinner, Nomination, Idol, BotwSettings
 from .channel_mirror import ChannelMirror
-from .custom_role import CustomRole
+from .custom_role import CustomRole, CustomRoleSettings
 from .greeter import Greeter, GreeterType
 from .guild_settings import GuildSettings, EmojiSettings, GuildCog
 from .log import CommandLog
@@ -15,6 +15,7 @@
     "BotwWinner",
     "ChannelMirror",
     "CustomRole",
+    "CustomRoleSettings",
     "GuildSettings",
     "GuildCog",
     "EmojiSettings",

models/custom_role.py

@@ -2,6 +2,7 @@
 from sqlalchemy.ext.hybrid import hybrid_property
 
 from models.base import Base
+from models.guild_settings import GuildSettingsMixin
 from models.role import RoleMixin
 
 
@@ -18,3 +19,10 @@ class CustomRole(RoleMixin, Base):
     @hybrid_property
     def member(self) -> Member:
         return self.guild.get_member(self._user)
+
+
+class CustomRoleSettings(GuildSettingsMixin, Base):
+    __tablename__ = "custom_role_settings"
+
+    """The role that gets to create custom roles/is the insertion point"""
+    _role = Column(BigInteger, nullable=False)

views/__init__.py

@@ -1,4 +1,4 @@
-from .role_creator import RoleCreatorView, RoleCreatorResult
+from .role_creator import RoleCreatorView, RoleCreatorResult, CustomRoleSetup
 
 
-__all__ = ("RoleCreatorView", "RoleCreatorResult")
+__all__ = ("RoleCreatorView", "RoleCreatorResult", "CustomRoleSetup")

views/role_creator.py

@@ -2,10 +2,12 @@
 
 import discord
 from discord import Interaction, Button
-from discord.ui import View, TextInput, Modal
+from discord.ui import View, TextInput, Modal, RoleSelect
 
 import db
 
+from models import CustomRoleSettings
+
 
 class RoleCreatorResult:
     name: str | None
@@ -28,6 +30,64 @@ def __init__(self, callback, result):
         self.result = result
 
 
+class RolePicker(RoleSelect):
+    def __init__(self):
+        super().__init__(placeholder="Choose one role...")
+
+    async def callback(self, interaction: Interaction) -> typing.Any:
+        role = self.values[0]
+        client_member = interaction.guild.get_member(interaction.client.user.id)
+        if role.position >= client_member.top_role.position:
+            await interaction.response.send_message(
+                "Please choose a role the bot can manage, i.e., one that is below its own highest role.",
+                ephemeral=True,
+            )
+
+        await interaction.response.defer()
+
+
+class CustomRoleSetup(View):
+    def __init__(self):
+        super().__init__()
+
+        self.add_item(RolePicker())
+
+    async def interaction_check(self, interaction: Interaction, /) -> bool:
+        if not interaction.user.guild_permissions.administrator:
+            await interaction.response.send_message(
+                "You are not allowed to set up custom roles.", ephemeral=True
+            )
+            return False
+
+        return True
+
+    @discord.ui.button(label="Confirm", style=discord.ButtonStyle.blurple, row=2)
+    async def confirm(self, interaction: Interaction, button: Button):
+        self.stop()
+
+        values = self.children[1].values
+
+        if len(values) != 1:
+            await interaction.response.send_message(
+                "Please try again and choose exactly one role!",
+                ephemeral=True,
+            )
+            return
+
+        role = values[0]
+
+        async with interaction.client.Session() as session:
+            custom_role_settings = CustomRoleSettings(
+                _role=role.id, _guild=interaction.guild_id
+            )
+            await session.merge(custom_role_settings)
+            await session.commit()
+
+            await interaction.response.send_message(
+                f"Members with the role {role.mention} will now be able to create custom roles!",
+            )
+
+
 class RoleCreatorView(CallbackView, View):
     @discord.ui.button(label="Click me to start", style=discord.ButtonStyle.blurple)
     async def create_role(self, interaction: Interaction, button: Button):
@@ -48,20 +108,23 @@ async def interaction_check(self, interaction: Interaction, /) -> bool:
                 )
                 return False
 
-        # FIXME remove me!
-        return True
+            custom_role_settings = await db.get_custom_role_settings(
+                session, interaction.guild_id
+            )
+            member = interaction.guild.get_member(interaction.user.id)
+
+            if member.guild_permissions.administrator or (
+                custom_role_settings is not None
+                and member is not None
+                and custom_role_settings._role in [role.id for role in member.roles]
+            ):
+                return True
 
-        if (
-            interaction.user.premium_since is None
-            and not interaction.user.guild_permissions.administrator
-        ):
             await interaction.response.send_message(
-                "You are not a server booster.", ephemeral=True
+                "You are missing a role to do this.", ephemeral=True
             )
             return False
 
-        return True
-
 
 class RoleCreatorNameConfirmationView(CallbackView, View):
     @discord.ui.button(label="I confirm", style=discord.ButtonStyle.blurple)