(guild admin): Autoroles and role persistence.

bot/modules/guild_admin/guild_config.py

@@ -11,7 +11,7 @@ from .module import module
 
 # Pages of configuration categories to display
 cat_pages = {
-    'Administration': ('Meta', 'Guild Roles', 'Greetings', 'Farewells'),
+    'Administration': ('Meta', 'Guild Roles', 'New Members'),
     'Moderation': ('Moderation', 'Video Channels'),
     'Productivity': ('Study Tracking', 'TODO List', 'Workout'),
     'Study Rooms': ('Rented Rooms', 'Accountability Rooms'),

bot/modules/guild_admin/new_members/__init__.py

@@ -1,2 +1,3 @@
 from . import settings
 from . import greetings
+from . import roles

bot/modules/guild_admin/new_members/data.py

@@ -0,0 +1,6 @@
+from data import Table, RowTable
+
+
+autoroles = Table('autoroles')
+bot_autoroles = Table('bot_autoroles')
+past_member_roles = Table('past_member_roles')

bot/modules/guild_admin/new_members/roles.py

@@ -0,0 +1,115 @@
+import asyncio
+import discord
+from collections import defaultdict
+
+from meta import client
+from core import Lion
+from settings import GuildSettings
+
+from .settings import autoroles, bot_autoroles, role_persistence
+from .data import past_member_roles
+
+
+# Locks to avoid storing the roles while adding them
+# The locking is cautious, leaving data unchanged upon collision
+locks = defaultdict(asyncio.Lock)
+
+
+@client.add_after_event('member_join')
+async def join_role_tracker(client, member):
+    """
+    Add autoroles or saved roles as needed.
+    """
+    guild = member.guild
+    if not guild.me.guild_permissions.manage_roles:
+        # We can't manage the roles here, don't try to give/restore the member roles
+        return
+
+    async with locks[(guild.id, member.id)]:
+        if role_persistence.get(guild.id).value and client.data.lions.fetch((guild.id, member.id)):
+            # Lookup stored roles
+            role_rows = past_member_roles.select_where(
+                guildid=guild.id,
+                userid=member.id
+            )
+            # Identify roles from roleids
+            roles = (guild.get_role(row['roleid']) for row in role_rows)
+            # Remove non-existent roles
+            roles = (role for role in roles if role is not None)
+            # Remove roles the client can't add
+            roles = [role for role in roles if role < guild.me.top_role]
+            if roles:
+                try:
+                    await member.add_roles(
+                        *roles,
+                        reason="Restoring saved roles.",
+                    )
+                except discord.HTTPException:
+                    # This shouldn't ususally happen, but there are valid cases where it can
+                    # E.g. the user left while we were restoring their roles
+                    pass
+                # Event log!
+                GuildSettings(guild.id).event_log.log(
+                    "Restored the following roles for returning member {}:\n{}".format(
+                        member.mention,
+                        ', '.join(role.mention for role in roles)
+                    ),
+                    title="Saved roles restored"
+                )
+        else:
+            # Add autoroles
+            roles = bot_autoroles.get(guild.id).value if member.bot else autoroles.get(guild.id).value
+            # Remove roles the client can't add
+            roles = [role for role in roles if role < guild.me.top_role]
+            if roles:
+                try:
+                    await member.add_roles(
+                        *roles,
+                        reason="Adding autoroles.",
+                    )
+                except discord.HTTPException:
+                    # This shouldn't ususally happen, but there are valid cases where it can
+                    # E.g. the user left while we were adding autoroles
+                    pass
+                # Event log!
+                GuildSettings(guild.id).event_log.log(
+                    "Gave {} the guild autoroles:\n{}".format(
+                        member.mention,
+                        ', '.join(role.mention for role in roles)
+                    ),
+                    titles="Autoroles added"
+                )
+
+
+@client.add_after_event('member_remove')
+async def left_role_tracker(client, member):
+    """
+    Delete and re-store member roles when they leave the server.
+    """
+    if (member.guild.id, member.id) in locks and locks[(member.guild.id, member.id)].locked():
+        # Currently processing a join event
+        # Which means the member left while we were adding their roles
+        # Cautiously return, not modifying the saved role data
+        return
+
+    # Delete existing member roles for this user
+    # NOTE: Not concurrency-safe
+    past_member_roles.delete_where(
+        guildid=member.guild.id,
+        userid=member.id,
+    )
+    if role_persistence.get(member.guild.id).value:
+        # Make sure the user has an associated lion, so we can detect when they rejoin
+        Lion.fetch(member.guild.id, member.id)
+
+        # Then insert the current member roles
+        values = [
+            (member.guild.id, member.id, role.id)
+            for role in member.roles
+            if not role.is_bot_managed() and not role.is_integration() and not role.is_default()
+        ]
+        if values:
+            past_member_roles.insert_many(
+                *values,
+                insert_keys=('guildid', 'userid', 'roleid')
+            )

bot/modules/guild_admin/new_members/settings.py

@@ -1,8 +1,12 @@
 import datetime
 import discord
 
+import settings
 from settings import GuildSettings, GuildSetting
 import settings.setting_types as stypes
+from wards import guild_admin
+
+from .data import autoroles, bot_autoroles
 
 
 @GuildSettings.attach_setting
@@ -16,7 +20,7 @@ class greeting_channel(stypes.Channel, GuildSetting):
     """
     DMCHANNEL = object()
 
-    category = "Greetings"
+    category = "New Members"
 
     attr_name = 'greeting_channel'
     _data_column = 'greeting_channel'
@@ -83,7 +87,7 @@ class greeting_channel(stypes.Channel, GuildSetting):
 
 @GuildSettings.attach_setting
 class greeting_message(stypes.Message, GuildSetting):
-    category = "Greetings"
+    category = "New Members"
 
     attr_name = 'greeting_message'
     _data_column = 'greeting_message'
@@ -134,7 +138,7 @@ class greeting_message(stypes.Message, GuildSetting):
 
 @GuildSettings.attach_setting
 class returning_message(stypes.Message, GuildSetting):
-    category = "Greetings"
+    category = "New Members"
 
     attr_name = 'returning_message'
     _data_column = 'returning_message'
@@ -187,7 +191,7 @@ class returning_message(stypes.Message, GuildSetting):
 
 @GuildSettings.attach_setting
 class starting_funds(stypes.Integer, GuildSetting):
-    category = "Greetings"
+    category = "New Members"
 
     attr_name = 'starting_funds'
     _data_column = 'starting_funds'
@@ -204,3 +208,96 @@ class starting_funds(stypes.Integer, GuildSetting):
     @property
     def success_response(self):
         return "Members will be given `{}` coins when they first join the server.".format(self.formatted)
+
+
+@GuildSettings.attach_setting
+class autoroles(stypes.RoleList, settings.ListData, settings.Setting):
+    category = "New Members"
+    write_ward = guild_admin
+
+    attr_name = 'autoroles'
+
+    _table_interface = autoroles
+    _id_column = 'guildid'
+    _data_column = 'roleid'
+
+    display_name = "autoroles"
+    desc = "Roles to give automatically to new members."
+
+    _force_unique = True
+
+    long_desc = (
+        "These roles will be given automatically to users when they join the server. "
+        "If `role_persistence` is enabled, the roles will only be given the first time a user joins the server."
+    )
+
+    # Flat cache, no need to expire
+    _cache = {}
+
+    @property
+    def success_response(self):
+        if self.value:
+            return "New members will be given the following roles:\n{}".format(self.formatted)
+        else:
+            return "New members will not automatically be given any roles."
+
+
+@GuildSettings.attach_setting
+class bot_autoroles(stypes.RoleList, settings.ListData, settings.Setting):
+    category = "New Members"
+    write_ward = guild_admin
+
+    attr_name = 'bot_autoroles'
+
+    _table_interface = bot_autoroles
+    _id_column = 'guildid'
+    _data_column = 'roleid'
+
+    display_name = "bot_autoroles"
+    desc = "Roles to give automatically to new bots."
+
+    _force_unique = True
+
+    long_desc = (
+        "These roles will be given automatically to bots when they join the server. "
+        "If `role_persistence` is enabled, the roles will only be given the first time a bot joins the server."
+    )
+
+    # Flat cache, no need to expire
+    _cache = {}
+
+    @property
+    def success_response(self):
+        if self.value:
+            return "New bots will be given the following roles:\n{}".format(self.formatted)
+        else:
+            return "New bots will not automatically be given any roles."
+
+
+@GuildSettings.attach_setting
+class role_persistence(stypes.Boolean, GuildSetting):
+    category = "New Members"
+
+    attr_name = "role_persistence"
+
+    _data_column = 'persist_roles'
+
+    display_name = "role_persistence"
+    desc = "Whether to remember member roles when they leave the server."
+    _outputs = {True: "Enabled", False: "Disabled"}
+    _default = True
+
+    long_desc = (
+        "When enabled, restores member roles when they rejoin the server.\n"
+        "This enables profile roles and purchased roles, such as field of study and colour roles, "
+        "as well as moderation roles, "
+        "such as the studyban and mute roles, to persist even when a member leaves and rejoins.\n"
+        "Note: Members who leave while this is disabled will not have their roles restored."
+    )
+
+    @property
+    def success_response(self):
+        if self.value:
+            return "Roles will now be restored when a member rejoins."
+        else:
+            return "Member roles will no longer be saved or restored."

data/migration/v3-v4/migration.sql

@@ -2,8 +2,30 @@ ALTER TABLE guild_config
   ADD COLUMN greeting_channel BIGINT,
   ADD COLUMN greeting_message TEXT,
   ADD COLUMN returning_message TEXT,
-  ADD COLUMN starting_funds INTEGER;
+  ADD COLUMN starting_funds INTEGER,
+  ADD COLUMN persist_roles BOOLEAN;
 
 CREATE INDEX rented_members_users ON rented_members (userid);
 
+CREATE TABLE autoroles(
+  guildid BIGINT NOT NULL ,
+  roleid BIGINT NOT NULL
+);
+CREATE INDEX autoroles_guilds ON autoroles (guildid);
+
+CREATE TABLE bot_autoroles(
+  guildid BIGINT NOT NULL ,
+  roleid BIGINT NOT NULL
+);
+CREATE INDEX bot_autoroles_guilds ON bot_autoroles (guildid);
+
+CREATE TABLE past_member_roles(
+  guildid BIGINT NOT NULL,
+  userid BIGINT NOT NULL,
+  roleid BIGINT NOT NULL,
+  _timestamp TIMESTAMPTZ DEFAULT now(),
+  FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid)
+);
+CREATE INDEX member_role_persistence_members ON past_member_roles (guildid, userid);
+
 INSERT INTO VersionHistory (version, author) VALUES (4, 'v3-v4 Migration');

data/schema.sql

@@ -4,7 +4,7 @@ CREATE TABLE VersionHistory(
   time TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP NOT NULL,
   author TEXT
 );
-INSERT INTO VersionHistory (version, author) VALUES (3, 'Initial Creation');
+INSERT INTO VersionHistory (version, author) VALUES (4, 'Initial Creation');
 
 
 CREATE OR REPLACE FUNCTION update_timestamp_column()
@@ -76,7 +76,8 @@ CREATE TABLE guild_config(
   greeting_channel BIGINT,
   greeting_message TEXT,
   returning_message TEXT,
-  starting_funds INTEGER
+  starting_funds INTEGER,
+  persist_roles BOOLEAN
 );
 
 CREATE TABLE ignored_members(
@@ -96,6 +97,18 @@ CREATE TABLE donator_roles(
   roleid BIGINT NOT NULL
 );
 CREATE INDEX donator_roles_guilds ON donator_roles (guildid);
+
+CREATE TABLE autoroles(
+  guildid BIGINT NOT NULL,
+  roleid BIGINT NOT NULL
+);
+CREATE INDEX autoroles_guilds ON autoroles (guildid);
+
+CREATE TABLE bot_autoroles(
+  guildid BIGINT NOT NULL ,
+  roleid BIGINT NOT NULL
+);
+CREATE INDEX bot_autoroles_guilds ON bot_autoroles (guildid);
 -- }}}
 
 -- Workout data {{{
@@ -478,4 +491,15 @@ CREATE VIEW accountability_open_slots AS
   WHERE closed_at IS NULL
   ORDER BY start_at ASC;
 -- }}}
+
+-- Member Role Data {{{
+CREATE TABLE past_member_roles(
+  guildid BIGINT NOT NULL,
+  userid BIGINT NOT NULL,
+  roleid BIGINT NOT NULL,
+  _timestamp TIMESTAMPTZ DEFAULT now(),
+  FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid)
+);
+CREATE INDEX member_role_persistence_members ON past_member_roles (guildid, userid);
+-- }}}
 -- vim: set fdm=marker: