Updated commands and events for levelling.

bot/extensions/levelling/__init__.py

@@ -1,9 +1,9 @@
 from bot.core import DiscordBot
 
 from .commands import Levelling
-from .events import LevelEvents
+from .events import LevellingEvents
 
 
 async def setup(bot: DiscordBot) -> None:
     await bot.add_cog(Levelling(bot=bot))
-    await bot.add_cog(LevelEvents(bot=bot))
+    await bot.add_cog(LevellingEvents(bot=bot))

bot/extensions/levelling/events.py

@@ -1,43 +1,60 @@
-from bisect import bisect
-
-import discord
 from discord.ext import commands
 
 from bot import core
-from bot.models import LevellingRole, PersistentRole
+from bot.models import LevellingRole, LevellingUser
 
 
-class LevelEvents(commands.Cog):
+class LevellingEvents(commands.Cog):
     """Events for Levelling in discord."""
 
     def __init__(self, bot: core.DiscordBot):
         self.bot = bot
 
     @commands.Cog.listener()
-    async def on_level_up(self, new_level, member: discord.Member):
-        """Add roles when user levels up"""
-        data = await LevellingRole.list_by_guild(guild_id=member.guild.id)
-        for i in range(len(data)):
-            # Adding roles when user levels up
-            if new_level >= data[i].level and member.guild.get_role(data[i].role_id) not in member.roles:
-                await member.add_roles(member.guild.get_role(data[i].role_id))
-                await PersistentRole.insert_by_guild(
-                    guild_id=member.guild.id, user_id=member.id, role_id=data[i].role_id
-                )
-            # Removing roles when users level down using remove_xp command
-            elif new_level <= data[i].level and member.guild.get_role(data[i].role_id) in member.roles:
-                await member.remove_roles(member.guild.get_role(data[i].role_id))
-                await PersistentRole.delete_by_guild(guild_id=member.guild.id, user_id=member.id)
-
-    @commands.Cog.listener()
-    async def on_xp_updated(self, data, member: discord.Member, required_xp):
-        """Function to check if user's level has changed and trigger the event to assign the roles"""
-        # Calculating old and new level
-        try:
-            old_level = bisect(required_xp, data.old_total_xp) - 1
-            new_level = bisect(required_xp, data.total_xp) - 1
-
-            if old_level != new_level:
-                self.bot.dispatch("level_up", new_level=new_level, member=member)
-        except AttributeError:
-            pass
+    async def on_xp_update(self, before: LevellingUser, after: LevellingUser):
+        if after.total_xp == before.total_xp:
+            return
+
+        elif after.total_xp > before.total_xp:
+            query = """
+                SELECT COALESCE(array_agg(role_id), '{}')
+                  FROM levelling_roles lr
+                 WHERE lr.guild_id = $1
+                   AND lr.required_xp <= $2
+                   AND lr.role_id NOT IN (
+                    SELECT pr.role_id
+                      FROM persisted_roles pr
+                     WHERE pr.guild_id = lr.guild_id
+                       AND pr.user_id = $3
+                   )
+            """
+            # Fetch role ids that the user qualifies for, but have not been persisted.
+            role_ids = await LevellingRole.fetchval(query, after.guild_id, after.total_xp, after.user_id)
+
+            if not role_ids:
+                return
+
+            self.bot.dispatch("persist_roles", guild_id=after.guild_id, user_id=after.user_id, role_ids=role_ids)
+
+        else:
+            query = """
+                SELECT COALESCE(array_agg(role_id), '{}')
+                  FROM levelling_roles lr
+                 WHERE lr.guild_id = $1
+                   AND lr.required_xp > $2
+                   AND lr.role_id IN (
+                    SELECT pr.role_id
+                      FROM persisted_roles pr
+                     WHERE pr.guild_id = lr.guild_id
+                       AND pr.user_id = $3
+                   )
+            """
+
+            role_ids = await LevellingRole.fetchval(query, after.guild_id, after.total_xp, after.user_id)
+
+            if not role_ids:
+                return
+
+            self.bot.dispatch(
+                "remove_persisted_roles", guild_id=after.guild_id, user_id=after.user_id, role_ids=role_ids
+            )

bot/extensions/levelling/utils.py

@@ -0,0 +1,15 @@
+from bisect import bisect
+
+required_xp = [0]
+
+for lvl in range(1001):
+    xp = 5 * (lvl**2) + (50 * lvl) + 100
+    required_xp.append(xp + required_xp[-1])
+
+
+def get_level_for_xp(user_xp: int) -> int:
+    return bisect(required_xp, user_xp) - 1
+
+
+def get_xp_for_level(level: int) -> int:
+    return required_xp[level]

bot/extensions/persistent_roles/events.py

@@ -1,7 +1,8 @@
+import discord
 from discord.ext import commands
 
 from bot import core
-from bot.models import PersistentRole
+from bot.models import PersistedRole
 
 
 class PersistentEvents(commands.Cog):
@@ -11,14 +12,81 @@ def __init__(self, bot: core.DiscordBot):
         self.bot = bot
 
     @commands.Cog.listener()
-    async def on_member_join(self, member):
-        """Add the persisted role to users if any on member join"""
-        # Get the data
-        data = await PersistentRole.list_by_guild(member.guild.id, member.id)
-        # Return if data for specified user and guild does not exist
-        if data is None:
+    async def on_remove_persisted_roles(self, guild_id: int, user_id: int, role_ids: list[int]):
+        """This event is called to remove persisted roles from a user."""
+        query = "DELETE FROM persisted_roles WHERE guild_id = $1 AND user_id = $2 AND role_id = ANY($3)"
+        await PersistedRole.execute(query, guild_id, user_id, role_ids)
+
+        guild = self.bot.get_guild(guild_id)
+
+        if guild is None:
+            return
+
+        member = guild.get_member(user_id)
+
+        if member is None:
+            return
+
+        roles = []
+
+        for role_id in role_ids:
+            role = member.get_role(role_id)
+
+            if role is None:
+                continue
+
+            roles.append(role)
+
+        await member.remove_roles(*roles, atomic=True)
+
+    @commands.Cog.listener()
+    async def on_persist_roles(self, guild_id: int, user_id: int, role_ids: list[int]):
+        """This event is called to trigger persist of roles."""
+        query = "INSERT INTO persisted_roles (guild_id, user_id, role_id) VALUES ($1, $2, $3)"
+        data = [(guild_id, user_id, role_id) for role_id in role_ids]
+        await PersistedRole.pool.executemany(query, data)
+
+        guild = self.bot.get_guild(guild_id)
+
+        if guild is None:
             return
-        # Add the roles to the user if data exists
-        else:
-            for i in range(len(data)):
-                await member.add_roles(member.guild.get_role(data[i].role_id))
+
+        member = guild.get_member(user_id)
+
+        if member is None:
+            return
+
+        roles = []
+
+        for role_id in role_ids:
+            role = guild.get_role(role_id)
+
+            if role is None:
+                continue
+
+            roles.append(role)
+
+        await member.add_roles(*roles, atomic=True)
+
+    @commands.Cog.listener()
+    async def on_member_join(self, member: discord.Member):
+        """Add the persisted role to users if any on member join"""
+        query = """
+            SELECT COALESCE(array_agg(role_id), '{}')
+              FROM persisted_roles
+             WHERE guild_id = $1
+               AND user_id = $2
+        """
+        role_ids = await PersistedRole.fetchval(query, member.guild.id, member.id)
+
+        roles = []
+
+        for role_id in role_ids:
+            role = member.guild.get_role(role_id)
+
+            if role is None:
+                continue
+
+            roles.append(role)
+
+        await member.add_roles(*roles)

bot/models/__init__.py

@@ -1,10 +1,10 @@
 from .gconfig import FilterConfig
 from .levelling_ignored_channels import IgnoredChannel
 from .levelling_roles import LevellingRole
-from .levelling_users import Levels
+from .levelling_users import LevellingUser
 from .message import Message
 from .model import Model
-from .persisted_role import PersistentRole
+from .persisted_role import PersistedRole
 from .rep import Rep
 from .tag import Tag
 from .user import User
@@ -16,8 +16,8 @@
     Rep,
     Tag,
     User,
-    Levels,
-    PersistentRole,
+    LevellingUser,
+    PersistedRole,
     IgnoredChannel,
     LevellingRole,
 )  # Fixes F401

bot/models/custom_roles.py

@@ -1,9 +1,24 @@
 from .model import Model
 
 
-class CustomRoles(Model):
+class CustomRole(Model):
     id: int
     guild_id: int
     role_id: int
     name: str
     color: str
+
+    @classmethod
+    async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: str):
+        """Inserts or updates the custom role."""
+        query = """
+            INSERT INTO custom_roles (guild_id, role_id, name, color)
+                 VALUES ($1, $2, $3, $4)
+            ON CONFLICT (guild_id, role_id)
+              DO UPDATE SET
+                   name = $3,
+                  color = $4
+             RETURNING *
+        """
+
+        return await cls.fetchrow(query, guild_id, role_id, name, color)

bot/models/migrations/003_up__snowflake_function.sql

@@ -25,7 +25,7 @@ CREATE OR REPLACE FUNCTION snowflake_to_timestamp(flake BIGINT)
     LANGUAGE 'plpgsql'
 AS $BODY$
 DECLARE
-    our_epoch BIGINT := 1621470600;
+    our_epoch BIGINT := 1609459200;
 BEGIN
     RETURN to_timestamp(((flake >> 22) + our_epoch) / 1000);
 END;

bot/models/migrations/004_up__levelling.sql

@@ -22,7 +22,7 @@ CREATE TABLE IF NOT EXISTS custom_roles
 
 CREATE TABLE IF NOT EXISTS levelling_roles
 (
-    id          SERIAL PRIMARY KEY DEFAULT create_snowflake(),
+    id          BIGINT PRIMARY KEY DEFAULT create_snowflake(),
     required_xp INTEGER NOT NULL,
     guild_id    BIGINT NOT NULL,
     role_id     BIGINT NOT NULL REFERENCES custom_roles(role_id),
@@ -31,7 +31,7 @@ CREATE TABLE IF NOT EXISTS levelling_roles
 
 CREATE TABLE IF NOT EXISTS persisted_roles
 (
-    id                 SERIAL PRIMARY KEY DEFAULT create_snowflake(),
+    id                 BIGINT PRIMARY KEY DEFAULT create_snowflake(),
     guild_id           BIGINT NOT NULL,
     user_id            BIGINT NOT NULL,
     role_id            BIGINT NOT NULL REFERENCES custom_roles(role_id),
@@ -41,7 +41,7 @@ CREATE TABLE IF NOT EXISTS persisted_roles
 
 CREATE TABLE IF NOT EXISTS levelling_ignored_channels
 (
-    id          SERIAL PRIMARY KEY DEFAULT create_snowflake(),
+    id          BIGINT PRIMARY KEY DEFAULT create_snowflake(),
     guild_id    BIGINT NOT NULL,
     channel_id  BIGINT NOT NULL,
     CONSTRAINT levelling_ignored_channels_guild_id_and_channel_id_key UNIQUE (guild_id, channel_id)

bot/models/model.py

@@ -9,6 +9,16 @@
 log = logging.getLogger(__name__)
 
 
+class CustomRecord(Record):
+    def __getattr__(self, item: str):
+        try:
+            return self[item]
+        except KeyError:
+            pass
+
+        return super().__getattr__(item)
+
+
 class Model(BaseModel):
     pool: ClassVar[Pool] = None
 
@@ -22,7 +32,9 @@ async def create_pool(
         loop: asyncio.AbstractEventLoop = None,
         **kwargs,
     ) -> None:
-        cls.pool = await create_pool(uri, min_size=min_con, max_size=max_con, loop=loop, **kwargs)
+        cls.pool = await create_pool(
+            uri, min_size=min_con, max_size=max_con, loop=loop, record_class=CustomRecord, **kwargs
+        )
         log.info(f"Established a pool with {min_con} - {max_con} connections\n")
 
     @classmethod
@@ -46,9 +58,12 @@ async def fetchrow(
     ) -> Union[BM, Record, None]:
         if con is None:
             con = cls.pool
+
         record = await con.fetchrow(query, *args)
+
         if cls is Model or record is None or convert is False:
             return record
+
         return cls(**record)
 
     @classmethod