From b0163ff9ae5dacf7088096d2351d6fa49b87c41f Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Wed, 11 Oct 2023 17:46:04 +0530 Subject: [PATCH 01/10] added persistent role,levelling features --- bot/cogs/levelling.py | 238 ++++++++++++++++++ bot/cogs/persistent_roles.py | 24 ++ bot/extensions/challenges/events.py | 15 ++ bot/models/__init__.py | 12 +- bot/models/ignored_channel.py | 23 ++ bot/models/levelling_role.py | 12 + bot/models/levels.py | 49 ++++ bot/models/migrations/001_down__delete_db.sql | 2 +- bot/models/migrations/001_up__create_db.sql | 48 +++- bot/models/persistent_role.py | 28 +++ cli.py | 2 + 11 files changed, 449 insertions(+), 4 deletions(-) create mode 100644 bot/cogs/levelling.py create mode 100644 bot/cogs/persistent_roles.py create mode 100644 bot/models/ignored_channel.py create mode 100644 bot/models/levelling_role.py create mode 100644 bot/models/levels.py create mode 100644 bot/models/persistent_role.py diff --git a/bot/cogs/levelling.py b/bot/cogs/levelling.py new file mode 100644 index 00000000..dec349b1 --- /dev/null +++ b/bot/cogs/levelling.py @@ -0,0 +1,238 @@ +import datetime +import random + +import discord +from discord import app_commands +from discord.ext import commands +from discord.ext.commands import has_permissions + +from bot import core +from bot.models import IgnoredChannel, Levels + +REQUIRED_XP = [ + 0, + 100, + 255, + 475, + 770, + 1150, + 1625, + 2205, + 2900, + 3720, + 4675, + 5775, + 7030, + 8450, + 10045, + 11825, + 13800, + 15980, + 18375, + 20995, + 23850, + 26950, + 30305, + 33925, + 37820, + 42000, + 46475, + 51255, + 56350, + 61770, + 67525, + 73625, + 80080, + 86900, + 94095, + 101675, + 109650, + 118030, + 126825, + 136045, + 145700, + 155800, + 166355, + 177375, + 188870, + 200850, + 213325, + 226305, + 239800, + 253820, + 268375, + 283475, + 299130, + 315350, + 332145, + 349525, + 367500, + 386080, + 405275, + 425095, + 445550, + 466650, + 488405, + 510825, + 533920, + 557700, + 582175, + 607355, + 633250, + 659870, + 687225, + 715325, + 744180, + 773800, + 804195, + 835375, + 867350, + 900130, + 933725, + 968145, + 1003400, + 1039500, + 1076455, + 1114275, + 1152970, + 1192550, + 1233025, + 1274405, + 1316700, + 1359920, + 1404075, + 1449175, + 1495230, + 1542250, + 1590245, + 1639225, + 1689200, + 1740180, + 1899250, + 1792175, + 1845195, +] +IGNORED_CHANNEL = {} + + +class Levelling(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @commands.Cog.listener() + async def on_ready(self): + for guild in self.bot.guilds: + channel_id = await IgnoredChannel.get(guild_id=guild.id) + if guild.id not in IGNORED_CHANNEL: + IGNORED_CHANNEL[guild.id] = [channel_id] + else: + IGNORED_CHANNEL[guild.id].append(channel_id) + print(IGNORED_CHANNEL) + + @commands.Cog.listener() + async def on_message(self, message): + if message.author.bot: + return + + # Check if message is sent in ignored channel + ignored_channel = await IgnoredChannel.get(message.guild.id) + for i in range(len(ignored_channel)): + if message.channel.id == ignored_channel[i].channel_id: + return + + # Generate random XP to be added + xp = random.randint(5, 25) + # Add the XP and update the DB + data = await Levels.update(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) + await self._check_level_up(data, message.author) + + @app_commands.command() + async def rank(self, interaction: core.InteractionType, member: discord.Member = None): + """Check the rank of another member or yourself + example: + - /rank @Noobmaster + - /rank""" + if member is None: + member = interaction.user + + query = """WITH ordered_users AS ( + SELECT + user_id,guild_id,total_xp, + ROW_NUMBER() OVER (ORDER BY levels.total_xp DESC) AS rank + FROM levels + WHERE guild_id = $2) + SELECT rank,total_xp,user_id,guild_id + FROM ordered_users WHERE ordered_users.user_id = $1;""" + data = await Levels.fetchrow(query, member.id, member.guild.id) + if data is None: + return await interaction.response.send_message(f"{member} is not ranked yet!") + for level, j in enumerate(REQUIRED_XP): + if data.total_xp <= j: + embed = discord.Embed( + title=f"Rank: {data.rank}\nLevel: {level - 1}\nTotal XP:{data.total_xp}", + timestamp=datetime.datetime.utcnow(), + colour=discord.Colour.blurple(), + ) + embed.set_thumbnail(url=member.avatar) + return await interaction.response.send_message(embed=embed) + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Add the channel to the ignored channel list to not gain XP + Example: + /ignore_channel #channel""" + await IgnoredChannel.insert(channel.guild.id, channel.id) + await interaction.response.send_message(f"{channel} has been ignored from gaining XP.") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + @has_permissions(administrator=True) + async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Remove channel from ignored channel list + Example: + /unignore_channel #channel""" + await IgnoredChannel.delete(channel.guild.id, channel.id) + await interaction.response.send_message(f"{channel} has been removed from ignored channel list") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): + """Give XP to specific user + Example: + /give_xp 1000 user""" + data = await Levels.update(member.guild.id, member.id, xp) + await self._check_level_up(data, member) + await interaction.response.send_message(f"{xp} XP has been added to user {member}") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): + """Remove XP from user + Example: + /remove_xp 100 user""" + data = await Levels.remove_xp(member.guild.id, member.id, xp) + await self._check_level_up(data, member) + await interaction.response.send_message(f"{xp} XP has been removed from user {member}") + + async def _check_level_up(self, data, member: discord.Member): + """Function to check if user's level has changed and trigger the event to assign the roles""" + # Calculating old and new level + try: + for level, j in enumerate(REQUIRED_XP): + if data.old_total_xp <= j: + old_level = level - 1 + break + for level, j in enumerate(REQUIRED_XP): + if data.total_xp <= j: + new_level = level - 1 + break + # If the level has changed call level_up to handle roles change + if old_level != new_level: + self.bot.dispatch("level_up", new_level=new_level, member=member) + except AttributeError: + pass + + +async def setup(bot: commands.Bot): + await bot.add_cog(Levelling(bot=bot)) diff --git a/bot/cogs/persistent_roles.py b/bot/cogs/persistent_roles.py new file mode 100644 index 00000000..60ea416f --- /dev/null +++ b/bot/cogs/persistent_roles.py @@ -0,0 +1,24 @@ +from discord.ext import commands + +from bot.models import PersistentRole + + +class PersistentRoles(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @commands.Cog.listener() + async def on_member_join(self, member): + # Get the data + data = await PersistentRole.get(member.guild.id, member.id) + # Return if data for specified user and guild does not exist + if data 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)) + + +async def setup(bot: commands.Bot): + await bot.add_cog(PersistentRoles(bot=bot)) diff --git a/bot/extensions/challenges/events.py b/bot/extensions/challenges/events.py index b88431ae..06574a2a 100644 --- a/bot/extensions/challenges/events.py +++ b/bot/extensions/challenges/events.py @@ -5,6 +5,7 @@ from bot import core from bot.config import settings +from bot.models import LevellingRole, PersistentRole log = logging.getLogger(__name__) @@ -48,3 +49,17 @@ async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent): await self.bot.http.add_role( guild_id=self.bot.guild.id, user_id=payload.user_id, role_id=self.participant_role.id ) + + @commands.Cog.listener() + async def on_level_up(self, new_level, member: discord.Member): + """Add roles when user levels up""" + data = await LevellingRole.get(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(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.remove(guild_id=member.guild.id, user_id=member.id) diff --git a/bot/models/__init__.py b/bot/models/__init__.py index ab0e68f6..ab9715e0 100644 --- a/bot/models/__init__.py +++ b/bot/models/__init__.py @@ -1,15 +1,23 @@ from .gconfig import FilterConfig +from .ignored_channel import IgnoredChannel +from .levelling_role import LevellingRole +from .levels import Levels from .message import Message from .model import Model +from .persistent_role import PersistentRole from .rep import Rep from .tag import Tag from .user import User -__all__ = ( # Fixes F401 +__all__ = ( Model, FilterConfig, Message, Rep, Tag, User, -) + Levels, + PersistentRole, + IgnoredChannel, + LevellingRole, +) # Fixes F401 diff --git a/bot/models/ignored_channel.py b/bot/models/ignored_channel.py new file mode 100644 index 00000000..ad11b1b3 --- /dev/null +++ b/bot/models/ignored_channel.py @@ -0,0 +1,23 @@ +from .model import Model + + +class IgnoredChannel(Model): + guild_id: int + channel_id: int + + @classmethod + async def get(cls, guild_id: int): + query = """ SELECT * FROM ignored_channels WHERE guild_id = $1""" + return await cls.fetch(query, guild_id) + + @classmethod + async def insert(cls, guild_id: int, channel_id: int): + query = """INSERT INTO ignored_channels (guild_id,channel_id) VALUES($1, $2) + ON CONFLICT (guild_id,channel_id) DO NOTHING + RETURNING guild_id, channel_id""" + return await cls.fetch(query, guild_id, channel_id) + + @classmethod + async def delete(cls, guild_id: int, channel_id: int): + query = """DELETE FROM ignored_channel WHERE guild_id= $1 and channel_id = $2""" + return await cls.fetch(query, guild_id, channel_id) diff --git a/bot/models/levelling_role.py b/bot/models/levelling_role.py new file mode 100644 index 00000000..33957427 --- /dev/null +++ b/bot/models/levelling_role.py @@ -0,0 +1,12 @@ +from .model import Model + + +class LevellingRole(Model): + guild_id: int + role_id: int + level: int + + @classmethod + async def get(cls, guild_id: int): + query = """SELECT guild_id, role_id, level FROM levelling_roles WHERE guild_id=$1""" + return await cls.fetch(query, guild_id) diff --git a/bot/models/levels.py b/bot/models/levels.py new file mode 100644 index 00000000..9ca9e23e --- /dev/null +++ b/bot/models/levels.py @@ -0,0 +1,49 @@ +from typing import Optional + +from .model import Model + + +class Levels(Model): + rank: Optional[int] + guild_id: int + user_id: int + old_total_xp: Optional[int] + total_xp: int + + @classmethod + async def update(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levels (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE + SET total_xp = levels.total_xp + $3, last_msg = CURRENT_TIMESTAMP + WHERE + levels.guild_id = $1 + AND levels.user_id = $2 + AND EXTRACT(EPOCH FROM (NOW() - levels.last_msg)) > 60 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id; +""" + + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def give_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """ INSERT INTO levels (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levels.total_xp + $3 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def remove_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levels (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levels.total_xp - $3 + RETURNING total_xp + $3 AS old_total_xp, total_xp, guild_id, user_id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) diff --git a/bot/models/migrations/001_down__delete_db.sql b/bot/models/migrations/001_down__delete_db.sql index dade8dc5..4eb41fa7 100644 --- a/bot/models/migrations/001_down__delete_db.sql +++ b/bot/models/migrations/001_down__delete_db.sql @@ -1 +1 @@ -DROP TABLE gconfigs, messages, reps, tags, users, polls; +DROP TABLE gconfigs, messages, reps, tags, users, polls, levels, roles, persistent_roles, ignored_channels, levelling_roles; diff --git a/bot/models/migrations/001_up__create_db.sql b/bot/models/migrations/001_up__create_db.sql index cc268740..92f21ced 100644 --- a/bot/models/migrations/001_up__create_db.sql +++ b/bot/models/migrations/001_up__create_db.sql @@ -55,4 +55,50 @@ CREATE TABLE IF NOT EXISTS polls created_at TIMESTAMP, channel_id BIGINT, message_id BIGINT PRIMARY KEY -) +); + +CREATE TABLE IF NOT EXISTS levels +( + guild_id BIGINT, + user_id BIGINT, + total_xp INT DEFAULT 0, + last_msg TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT PK_user PRIMARY KEY(guild_id,user_id) +); + +CREATE TABLE IF NOT EXISTS roles +( + role_id BIGINT PRIMARY KEY, + guild_id BIGINT, + name VARCHAR, + color VARCHAR +); + +CREATE TABLE IF NOT EXISTS persistent_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + user_id BIGINT, + role_id BIGINT, + CONSTRAINT fk_roles FOREIGN KEY(role_id) REFERENCES roles(role_id), + CONSTRAINT unique_guild_user_role UNIQUE (guild_id, user_id, role_id) +); + + +CREATE TABLE IF NOT EXISTS ignored_channels +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + channel_id BIGINT, + CONSTRAINT unique_ignored_channel UNIQUE (guild_id, channel_id) + +); + +CREATE TABLE IF NOT EXISTS levelling_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + role_id BIGINT, + level INT, + CONSTRAINT fk_levelling_roles FOREIGN KEY(role_id) REFERENCES roles(role_id) +); diff --git a/bot/models/persistent_role.py b/bot/models/persistent_role.py new file mode 100644 index 00000000..077ce0a5 --- /dev/null +++ b/bot/models/persistent_role.py @@ -0,0 +1,28 @@ +from typing import Optional + +from .model import Model + + +class PersistentRole(Model): + guild_id: int + user_id: int + role_id: Optional[int] + + @classmethod + async def get(cls, guild_id: int, user_id: int): + query = """SELECT * FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" + return await cls.fetch(query, guild_id, user_id) + + @classmethod + async def insert(cls, guild_id: int, user_id: int, role_id: int): + query = """INSERT INTO persistent_roles (guild_id, user_id, role_id) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id, role_id) + DO NOTHING""" + + await cls.execute(query, guild_id, user_id, role_id) + + @classmethod + async def remove(cls, guild_id: int, user_id: int): + query = """DELETE FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" + return await cls.fetch(query, guild_id, user_id) diff --git a/cli.py b/cli.py index e3d101c8..359c449f 100644 --- a/cli.py +++ b/cli.py @@ -139,6 +139,8 @@ async def main(ctx): "bot.cogs.clashofcode", "bot.cogs.roles", "bot.cogs.poll", + "bot.cogs.levelling", + "bot.cogs.persistent_roles", ) intents = discord.Intents.all() From 978b60824f19d65e22acc9bdd0eb540927b1c5b6 Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Wed, 11 Oct 2023 21:31:41 +0530 Subject: [PATCH 02/10] added ignored_channel as attribute, moved cache logic to cog_load event and added condition to check message was not sent to bot DM --- bot/extensions/__init__.py | 0 bot/extensions/levelling/__init__.py | 0 .../levelling/commands.py} | 49 +++++++------------ 3 files changed, 18 insertions(+), 31 deletions(-) create mode 100644 bot/extensions/__init__.py create mode 100644 bot/extensions/levelling/__init__.py rename bot/{cogs/levelling.py => extensions/levelling/commands.py} (83%) diff --git a/bot/extensions/__init__.py b/bot/extensions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/extensions/levelling/__init__.py b/bot/extensions/levelling/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/cogs/levelling.py b/bot/extensions/levelling/commands.py similarity index 83% rename from bot/cogs/levelling.py rename to bot/extensions/levelling/commands.py index dec349b1..2816d765 100644 --- a/bot/cogs/levelling.py +++ b/bot/extensions/levelling/commands.py @@ -112,33 +112,31 @@ 1792175, 1845195, ] -IGNORED_CHANNEL = {} class Levelling(commands.Cog): def __init__(self, bot): self.bot = bot + self.ignored_channel = {} - @commands.Cog.listener() - async def on_ready(self): + async def cog_load(self): for guild in self.bot.guilds: - channel_id = await IgnoredChannel.get(guild_id=guild.id) - if guild.id not in IGNORED_CHANNEL: - IGNORED_CHANNEL[guild.id] = [channel_id] - else: - IGNORED_CHANNEL[guild.id].append(channel_id) - print(IGNORED_CHANNEL) + data = await IgnoredChannel.get(guild_id=guild.id) + for i in data: + if guild.id not in self.ignored_channel: + self.ignored_channel[guild.id] = [i.channel_id] + else: + self.ignored_channel[guild.id].append(i.channel_id) @commands.Cog.listener() async def on_message(self, message): - if message.author.bot: + # Return if message was sent by Bot or sent in DMs + if message.author.bot or message.guild is None: return # Check if message is sent in ignored channel - ignored_channel = await IgnoredChannel.get(message.guild.id) - for i in range(len(ignored_channel)): - if message.channel.id == ignored_channel[i].channel_id: - return + if message.channel.id in self.ignored_channel[message.guild.id]: + return # Generate random XP to be added xp = random.randint(5, 25) @@ -148,10 +146,7 @@ async def on_message(self, message): @app_commands.command() async def rank(self, interaction: core.InteractionType, member: discord.Member = None): - """Check the rank of another member or yourself - example: - - /rank @Noobmaster - - /rank""" + """Check the rank of another member or yourself""" if member is None: member = interaction.user @@ -165,7 +160,7 @@ async def rank(self, interaction: core.InteractionType, member: discord.Member = FROM ordered_users WHERE ordered_users.user_id = $1;""" data = await Levels.fetchrow(query, member.id, member.guild.id) if data is None: - return await interaction.response.send_message(f"{member} is not ranked yet!") + return await interaction.response.send_message("You are not ranked yet!", ephemeral=True) for level, j in enumerate(REQUIRED_XP): if data.total_xp <= j: embed = discord.Embed( @@ -179,9 +174,7 @@ async def rank(self, interaction: core.InteractionType, member: discord.Member = @app_commands.command() @app_commands.checks.has_permissions(administrator=True) async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): - """Add the channel to the ignored channel list to not gain XP - Example: - /ignore_channel #channel""" + """Add the channel to the ignored channel list to not gain XP""" await IgnoredChannel.insert(channel.guild.id, channel.id) await interaction.response.send_message(f"{channel} has been ignored from gaining XP.") @@ -189,18 +182,14 @@ async def ignore_channel(self, interaction: core.InteractionType, channel: disco @app_commands.checks.has_permissions(administrator=True) @has_permissions(administrator=True) async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): - """Remove channel from ignored channel list - Example: - /unignore_channel #channel""" + """Remove channel from ignored channel list""" await IgnoredChannel.delete(channel.guild.id, channel.id) await interaction.response.send_message(f"{channel} has been removed from ignored channel list") @app_commands.command() @app_commands.checks.has_permissions(administrator=True) async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): - """Give XP to specific user - Example: - /give_xp 1000 user""" + """Give XP to specific user""" data = await Levels.update(member.guild.id, member.id, xp) await self._check_level_up(data, member) await interaction.response.send_message(f"{xp} XP has been added to user {member}") @@ -208,9 +197,7 @@ async def give_xp(self, interaction: core.InteractionType, xp: int, member: disc @app_commands.command() @app_commands.checks.has_permissions(administrator=True) async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): - """Remove XP from user - Example: - /remove_xp 100 user""" + """Remove XP from user""" data = await Levels.remove_xp(member.guild.id, member.id, xp) await self._check_level_up(data, member) await interaction.response.send_message(f"{xp} XP has been removed from user {member}") From 75283dc1cfc026efae55f16dc3db33ee390d0855 Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Fri, 13 Oct 2023 16:59:19 +0530 Subject: [PATCH 03/10] renaming the model variable to be more meaningful and refactoring the cog to work be in extensions format --- bot/cogs/persistent_roles.py | 2 +- bot/extensions/challenges/events.py | 15 --------------- bot/extensions/levelling/__init__.py | 9 +++++++++ bot/extensions/levelling/commands.py | 6 +++--- bot/extensions/levelling/events.py | 28 ++++++++++++++++++++++++++++ bot/models/ignored_channel.py | 12 ++++++------ bot/models/levelling_role.py | 2 +- bot/models/persistent_role.py | 12 ++++++------ cli.py | 2 +- 9 files changed, 55 insertions(+), 33 deletions(-) create mode 100644 bot/extensions/levelling/events.py diff --git a/bot/cogs/persistent_roles.py b/bot/cogs/persistent_roles.py index 60ea416f..f7b9a305 100644 --- a/bot/cogs/persistent_roles.py +++ b/bot/cogs/persistent_roles.py @@ -10,7 +10,7 @@ def __init__(self, bot): @commands.Cog.listener() async def on_member_join(self, member): # Get the data - data = await PersistentRole.get(member.guild.id, member.id) + 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: return diff --git a/bot/extensions/challenges/events.py b/bot/extensions/challenges/events.py index 06574a2a..b88431ae 100644 --- a/bot/extensions/challenges/events.py +++ b/bot/extensions/challenges/events.py @@ -5,7 +5,6 @@ from bot import core from bot.config import settings -from bot.models import LevellingRole, PersistentRole log = logging.getLogger(__name__) @@ -49,17 +48,3 @@ async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent): await self.bot.http.add_role( guild_id=self.bot.guild.id, user_id=payload.user_id, role_id=self.participant_role.id ) - - @commands.Cog.listener() - async def on_level_up(self, new_level, member: discord.Member): - """Add roles when user levels up""" - data = await LevellingRole.get(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(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.remove(guild_id=member.guild.id, user_id=member.id) diff --git a/bot/extensions/levelling/__init__.py b/bot/extensions/levelling/__init__.py index e69de29b..0e10853a 100644 --- a/bot/extensions/levelling/__init__.py +++ b/bot/extensions/levelling/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .commands import Levelling +from .events import LevelEvents + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(Levelling(bot=bot)) + await bot.add_cog(LevelEvents(bot=bot)) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py index 2816d765..d94ca5a3 100644 --- a/bot/extensions/levelling/commands.py +++ b/bot/extensions/levelling/commands.py @@ -121,7 +121,7 @@ def __init__(self, bot): async def cog_load(self): for guild in self.bot.guilds: - data = await IgnoredChannel.get(guild_id=guild.id) + data = await IgnoredChannel.list_by_guild(guild_id=guild.id) for i in data: if guild.id not in self.ignored_channel: self.ignored_channel[guild.id] = [i.channel_id] @@ -175,7 +175,7 @@ async def rank(self, interaction: core.InteractionType, member: discord.Member = @app_commands.checks.has_permissions(administrator=True) async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): """Add the channel to the ignored channel list to not gain XP""" - await IgnoredChannel.insert(channel.guild.id, channel.id) + await IgnoredChannel.insert_by_guild(channel.guild.id, channel.id) await interaction.response.send_message(f"{channel} has been ignored from gaining XP.") @app_commands.command() @@ -183,7 +183,7 @@ async def ignore_channel(self, interaction: core.InteractionType, channel: disco @has_permissions(administrator=True) async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): """Remove channel from ignored channel list""" - await IgnoredChannel.delete(channel.guild.id, channel.id) + await IgnoredChannel.delete_by_guild(channel.guild.id, channel.id) await interaction.response.send_message(f"{channel} has been removed from ignored channel list") @app_commands.command() diff --git a/bot/extensions/levelling/events.py b/bot/extensions/levelling/events.py new file mode 100644 index 00000000..acfa410f --- /dev/null +++ b/bot/extensions/levelling/events.py @@ -0,0 +1,28 @@ +import discord +from discord.ext import commands + +from bot import core +from bot.models import LevellingRole, PersistentRole + + +class LevelEvents(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) diff --git a/bot/models/ignored_channel.py b/bot/models/ignored_channel.py index ad11b1b3..85f4e402 100644 --- a/bot/models/ignored_channel.py +++ b/bot/models/ignored_channel.py @@ -6,18 +6,18 @@ class IgnoredChannel(Model): channel_id: int @classmethod - async def get(cls, guild_id: int): - query = """ SELECT * FROM ignored_channels WHERE guild_id = $1""" + async def list_by_guild(cls, guild_id: int): + query = """SELECT * FROM ignored_channels WHERE guild_id = $1""" return await cls.fetch(query, guild_id) @classmethod - async def insert(cls, guild_id: int, channel_id: int): + async def insert_by_guild(cls, guild_id: int, channel_id: int): query = """INSERT INTO ignored_channels (guild_id,channel_id) VALUES($1, $2) - ON CONFLICT (guild_id,channel_id) DO NOTHING - RETURNING guild_id, channel_id""" + ON CONFLICT (guild_id,channel_id) DO NOTHING + RETURNING guild_id, channel_id""" return await cls.fetch(query, guild_id, channel_id) @classmethod - async def delete(cls, guild_id: int, channel_id: int): + async def delete_by_guild(cls, guild_id: int, channel_id: int): query = """DELETE FROM ignored_channel WHERE guild_id= $1 and channel_id = $2""" return await cls.fetch(query, guild_id, channel_id) diff --git a/bot/models/levelling_role.py b/bot/models/levelling_role.py index 33957427..c9ee430b 100644 --- a/bot/models/levelling_role.py +++ b/bot/models/levelling_role.py @@ -7,6 +7,6 @@ class LevellingRole(Model): level: int @classmethod - async def get(cls, guild_id: int): + async def list_by_guild(cls, guild_id: int): query = """SELECT guild_id, role_id, level FROM levelling_roles WHERE guild_id=$1""" return await cls.fetch(query, guild_id) diff --git a/bot/models/persistent_role.py b/bot/models/persistent_role.py index 077ce0a5..599666fc 100644 --- a/bot/models/persistent_role.py +++ b/bot/models/persistent_role.py @@ -9,20 +9,20 @@ class PersistentRole(Model): role_id: Optional[int] @classmethod - async def get(cls, guild_id: int, user_id: int): + async def list_by_guild(cls, guild_id: int, user_id: int): query = """SELECT * FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" return await cls.fetch(query, guild_id, user_id) @classmethod - async def insert(cls, guild_id: int, user_id: int, role_id: int): + async def insert_by_guild(cls, guild_id: int, user_id: int, role_id: int): query = """INSERT INTO persistent_roles (guild_id, user_id, role_id) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id, role_id) - DO NOTHING""" + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id, role_id) + DO NOTHING""" await cls.execute(query, guild_id, user_id, role_id) @classmethod - async def remove(cls, guild_id: int, user_id: int): + async def delete_by_guild(cls, guild_id: int, user_id: int): query = """DELETE FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" return await cls.fetch(query, guild_id, user_id) diff --git a/cli.py b/cli.py index 359c449f..8f41df78 100644 --- a/cli.py +++ b/cli.py @@ -139,7 +139,7 @@ async def main(ctx): "bot.cogs.clashofcode", "bot.cogs.roles", "bot.cogs.poll", - "bot.cogs.levelling", + "bot.extensions.levelling", "bot.cogs.persistent_roles", ) From 9d1177c88031edd6a7c1e76047f7226fcea0ceee Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Sun, 15 Oct 2023 11:40:52 +0530 Subject: [PATCH 04/10] refactored code, fixed models and added new commands for levelling role --- bot/extensions/levelling/commands.py | 196 +++++------------- bot/extensions/levelling/events.py | 15 ++ bot/extensions/persistent_roles/__init__.py | 9 + bot/extensions/persistent_roles/commands.py | 10 + .../persistent_roles/events.py} | 12 +- bot/models/__init__.py | 8 +- bot/models/custom_roles.py | 20 ++ ...annel.py => levelling_ignored_channels.py} | 9 +- bot/models/levelling_role.py | 12 -- bot/models/levelling_roles.py | 26 +++ bot/models/levelling_users.py | 48 +++++ bot/models/levels.py | 49 ----- bot/models/migrations/001_down__delete_db.sql | 2 +- bot/models/migrations/001_up__create_db.sql | 48 +---- .../migrations/002_down__levelling_db.sql | 1 + .../migrations/002_up__levelling_db.sql | 48 +++++ .../{persistent_role.py => persisted_role.py} | 9 +- cli.py | 2 +- 18 files changed, 252 insertions(+), 272 deletions(-) create mode 100644 bot/extensions/persistent_roles/__init__.py create mode 100644 bot/extensions/persistent_roles/commands.py rename bot/{cogs/persistent_roles.py => extensions/persistent_roles/events.py} (72%) create mode 100644 bot/models/custom_roles.py rename bot/models/{ignored_channel.py => levelling_ignored_channels.py} (61%) delete mode 100644 bot/models/levelling_role.py create mode 100644 bot/models/levelling_roles.py create mode 100644 bot/models/levelling_users.py delete mode 100644 bot/models/levels.py create mode 100644 bot/models/migrations/002_down__levelling_db.sql create mode 100644 bot/models/migrations/002_up__levelling_db.sql rename bot/models/{persistent_role.py => persisted_role.py} (67%) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py index d94ca5a3..b7d5046a 100644 --- a/bot/extensions/levelling/commands.py +++ b/bot/extensions/levelling/commands.py @@ -1,123 +1,21 @@ import datetime import random +import asyncpg.exceptions import discord from discord import app_commands from discord.ext import commands -from discord.ext.commands import has_permissions from bot import core -from bot.models import IgnoredChannel, Levels - -REQUIRED_XP = [ - 0, - 100, - 255, - 475, - 770, - 1150, - 1625, - 2205, - 2900, - 3720, - 4675, - 5775, - 7030, - 8450, - 10045, - 11825, - 13800, - 15980, - 18375, - 20995, - 23850, - 26950, - 30305, - 33925, - 37820, - 42000, - 46475, - 51255, - 56350, - 61770, - 67525, - 73625, - 80080, - 86900, - 94095, - 101675, - 109650, - 118030, - 126825, - 136045, - 145700, - 155800, - 166355, - 177375, - 188870, - 200850, - 213325, - 226305, - 239800, - 253820, - 268375, - 283475, - 299130, - 315350, - 332145, - 349525, - 367500, - 386080, - 405275, - 425095, - 445550, - 466650, - 488405, - 510825, - 533920, - 557700, - 582175, - 607355, - 633250, - 659870, - 687225, - 715325, - 744180, - 773800, - 804195, - 835375, - 867350, - 900130, - 933725, - 968145, - 1003400, - 1039500, - 1076455, - 1114275, - 1152970, - 1192550, - 1233025, - 1274405, - 1316700, - 1359920, - 1404075, - 1449175, - 1495230, - 1542250, - 1590245, - 1639225, - 1689200, - 1740180, - 1899250, - 1792175, - 1845195, -] +from bot.models import IgnoredChannel, LevellingRole, Levels +from bot.models.custom_roles import CustomRoles class Levelling(commands.Cog): def __init__(self, bot): self.bot = bot self.ignored_channel = {} + self.required_xp = [0] async def cog_load(self): for guild in self.bot.guilds: @@ -128,6 +26,11 @@ async def cog_load(self): else: self.ignored_channel[guild.id].append(i.channel_id) + # Calculating required_XP for next level and storing in a list, list-index corresponds to the level + for lvl in range(101): + xp = 5 * (lvl**2) + (50 * lvl) + 100 + self.required_xp.append(xp + self.required_xp[-1]) + @commands.Cog.listener() async def on_message(self, message): # Return if message was sent by Bot or sent in DMs @@ -135,33 +38,34 @@ async def on_message(self, message): return # Check if message is sent in ignored channel - if message.channel.id in self.ignored_channel[message.guild.id]: - return + try: + if message.channel.id in self.ignored_channel[message.guild.id]: + return + except KeyError: + pass # Generate random XP to be added xp = random.randint(5, 25) # Add the XP and update the DB - data = await Levels.update(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) - await self._check_level_up(data, message.author) + data = await Levels.insert_by_guild(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) + self.bot.dispatch("xp_updated", data=data, member=message.author, required_xp=self.required_xp) @app_commands.command() async def rank(self, interaction: core.InteractionType, member: discord.Member = None): """Check the rank of another member or yourself""" if member is None: member = interaction.user - query = """WITH ordered_users AS ( - SELECT - user_id,guild_id,total_xp, - ROW_NUMBER() OVER (ORDER BY levels.total_xp DESC) AS rank - FROM levels - WHERE guild_id = $2) - SELECT rank,total_xp,user_id,guild_id - FROM ordered_users WHERE ordered_users.user_id = $1;""" + SELECT id,user_id,guild_id,total_xp, + ROW_NUMBER() OVER (ORDER BY levelling_users.total_xp DESC) AS rank + FROM levelling_users + WHERE guild_id = $2) + SELECT id,rank,total_xp,user_id,guild_id + FROM ordered_users WHERE ordered_users.user_id = $1;""" data = await Levels.fetchrow(query, member.id, member.guild.id) if data is None: - return await interaction.response.send_message("You are not ranked yet!", ephemeral=True) - for level, j in enumerate(REQUIRED_XP): + return await interaction.response.send_message("User Not ranked yet!", ephemeral=True) + for level, j in enumerate(self.required_xp): if data.total_xp <= j: embed = discord.Embed( title=f"Rank: {data.rank}\nLevel: {level - 1}\nTotal XP:{data.total_xp}", @@ -180,7 +84,6 @@ async def ignore_channel(self, interaction: core.InteractionType, channel: disco @app_commands.command() @app_commands.checks.has_permissions(administrator=True) - @has_permissions(administrator=True) async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): """Remove channel from ignored channel list""" await IgnoredChannel.delete_by_guild(channel.guild.id, channel.id) @@ -190,35 +93,40 @@ async def unignore_channel(self, interaction: core.InteractionType, channel: dis @app_commands.checks.has_permissions(administrator=True) async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): """Give XP to specific user""" - data = await Levels.update(member.guild.id, member.id, xp) - await self._check_level_up(data, member) - await interaction.response.send_message(f"{xp} XP has been added to user {member}") + if xp <= 0: + return await interaction.response.send_message("XP can not be less than 0") + try: + data = await Levels.give_xp(member.guild.id, member.id, xp) + self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) + await interaction.response.send_message(f"{xp} XP has been added to user {member}") + except asyncpg.exceptions.DataError: + return await interaction.response.send_message("Invalid XP provided") @app_commands.command() @app_commands.checks.has_permissions(administrator=True) async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): """Remove XP from user""" - data = await Levels.remove_xp(member.guild.id, member.id, xp) - await self._check_level_up(data, member) - await interaction.response.send_message(f"{xp} XP has been removed from user {member}") - - async def _check_level_up(self, data, member: discord.Member): - """Function to check if user's level has changed and trigger the event to assign the roles""" - # Calculating old and new level + if xp <= 0: + return await interaction.response.send_message("XP can not be less than 0") try: - for level, j in enumerate(REQUIRED_XP): - if data.old_total_xp <= j: - old_level = level - 1 - break - for level, j in enumerate(REQUIRED_XP): - if data.total_xp <= j: - new_level = level - 1 - break - # If the level has changed call level_up to handle roles change - if old_level != new_level: - self.bot.dispatch("level_up", new_level=new_level, member=member) - except AttributeError: - pass + data = await Levels.remove_xp(member.guild.id, member.id, xp) + self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) + await interaction.response.send_message(f"{xp} XP has been removed from user {member}") + except asyncpg.exceptions.DataError: + return await interaction.response.send_message("Invalid XP provided") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def levelling_rewards_add(self, interaction: core.InteractionType, role: discord.Role, level: int): + await CustomRoles.insert_by_guild(role.id, role.guild.id, role.name, str(role.colour)) + await LevellingRole.insert_by_guild(role.guild.id, role.id, level) + return await interaction.response.send_message("Levelling reward role added") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def levelling_rewards_remove(self, interaction: core.InteractionType, role: discord.Role): + await LevellingRole.delete_by_guild(role.guild.id, role.id) + return await interaction.response.send_message("Levelling reward role removed") async def setup(bot: commands.Bot): diff --git a/bot/extensions/levelling/events.py b/bot/extensions/levelling/events.py index acfa410f..80ab43bf 100644 --- a/bot/extensions/levelling/events.py +++ b/bot/extensions/levelling/events.py @@ -1,3 +1,5 @@ +from bisect import bisect + import discord from discord.ext import commands @@ -26,3 +28,16 @@ async def on_level_up(self, new_level, member: discord.Member): 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 diff --git a/bot/extensions/persistent_roles/__init__.py b/bot/extensions/persistent_roles/__init__.py new file mode 100644 index 00000000..570498bc --- /dev/null +++ b/bot/extensions/persistent_roles/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .commands import PersistentRoles +from .events import PersistentEvents + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(PersistentRoles(bot=bot)) + await bot.add_cog(PersistentEvents(bot=bot)) diff --git a/bot/extensions/persistent_roles/commands.py b/bot/extensions/persistent_roles/commands.py new file mode 100644 index 00000000..72769c53 --- /dev/null +++ b/bot/extensions/persistent_roles/commands.py @@ -0,0 +1,10 @@ +from discord.ext import commands + + +class PersistentRoles(commands.Cog): + def __init__(self, bot): + self.bot = bot + + +async def setup(bot: commands.Bot): + await bot.add_cog(PersistentRoles(bot=bot)) diff --git a/bot/cogs/persistent_roles.py b/bot/extensions/persistent_roles/events.py similarity index 72% rename from bot/cogs/persistent_roles.py rename to bot/extensions/persistent_roles/events.py index f7b9a305..3e3bf07c 100644 --- a/bot/cogs/persistent_roles.py +++ b/bot/extensions/persistent_roles/events.py @@ -1,14 +1,18 @@ from discord.ext import commands +from bot import core from bot.models import PersistentRole -class PersistentRoles(commands.Cog): - def __init__(self, bot): +class PersistentEvents(commands.Cog): + """Events for Persisted roles.""" + + 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 @@ -18,7 +22,3 @@ async def on_member_join(self, member): else: for i in range(len(data)): await member.add_roles(member.guild.get_role(data[i].role_id)) - - -async def setup(bot: commands.Bot): - await bot.add_cog(PersistentRoles(bot=bot)) diff --git a/bot/models/__init__.py b/bot/models/__init__.py index ab9715e0..c80a4bfa 100644 --- a/bot/models/__init__.py +++ b/bot/models/__init__.py @@ -1,10 +1,10 @@ from .gconfig import FilterConfig -from .ignored_channel import IgnoredChannel -from .levelling_role import LevellingRole -from .levels import Levels +from .levelling_ignored_channels import IgnoredChannel +from .levelling_roles import LevellingRole +from .levelling_users import Levels from .message import Message from .model import Model -from .persistent_role import PersistentRole +from .persisted_role import PersistentRole from .rep import Rep from .tag import Tag from .user import User diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py new file mode 100644 index 00000000..a764ea66 --- /dev/null +++ b/bot/models/custom_roles.py @@ -0,0 +1,20 @@ +from .model import Model + + +class CustomRoles(Model): + id: int + role_id: int + guild_id: int + name: str + color: str + + @classmethod + async def insert_by_guild(cls, role_id: int, guild_id: int, name: str, color: str): + query = """INSERT INTO custom_roles (role_id, guild_id, name, color) + VALUES($1, $2, $3, $4)""" + await cls.execute(query, role_id, guild_id, name, color) + + @classmethod + async def delete_by_guild(cls, guild_id: int, role_id: int): + query = """DELETE FROM custom_roles WHERE guild_id = $1 and role_id = $2""" + await cls.execute(query, guild_id, role_id) diff --git a/bot/models/ignored_channel.py b/bot/models/levelling_ignored_channels.py similarity index 61% rename from bot/models/ignored_channel.py rename to bot/models/levelling_ignored_channels.py index 85f4e402..931ac7e4 100644 --- a/bot/models/ignored_channel.py +++ b/bot/models/levelling_ignored_channels.py @@ -2,22 +2,23 @@ class IgnoredChannel(Model): + id: int guild_id: int channel_id: int @classmethod async def list_by_guild(cls, guild_id: int): - query = """SELECT * FROM ignored_channels WHERE guild_id = $1""" + query = """SELECT * FROM levelling_ignored_channels WHERE guild_id = $1""" return await cls.fetch(query, guild_id) @classmethod async def insert_by_guild(cls, guild_id: int, channel_id: int): - query = """INSERT INTO ignored_channels (guild_id,channel_id) VALUES($1, $2) + query = """INSERT INTO levelling_ignored_channels (guild_id,channel_id) VALUES($1, $2) ON CONFLICT (guild_id,channel_id) DO NOTHING - RETURNING guild_id, channel_id""" + RETURNING id, guild_id, channel_id""" return await cls.fetch(query, guild_id, channel_id) @classmethod async def delete_by_guild(cls, guild_id: int, channel_id: int): - query = """DELETE FROM ignored_channel WHERE guild_id= $1 and channel_id = $2""" + query = """DELETE FROM levelling_ignored_channels WHERE guild_id= $1 and channel_id = $2""" return await cls.fetch(query, guild_id, channel_id) diff --git a/bot/models/levelling_role.py b/bot/models/levelling_role.py deleted file mode 100644 index c9ee430b..00000000 --- a/bot/models/levelling_role.py +++ /dev/null @@ -1,12 +0,0 @@ -from .model import Model - - -class LevellingRole(Model): - guild_id: int - role_id: int - level: int - - @classmethod - async def list_by_guild(cls, guild_id: int): - query = """SELECT guild_id, role_id, level FROM levelling_roles WHERE guild_id=$1""" - return await cls.fetch(query, guild_id) diff --git a/bot/models/levelling_roles.py b/bot/models/levelling_roles.py new file mode 100644 index 00000000..83340353 --- /dev/null +++ b/bot/models/levelling_roles.py @@ -0,0 +1,26 @@ +from .model import Model + + +class LevellingRole(Model): + id: int + guild_id: int + role_id: int + level: int + + @classmethod + async def list_by_guild(cls, guild_id: int): + query = """SELECT * FROM levelling_roles WHERE guild_id=$1""" + return await cls.fetch(query, guild_id) + + @classmethod + async def insert_by_guild(cls, guild_id: int, role_id: int, level: int): + query = """INSERT INTO levelling_roles (guild_id, role_id, level) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, role_id) + DO NOTHING""" + await cls.execute(query, guild_id, role_id, level) + + @classmethod + async def delete_by_guild(cls, guild_id: int, role_id: int): + query = """DELETE FROM levelling_roles WHERE guild_id = $1 and role_id = $2""" + await cls.execute(query, guild_id, role_id) diff --git a/bot/models/levelling_users.py b/bot/models/levelling_users.py new file mode 100644 index 00000000..fa66f1a6 --- /dev/null +++ b/bot/models/levelling_users.py @@ -0,0 +1,48 @@ +from typing import Optional + +from .model import Model + + +class Levels(Model): + id: int + rank: Optional[int] + guild_id: int + user_id: int + old_total_xp: Optional[int] + total_xp: int + + @classmethod + async def insert_by_guild(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE + SET total_xp = levelling_users.total_xp + $3, last_msg = CURRENT_TIMESTAMP + WHERE levelling_users.guild_id = $1 + AND levelling_users.user_id = $2 + AND EXTRACT(EPOCH FROM (NOW() - levelling_users.last_msg)) > 60 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; +""" + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def give_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levelling_users.total_xp + $3 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def remove_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levelling_users.total_xp - $3 + RETURNING total_xp + $3 AS old_total_xp, total_xp, guild_id, user_id, id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) diff --git a/bot/models/levels.py b/bot/models/levels.py deleted file mode 100644 index 9ca9e23e..00000000 --- a/bot/models/levels.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import Optional - -from .model import Model - - -class Levels(Model): - rank: Optional[int] - guild_id: int - user_id: int - old_total_xp: Optional[int] - total_xp: int - - @classmethod - async def update(cls, guild_id: int, user_id: int, total_xp: int): - query = """INSERT INTO levels (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE - SET total_xp = levels.total_xp + $3, last_msg = CURRENT_TIMESTAMP - WHERE - levels.guild_id = $1 - AND levels.user_id = $2 - AND EXTRACT(EPOCH FROM (NOW() - levels.last_msg)) > 60 - RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id; -""" - - return await cls.fetchrow(query, guild_id, user_id, total_xp) - - @classmethod - async def give_xp(cls, guild_id: int, user_id: int, total_xp: int): - query = """ INSERT INTO levels (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE SET total_xp = levels.total_xp + $3 - RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id; - """ - - return await cls.fetchrow(query, guild_id, user_id, total_xp) - - @classmethod - async def remove_xp(cls, guild_id: int, user_id: int, total_xp: int): - query = """INSERT INTO levels (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE SET total_xp = levels.total_xp - $3 - RETURNING total_xp + $3 AS old_total_xp, total_xp, guild_id, user_id; - """ - - return await cls.fetchrow(query, guild_id, user_id, total_xp) diff --git a/bot/models/migrations/001_down__delete_db.sql b/bot/models/migrations/001_down__delete_db.sql index 4eb41fa7..dade8dc5 100644 --- a/bot/models/migrations/001_down__delete_db.sql +++ b/bot/models/migrations/001_down__delete_db.sql @@ -1 +1 @@ -DROP TABLE gconfigs, messages, reps, tags, users, polls, levels, roles, persistent_roles, ignored_channels, levelling_roles; +DROP TABLE gconfigs, messages, reps, tags, users, polls; diff --git a/bot/models/migrations/001_up__create_db.sql b/bot/models/migrations/001_up__create_db.sql index 92f21ced..cc268740 100644 --- a/bot/models/migrations/001_up__create_db.sql +++ b/bot/models/migrations/001_up__create_db.sql @@ -55,50 +55,4 @@ CREATE TABLE IF NOT EXISTS polls created_at TIMESTAMP, channel_id BIGINT, message_id BIGINT PRIMARY KEY -); - -CREATE TABLE IF NOT EXISTS levels -( - guild_id BIGINT, - user_id BIGINT, - total_xp INT DEFAULT 0, - last_msg TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - CONSTRAINT PK_user PRIMARY KEY(guild_id,user_id) -); - -CREATE TABLE IF NOT EXISTS roles -( - role_id BIGINT PRIMARY KEY, - guild_id BIGINT, - name VARCHAR, - color VARCHAR -); - -CREATE TABLE IF NOT EXISTS persistent_roles -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - user_id BIGINT, - role_id BIGINT, - CONSTRAINT fk_roles FOREIGN KEY(role_id) REFERENCES roles(role_id), - CONSTRAINT unique_guild_user_role UNIQUE (guild_id, user_id, role_id) -); - - -CREATE TABLE IF NOT EXISTS ignored_channels -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - channel_id BIGINT, - CONSTRAINT unique_ignored_channel UNIQUE (guild_id, channel_id) - -); - -CREATE TABLE IF NOT EXISTS levelling_roles -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - role_id BIGINT, - level INT, - CONSTRAINT fk_levelling_roles FOREIGN KEY(role_id) REFERENCES roles(role_id) -); +) diff --git a/bot/models/migrations/002_down__levelling_db.sql b/bot/models/migrations/002_down__levelling_db.sql new file mode 100644 index 00000000..5b043205 --- /dev/null +++ b/bot/models/migrations/002_down__levelling_db.sql @@ -0,0 +1 @@ +DROP TABLE levels, roles, persistent_roles, ignored_channels, levelling_roles; diff --git a/bot/models/migrations/002_up__levelling_db.sql b/bot/models/migrations/002_up__levelling_db.sql new file mode 100644 index 00000000..5a59b781 --- /dev/null +++ b/bot/models/migrations/002_up__levelling_db.sql @@ -0,0 +1,48 @@ +CREATE TABLE IF NOT EXISTS levelling_users +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + user_id BIGINT, + total_xp INT DEFAULT 0, + last_msg TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT levelling_roles_guild_id_and_user_id_key UNIQUE (guild_id, user_id) +); + +CREATE TABLE IF NOT EXISTS custom_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + role_id BIGINT UNIQUE, + name VARCHAR, + color VARCHAR +); + +CREATE TABLE IF NOT EXISTS persisted_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + user_id BIGINT, + role_id BIGINT, + CONSTRAINT persisted_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), + CONSTRAINT persisted_roles_guild_id_and_user_id_and_role_id_key UNIQUE (guild_id, user_id, role_id) +); + + +CREATE TABLE IF NOT EXISTS levelling_ignored_channels +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + channel_id BIGINT, + CONSTRAINT levelling_ignored_channels_guild_id_and_channel_id_key UNIQUE (guild_id, channel_id) + +); + +CREATE TABLE IF NOT EXISTS levelling_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + role_id BIGINT, + level INT, + CONSTRAINT levelling_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), + CONSTRAINT levelling_roles_guild_id_and_role_id_key UNIQUE (guild_id, role_id) +); diff --git a/bot/models/persistent_role.py b/bot/models/persisted_role.py similarity index 67% rename from bot/models/persistent_role.py rename to bot/models/persisted_role.py index 599666fc..3c28b0d6 100644 --- a/bot/models/persistent_role.py +++ b/bot/models/persisted_role.py @@ -4,18 +4,19 @@ class PersistentRole(Model): + id: int guild_id: int user_id: int role_id: Optional[int] @classmethod async def list_by_guild(cls, guild_id: int, user_id: int): - query = """SELECT * FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" + query = """SELECT * FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" return await cls.fetch(query, guild_id, user_id) @classmethod async def insert_by_guild(cls, guild_id: int, user_id: int, role_id: int): - query = """INSERT INTO persistent_roles (guild_id, user_id, role_id) + query = """INSERT INTO persisted_roles (guild_id, user_id, role_id) VALUES ($1, $2, $3) ON CONFLICT (guild_id, user_id, role_id) DO NOTHING""" @@ -24,5 +25,5 @@ async def insert_by_guild(cls, guild_id: int, user_id: int, role_id: int): @classmethod async def delete_by_guild(cls, guild_id: int, user_id: int): - query = """DELETE FROM persistent_roles WHERE guild_id = $1 and user_id = $2""" - return await cls.fetch(query, guild_id, user_id) + query = """DELETE FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" + return await cls.execute(query, guild_id, user_id) diff --git a/cli.py b/cli.py index 8f41df78..73ed3348 100644 --- a/cli.py +++ b/cli.py @@ -140,7 +140,7 @@ async def main(ctx): "bot.cogs.roles", "bot.cogs.poll", "bot.extensions.levelling", - "bot.cogs.persistent_roles", + "bot.extensions.persistent_roles", ) intents = discord.Intents.all() From 4bb9a38b55da9d4f26572a2114c84507dd9d13a8 Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Sun, 15 Oct 2023 11:54:24 +0530 Subject: [PATCH 05/10] added levelling rewards list commnad --- bot/extensions/levelling/commands.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py index b7d5046a..f994227e 100644 --- a/bot/extensions/levelling/commands.py +++ b/bot/extensions/levelling/commands.py @@ -128,6 +128,18 @@ async def levelling_rewards_remove(self, interaction: core.InteractionType, role await LevellingRole.delete_by_guild(role.guild.id, role.id) return await interaction.response.send_message("Levelling reward role removed") + @app_commands.command() + async def levelling_rewards_list(self, interaction: core.InteractionType): + data = await LevellingRole.list_by_guild(interaction.guild.id) + res = "| {:<10} | {:<10} | {:<5} |".format("Guild ID", "Role ID", "Level") + res += "\n|" + "-" * 12 + "|" + "-" * 12 + "|" + "-" * 7 + "|" + + # Print data + for record in data: + res += "\n| {:<10} | {:<10} | {:<5} |".format(record["guild_id"], record["role_id"], record["level"]) + + await interaction.response.send_message(f"{res}") + async def setup(bot: commands.Bot): await bot.add_cog(Levelling(bot=bot)) From e70b64b76e703c82d820717a1b910891c16dc6d6 Mon Sep 17 00:00:00 2001 From: Sarthak singh Date: Tue, 17 Oct 2023 20:23:30 +0530 Subject: [PATCH 06/10] Added XP multiplier --- bot/extensions/levelling/commands.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py index f994227e..ceafbe15 100644 --- a/bot/extensions/levelling/commands.py +++ b/bot/extensions/levelling/commands.py @@ -16,6 +16,7 @@ def __init__(self, bot): self.bot = bot self.ignored_channel = {} self.required_xp = [0] + self.xp_boost = 1 async def cog_load(self): for guild in self.bot.guilds: @@ -45,7 +46,7 @@ async def on_message(self, message): pass # Generate random XP to be added - xp = random.randint(5, 25) + xp = random.randint(5, 25) * self.xp_boost # Add the XP and update the DB data = await Levels.insert_by_guild(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) self.bot.dispatch("xp_updated", data=data, member=message.author, required_xp=self.required_xp) @@ -140,6 +141,16 @@ async def levelling_rewards_list(self, interaction: core.InteractionType): await interaction.response.send_message(f"{res}") + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def xp_multiplier(self, interaction: core.InteractionType, multiplier: int): + """Increase XP gain per message""" + if multiplier <= 0 or multiplier > 5: + return await interaction.response.send_message("Invalid multiplier value.(Max. 5)") + + self.xp_boost = multiplier + return await interaction.response.send_message(f"XP multiplied by {multiplier}x.") + async def setup(bot: commands.Bot): await bot.add_cog(Levelling(bot=bot)) From 85f25f45fb3ebb5d90a6afa28cdbdb05935ce845 Mon Sep 17 00:00:00 2001 From: SylteA Date: Sat, 21 Oct 2023 22:52:33 +0200 Subject: [PATCH 07/10] Update migrations --- .../migrations/002_down__levelling_db.sql | 1 - .../migrations/002_up__levelling_db.sql | 48 ------------------- .../003_down__snowflake_function.sql | 3 ++ .../migrations/003_up__snowflake_function.sql | 32 +++++++++++++ bot/models/migrations/004_down__levelling.sql | 1 + bot/models/migrations/004_up__levelling.sql | 48 +++++++++++++++++++ cli.py | 6 +-- 7 files changed, 87 insertions(+), 52 deletions(-) delete mode 100644 bot/models/migrations/002_down__levelling_db.sql delete mode 100644 bot/models/migrations/002_up__levelling_db.sql create mode 100644 bot/models/migrations/003_down__snowflake_function.sql create mode 100644 bot/models/migrations/003_up__snowflake_function.sql create mode 100644 bot/models/migrations/004_down__levelling.sql create mode 100644 bot/models/migrations/004_up__levelling.sql diff --git a/bot/models/migrations/002_down__levelling_db.sql b/bot/models/migrations/002_down__levelling_db.sql deleted file mode 100644 index 5b043205..00000000 --- a/bot/models/migrations/002_down__levelling_db.sql +++ /dev/null @@ -1 +0,0 @@ -DROP TABLE levels, roles, persistent_roles, ignored_channels, levelling_roles; diff --git a/bot/models/migrations/002_up__levelling_db.sql b/bot/models/migrations/002_up__levelling_db.sql deleted file mode 100644 index 5a59b781..00000000 --- a/bot/models/migrations/002_up__levelling_db.sql +++ /dev/null @@ -1,48 +0,0 @@ -CREATE TABLE IF NOT EXISTS levelling_users -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - user_id BIGINT, - total_xp INT DEFAULT 0, - last_msg TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - CONSTRAINT levelling_roles_guild_id_and_user_id_key UNIQUE (guild_id, user_id) -); - -CREATE TABLE IF NOT EXISTS custom_roles -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - role_id BIGINT UNIQUE, - name VARCHAR, - color VARCHAR -); - -CREATE TABLE IF NOT EXISTS persisted_roles -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - user_id BIGINT, - role_id BIGINT, - CONSTRAINT persisted_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), - CONSTRAINT persisted_roles_guild_id_and_user_id_and_role_id_key UNIQUE (guild_id, user_id, role_id) -); - - -CREATE TABLE IF NOT EXISTS levelling_ignored_channels -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - channel_id BIGINT, - CONSTRAINT levelling_ignored_channels_guild_id_and_channel_id_key UNIQUE (guild_id, channel_id) - -); - -CREATE TABLE IF NOT EXISTS levelling_roles -( - id SERIAL PRIMARY KEY, - guild_id BIGINT, - role_id BIGINT, - level INT, - CONSTRAINT levelling_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), - CONSTRAINT levelling_roles_guild_id_and_role_id_key UNIQUE (guild_id, role_id) -); diff --git a/bot/models/migrations/003_down__snowflake_function.sql b/bot/models/migrations/003_down__snowflake_function.sql new file mode 100644 index 00000000..66ac98b5 --- /dev/null +++ b/bot/models/migrations/003_down__snowflake_function.sql @@ -0,0 +1,3 @@ +DROP FUNCTION IF EXISTS snowflake_to_timestamp(flake bigint); +DROP FUNCTION IF EXISTS create_snowflake(shard_id integer); +DROP SEQUENCE IF EXISTS global_snowflake_id_seq; diff --git a/bot/models/migrations/003_up__snowflake_function.sql b/bot/models/migrations/003_up__snowflake_function.sql new file mode 100644 index 00000000..0af7a145 --- /dev/null +++ b/bot/models/migrations/003_up__snowflake_function.sql @@ -0,0 +1,32 @@ +CREATE SEQUENCE IF NOT EXISTS global_snowflake_id_seq; + +CREATE OR REPLACE FUNCTION create_snowflake(shard_id INT DEFAULT 1) + RETURNS bigint + LANGUAGE 'plpgsql' +AS $$ +DECLARE + our_epoch bigint := 1609459200; + seq_id bigint; + now_millis bigint; + result bigint:= 0; +BEGIN + SELECT nextval('global_snowflake_id_seq') % 1024 INTO seq_id; + + SELECT FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000) INTO now_millis; + result := (now_millis - our_epoch) << 22; + result := result | (shard_id << 9); + result := result | (seq_id); + return result; +END; +$$; + +CREATE OR REPLACE FUNCTION snowflake_to_timestamp(flake BIGINT) + RETURNS TIMESTAMP + LANGUAGE 'plpgsql' +AS $BODY$ +DECLARE + our_epoch BIGINT := 1621470600; +BEGIN + RETURN to_timestamp(((flake >> 22) + our_epoch) / 1000); +END; +$BODY$; diff --git a/bot/models/migrations/004_down__levelling.sql b/bot/models/migrations/004_down__levelling.sql new file mode 100644 index 00000000..da44ff4e --- /dev/null +++ b/bot/models/migrations/004_down__levelling.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS levelling_users, levelling_roles, levelling_ignored_channels, custom_roles, persisted_roles; diff --git a/bot/models/migrations/004_up__levelling.sql b/bot/models/migrations/004_up__levelling.sql new file mode 100644 index 00000000..ea2b2414 --- /dev/null +++ b/bot/models/migrations/004_up__levelling.sql @@ -0,0 +1,48 @@ +CREATE TABLE IF NOT EXISTS levelling_users +( + id BIGINT PRIMARY KEY, + guild_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + total_xp INT NOT NULL DEFAULT 0, + last_msg BIGINT NOT NULL DEFAULT create_snowflake(), + CONSTRAINT levelling_users_guild_id_user_id_key UNIQUE (guild_id, user_id) +); + + +CREATE TABLE IF NOT EXISTS custom_roles +( + id BIGINT PRIMARY KEY, + guild_id BIGINT NOT NULL, + role_id BIGINT NOT NULL, + name VARCHAR NOT NULL, + color VARCHAR NOT NULL, + CONSTRAINT custom_roles_role_id UNIQUE (role_id), + CONSTRAINT custom_roles_guild_id_role_id UNIQUE (guild_id, role_id) +); + +CREATE TABLE IF NOT EXISTS levelling_roles +( + id SERIAL PRIMARY KEY, + required_xp INTEGER NOT NULL, + guild_id BIGINT NOT NULL, + role_id BIGINT NOT NULL REFERENCES custom_roles(role_id), + CONSTRAINT levelling_roles_guild_id_role_id_key UNIQUE (guild_id, role_id) +); + +CREATE TABLE IF NOT EXISTS persisted_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + role_id BIGINT NOT NULL REFERENCES custom_roles(role_id), + CONSTRAINT persisted_roles_guild_id_role_id_user_id__key UNIQUE (guild_id, role_id, user_id) +); + + +CREATE TABLE IF NOT EXISTS levelling_ignored_channels +( + id SERIAL PRIMARY KEY, + 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) +); diff --git a/cli.py b/cli.py index 2f7a70e3..f90ae121 100644 --- a/cli.py +++ b/cli.py @@ -135,12 +135,12 @@ async def main(ctx): "bot.extensions.suggestions", "bot.extensions.github", "bot.extensions.tags", + "bot.extensions.levelling", + "bot.extensions.persistent_roles", "bot.cogs._help", "bot.cogs.clashofcode", "bot.cogs.roles", "bot.cogs.poll", - "bot.extensions.levelling", - "bot.extensions.persistent_roles", ) intents = discord.Intents.all() @@ -160,7 +160,7 @@ async def run_migration(file: str = "000_up__migrations.sql") -> Migration: mig = Migration.from_match(match) await Model.execute(query) - click.echo(f"{file} was executed.") + log.info(f"Migration {mig.name} ({file}) was executed.") await mig.post() return mig From 72c74edb58d0b959f7c815bfc392c05bde4a3b81 Mon Sep 17 00:00:00 2001 From: SylteA Date: Sat, 21 Oct 2023 22:58:08 +0200 Subject: [PATCH 08/10] Temporarily update models --- bot/models/custom_roles.py | 13 +------- bot/models/levelling_ignored_channels.py | 17 ---------- bot/models/levelling_roles.py | 20 +---------- bot/models/levelling_users.py | 42 +----------------------- bot/models/persisted_role.py | 25 ++------------ 5 files changed, 5 insertions(+), 112 deletions(-) diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py index a764ea66..4c650291 100644 --- a/bot/models/custom_roles.py +++ b/bot/models/custom_roles.py @@ -3,18 +3,7 @@ class CustomRoles(Model): id: int - role_id: int guild_id: int + role_id: int name: str color: str - - @classmethod - async def insert_by_guild(cls, role_id: int, guild_id: int, name: str, color: str): - query = """INSERT INTO custom_roles (role_id, guild_id, name, color) - VALUES($1, $2, $3, $4)""" - await cls.execute(query, role_id, guild_id, name, color) - - @classmethod - async def delete_by_guild(cls, guild_id: int, role_id: int): - query = """DELETE FROM custom_roles WHERE guild_id = $1 and role_id = $2""" - await cls.execute(query, guild_id, role_id) diff --git a/bot/models/levelling_ignored_channels.py b/bot/models/levelling_ignored_channels.py index 931ac7e4..5eee61b9 100644 --- a/bot/models/levelling_ignored_channels.py +++ b/bot/models/levelling_ignored_channels.py @@ -5,20 +5,3 @@ class IgnoredChannel(Model): id: int guild_id: int channel_id: int - - @classmethod - async def list_by_guild(cls, guild_id: int): - query = """SELECT * FROM levelling_ignored_channels WHERE guild_id = $1""" - return await cls.fetch(query, guild_id) - - @classmethod - async def insert_by_guild(cls, guild_id: int, channel_id: int): - query = """INSERT INTO levelling_ignored_channels (guild_id,channel_id) VALUES($1, $2) - ON CONFLICT (guild_id,channel_id) DO NOTHING - RETURNING id, guild_id, channel_id""" - return await cls.fetch(query, guild_id, channel_id) - - @classmethod - async def delete_by_guild(cls, guild_id: int, channel_id: int): - query = """DELETE FROM levelling_ignored_channels WHERE guild_id= $1 and channel_id = $2""" - return await cls.fetch(query, guild_id, channel_id) diff --git a/bot/models/levelling_roles.py b/bot/models/levelling_roles.py index 83340353..c11216d9 100644 --- a/bot/models/levelling_roles.py +++ b/bot/models/levelling_roles.py @@ -3,24 +3,6 @@ class LevellingRole(Model): id: int + required_xp: int guild_id: int role_id: int - level: int - - @classmethod - async def list_by_guild(cls, guild_id: int): - query = """SELECT * FROM levelling_roles WHERE guild_id=$1""" - return await cls.fetch(query, guild_id) - - @classmethod - async def insert_by_guild(cls, guild_id: int, role_id: int, level: int): - query = """INSERT INTO levelling_roles (guild_id, role_id, level) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, role_id) - DO NOTHING""" - await cls.execute(query, guild_id, role_id, level) - - @classmethod - async def delete_by_guild(cls, guild_id: int, role_id: int): - query = """DELETE FROM levelling_roles WHERE guild_id = $1 and role_id = $2""" - await cls.execute(query, guild_id, role_id) diff --git a/bot/models/levelling_users.py b/bot/models/levelling_users.py index fa66f1a6..5073a9c0 100644 --- a/bot/models/levelling_users.py +++ b/bot/models/levelling_users.py @@ -1,48 +1,8 @@ -from typing import Optional - from .model import Model -class Levels(Model): +class LevellingUser(Model): id: int - rank: Optional[int] guild_id: int user_id: int - old_total_xp: Optional[int] total_xp: int - - @classmethod - async def insert_by_guild(cls, guild_id: int, user_id: int, total_xp: int): - query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE - SET total_xp = levelling_users.total_xp + $3, last_msg = CURRENT_TIMESTAMP - WHERE levelling_users.guild_id = $1 - AND levelling_users.user_id = $2 - AND EXTRACT(EPOCH FROM (NOW() - levelling_users.last_msg)) > 60 - RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; -""" - return await cls.fetchrow(query, guild_id, user_id, total_xp) - - @classmethod - async def give_xp(cls, guild_id: int, user_id: int, total_xp: int): - query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE SET total_xp = levelling_users.total_xp + $3 - RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; - """ - - return await cls.fetchrow(query, guild_id, user_id, total_xp) - - @classmethod - async def remove_xp(cls, guild_id: int, user_id: int, total_xp: int): - query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id) - DO UPDATE SET total_xp = levelling_users.total_xp - $3 - RETURNING total_xp + $3 AS old_total_xp, total_xp, guild_id, user_id, id; - """ - - return await cls.fetchrow(query, guild_id, user_id, total_xp) diff --git a/bot/models/persisted_role.py b/bot/models/persisted_role.py index 3c28b0d6..d16c49cd 100644 --- a/bot/models/persisted_role.py +++ b/bot/models/persisted_role.py @@ -1,29 +1,8 @@ -from typing import Optional - from .model import Model -class PersistentRole(Model): +class PersistedRole(Model): id: int guild_id: int user_id: int - role_id: Optional[int] - - @classmethod - async def list_by_guild(cls, guild_id: int, user_id: int): - query = """SELECT * FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" - return await cls.fetch(query, guild_id, user_id) - - @classmethod - async def insert_by_guild(cls, guild_id: int, user_id: int, role_id: int): - query = """INSERT INTO persisted_roles (guild_id, user_id, role_id) - VALUES ($1, $2, $3) - ON CONFLICT (guild_id, user_id, role_id) - DO NOTHING""" - - await cls.execute(query, guild_id, user_id, role_id) - - @classmethod - async def delete_by_guild(cls, guild_id: int, user_id: int): - query = """DELETE FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" - return await cls.execute(query, guild_id, user_id) + role_id: int From 75f8e6b8b0a005dfbe353bfb2ae18fd22cce2c58 Mon Sep 17 00:00:00 2001 From: SylteA Date: Sat, 21 Oct 2023 23:13:58 +0200 Subject: [PATCH 09/10] Add defaults to all new ids in tables --- bot/models/migrations/004_up__levelling.sql | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/bot/models/migrations/004_up__levelling.sql b/bot/models/migrations/004_up__levelling.sql index ea2b2414..5f8f70bc 100644 --- a/bot/models/migrations/004_up__levelling.sql +++ b/bot/models/migrations/004_up__levelling.sql @@ -1,6 +1,6 @@ CREATE TABLE IF NOT EXISTS levelling_users ( - id BIGINT PRIMARY KEY, + id BIGINT PRIMARY KEY DEFAULT create_snowflake(), guild_id BIGINT NOT NULL, user_id BIGINT NOT NULL, total_xp INT NOT NULL DEFAULT 0, @@ -11,7 +11,7 @@ CREATE TABLE IF NOT EXISTS levelling_users CREATE TABLE IF NOT EXISTS custom_roles ( - id BIGINT PRIMARY KEY, + id BIGINT PRIMARY KEY DEFAULT create_snowflake(), guild_id BIGINT NOT NULL, role_id BIGINT NOT NULL, name VARCHAR NOT NULL, @@ -22,7 +22,7 @@ CREATE TABLE IF NOT EXISTS custom_roles CREATE TABLE IF NOT EXISTS levelling_roles ( - id SERIAL PRIMARY KEY, + id SERIAL 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, + id SERIAL 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, + id SERIAL 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) From 9aae19e537467ece151d4b4b2a5f3b94e445dfd2 Mon Sep 17 00:00:00 2001 From: SylteA Date: Sun, 22 Oct 2023 14:27:14 +0200 Subject: [PATCH 10/10] Updated commands and events for levelling. --- bot/extensions/levelling/__init__.py | 4 +- bot/extensions/levelling/commands.py | 363 ++++++++++++------ bot/extensions/levelling/events.py | 81 ++-- bot/extensions/levelling/utils.py | 15 + bot/extensions/persistent_roles/events.py | 90 ++++- bot/models/__init__.py | 8 +- bot/models/custom_roles.py | 17 +- .../migrations/003_up__snowflake_function.sql | 2 +- bot/models/migrations/004_up__levelling.sql | 6 +- bot/models/model.py | 17 +- 10 files changed, 439 insertions(+), 164 deletions(-) create mode 100644 bot/extensions/levelling/utils.py diff --git a/bot/extensions/levelling/__init__.py b/bot/extensions/levelling/__init__.py index 0e10853a..3ba1b473 100644 --- a/bot/extensions/levelling/__init__.py +++ b/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)) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py index ceafbe15..4c4ff0af 100644 --- a/bot/extensions/levelling/commands.py +++ b/bot/extensions/levelling/commands.py @@ -1,156 +1,301 @@ import datetime import random -import asyncpg.exceptions import discord from discord import app_commands from discord.ext import commands from bot import core -from bot.models import IgnoredChannel, LevellingRole, Levels -from bot.models.custom_roles import CustomRoles +from bot.extensions.levelling import utils +from bot.models import IgnoredChannel, LevellingRole, LevellingUser +from bot.models.custom_roles import CustomRole class Levelling(commands.Cog): - def __init__(self, bot): + admin_commands = app_commands.Group( + name="levelling", + description="Levelling commands for staff", + default_permissions=discord.Permissions(administrator=True), + ) + + ignored_channels = app_commands.Group( + parent=admin_commands, + name="ignored_channels", + description="Commands to manage XP ignored channels.", + default_permissions=discord.Permissions(administrator=True), + ) + + xp = app_commands.Group( + parent=admin_commands, + name="experience", + description="Manually update the XP of a player.", + default_permissions=discord.Permissions(administrator=True), + ) + + rewards = app_commands.Group( + parent=admin_commands, + name="rewards", + description="Manage the roles given at a certain amount of xp.", + default_permissions=discord.Permissions(administrator=True), + ) + + def __init__(self, bot: core.DiscordBot): self.bot = bot - self.ignored_channel = {} + + self.ignored_channels: dict[int, list[int]] = {} self.required_xp = [0] self.xp_boost = 1 async def cog_load(self): - for guild in self.bot.guilds: - data = await IgnoredChannel.list_by_guild(guild_id=guild.id) - for i in data: - if guild.id not in self.ignored_channel: - self.ignored_channel[guild.id] = [i.channel_id] - else: - self.ignored_channel[guild.id].append(i.channel_id) - - # Calculating required_XP for next level and storing in a list, list-index corresponds to the level + query = """ + SELECT * + FROM levelling_ignored_channels + WHERE guild_id = ANY($1) + """ + guild_ids = [guild.id for guild in self.bot.guilds] + self.ignored_channels = {guild.id: [] for guild in self.bot.guilds} + + channels = await IgnoredChannel.fetch(query, guild_ids) + + for channel in channels: + self.ignored_channels[channel.guild_id].append(channel.channel_id) + for lvl in range(101): xp = 5 * (lvl**2) + (50 * lvl) + 100 self.required_xp.append(xp + self.required_xp[-1]) @commands.Cog.listener() async def on_message(self, message): - # Return if message was sent by Bot or sent in DMs if message.author.bot or message.guild is None: return - # Check if message is sent in ignored channel - try: - if message.channel.id in self.ignored_channel[message.guild.id]: + if message.guild.id in self.ignored_channels: + if message.channel.id in self.ignored_channels[message.guild.id]: return - except KeyError: - pass - # Generate random XP to be added + query = """ + INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET + total_xp = levelling_users.total_xp + $3, + last_msg = create_snowflake() + WHERE levelling_users.guild_id = $1 + AND levelling_users.user_id = $2 + AND snowflake_to_timestamp(levelling_users.last_msg) < NOW() - INTERVAL '1 min' + RETURNING *; + """ + + # TODO: Allow each guild to set custom xp range and boost. xp = random.randint(5, 25) * self.xp_boost - # Add the XP and update the DB - data = await Levels.insert_by_guild(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) - self.bot.dispatch("xp_updated", data=data, member=message.author, required_xp=self.required_xp) + after = await LevellingUser.fetchrow(query, message.guild.id, message.author.id, xp) + + if after is None: + return # Last message was less than a minute ago. + + before = after.copy(update={"total_xp": after.total_xp - xp}) + + self.bot.dispatch("xp_update", before=before, after=after) @app_commands.command() async def rank(self, interaction: core.InteractionType, member: discord.Member = None): """Check the rank of another member or yourself""" - if member is None: - member = interaction.user - query = """WITH ordered_users AS ( - SELECT id,user_id,guild_id,total_xp, - ROW_NUMBER() OVER (ORDER BY levelling_users.total_xp DESC) AS rank - FROM levelling_users - WHERE guild_id = $2) - SELECT id,rank,total_xp,user_id,guild_id - FROM ordered_users WHERE ordered_users.user_id = $1;""" - data = await Levels.fetchrow(query, member.id, member.guild.id) - if data is None: + member = member or interaction.user + + query = """ + WITH "user" AS ( + SELECT total_xp + FROM levelling_users + WHERE guild_id = $1 + AND user_id = $2 + ) + SELECT (SELECT total_xp FROM "user"), COUNT(*) + FROM levelling_users + WHERE guild_id = $1 + AND total_xp > (SELECT total_xp FROM "user"); + """ + + record = await LevellingUser.pool.fetchrow(query, interaction.guild.id, member.id) + + if record is None: return await interaction.response.send_message("User Not ranked yet!", ephemeral=True) - for level, j in enumerate(self.required_xp): - if data.total_xp <= j: - embed = discord.Embed( - title=f"Rank: {data.rank}\nLevel: {level - 1}\nTotal XP:{data.total_xp}", - timestamp=datetime.datetime.utcnow(), - colour=discord.Colour.blurple(), - ) - embed.set_thumbnail(url=member.avatar) - return await interaction.response.send_message(embed=embed) - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): - """Add the channel to the ignored channel list to not gain XP""" - await IgnoredChannel.insert_by_guild(channel.guild.id, channel.id) - await interaction.response.send_message(f"{channel} has been ignored from gaining XP.") + level = utils.get_level_for_xp(user_xp=record.total_xp) - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): - """Remove channel from ignored channel list""" - await IgnoredChannel.delete_by_guild(channel.guild.id, channel.id) - await interaction.response.send_message(f"{channel} has been removed from ignored channel list") + embed = discord.Embed( + title=f"Rank: {record.count + 1}\nLevel: {level}, Total XP: {record.total_xp}", + timestamp=datetime.datetime.utcnow(), + color=discord.Color.blurple(), + ) + embed.set_thumbnail(url=member.avatar) + return await interaction.response.send_message(embed=embed) - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): - """Give XP to specific user""" - if xp <= 0: - return await interaction.response.send_message("XP can not be less than 0") - try: - data = await Levels.give_xp(member.guild.id, member.id, xp) - self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) - await interaction.response.send_message(f"{xp} XP has been added to user {member}") - except asyncpg.exceptions.DataError: - return await interaction.response.send_message("Invalid XP provided") + @ignored_channels.command() + @app_commands.describe(channel="The channel to ignore.") + async def add(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Add a channel to the list of ignored channels.""" + if channel.guild.id not in self.ignored_channels: + self.ignored_channels[channel.guild.id] = [] - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): - """Remove XP from user""" - if xp <= 0: - return await interaction.response.send_message("XP can not be less than 0") - try: - data = await Levels.remove_xp(member.guild.id, member.id, xp) - self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) - await interaction.response.send_message(f"{xp} XP has been removed from user {member}") - except asyncpg.exceptions.DataError: - return await interaction.response.send_message("Invalid XP provided") + if channel.id in self.ignored_channels[channel.guild.id]: + return await interaction.response.send_message("That channel is already ignored.", ephemeral=True) - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def levelling_rewards_add(self, interaction: core.InteractionType, role: discord.Role, level: int): - await CustomRoles.insert_by_guild(role.id, role.guild.id, role.name, str(role.colour)) - await LevellingRole.insert_by_guild(role.guild.id, role.id, level) - return await interaction.response.send_message("Levelling reward role added") + self.ignored_channels[channel.guild.id].append(channel.id) - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def levelling_rewards_remove(self, interaction: core.InteractionType, role: discord.Role): - await LevellingRole.delete_by_guild(role.guild.id, role.id) - return await interaction.response.send_message("Levelling reward role removed") + query = """ + INSERT INTO levelling_ignored_channels (guild_id, channel_id) + VALUES ($1, $2) + ON CONFLICT (guild_id, channel_id) DO NOTHING + """ + await IgnoredChannel.execute(query, channel.guild.id, channel.id) + return await interaction.response.send_message(f"I'll now ignore messages in {channel.mention}.") - @app_commands.command() - async def levelling_rewards_list(self, interaction: core.InteractionType): - data = await LevellingRole.list_by_guild(interaction.guild.id) - res = "| {:<10} | {:<10} | {:<5} |".format("Guild ID", "Role ID", "Level") - res += "\n|" + "-" * 12 + "|" + "-" * 12 + "|" + "-" * 7 + "|" + @ignored_channels.command() + @app_commands.describe(channel="The channel to remove.") + async def remove(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Remove a channel from the list of ignored channels.""" + if channel.id not in self.ignored_channels.get(channel.guild.id, []): + return await interaction.response.send_message("That channel is not ignored", ephemeral=True) - # Print data - for record in data: - res += "\n| {:<10} | {:<10} | {:<5} |".format(record["guild_id"], record["role_id"], record["level"]) + self.ignored_channels[channel.guild.id].remove(channel.id) - await interaction.response.send_message(f"{res}") + query = """ + DELETE FROM levelling_ignored_channels + WHERE channel_id = $1 + """ + await IgnoredChannel.execute(query, channel.id) + return await interaction.response.send_message(f"No longer ignoring messages in {channel.mention}.") - @app_commands.command() - @app_commands.checks.has_permissions(administrator=True) - async def xp_multiplier(self, interaction: core.InteractionType, multiplier: int): - """Increase XP gain per message""" - if multiplier <= 0 or multiplier > 5: - return await interaction.response.send_message("Invalid multiplier value.(Max. 5)") + @ignored_channels.command(name="list") + @app_commands.describe(ephemeral="If true, only you can see the response.") + async def list_channels(self, interaction: core.InteractionType, ephemeral: bool = True): + """Lists the ignored channels in this guild.""" + ignored_channels = self.ignored_channels.get(interaction.guild.id, []) + + response = f"## Listing `{len(ignored_channels)}` ignored channels\n\n" + + for channel_id in ignored_channels: + channel = interaction.guild.get_channel(channel_id) + + if channel is None: + response += f"- Unknown Channel (`{channel_id}`)\n" + else: + response += f"- {channel.mention}\n" + + return await interaction.response.send_message(response, ephemeral=ephemeral) + + async def update_user_xp(self, guild_id: int, user_id: int, amount: int): + query = """ + INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET + total_xp = GREATEST(levelling_users.total_xp + $3, 0), + last_msg = create_snowflake() + WHERE levelling_users.guild_id = $1 + AND levelling_users.user_id = $2 + RETURNING *; + """ + + after = await LevellingUser.fetchrow(query, guild_id, user_id, amount) + before = after.copy(update={"total_xp": after.total_xp - amount}) + + self.bot.dispatch("xp_update", before=before, after=after) + + @xp.command(name="add") + @app_commands.describe(member="The member member to add xp to.", amount="The amount of xp to give.") + async def add_xp(self, interaction: core.InteractionType, member: discord.Member, amount: int): + """Give XP to the specified user""" + if amount <= 0: + return await interaction.response.send_message("Amount must be a positive integer.", ephemeral=True) + + if member.bot: + return await interaction.response.send_message("Cannot add experience to a bot.", ephemeral=True) + + await self.update_user_xp(guild_id=member.guild.id, user_id=member.id, amount=amount) + return await interaction.response.send_message(f"Added `{amount}` XP to {member.display_name}") + + @xp.command(name="remove") + @app_commands.describe(member="The member member to add xp to.", amount="The amount of xp to give.") + async def remove_xp(self, interaction: core.InteractionType, member: discord.Member, amount: int): + """Remove XP from the specified user.""" + if amount <= 0: + return await interaction.response.send_message("Amount must be a positive integer.", ephemeral=True) + + if member.bot: + return await interaction.response.send_message("Cannot remove experience from a bot.", ephemeral=True) + + await self.update_user_xp(guild_id=member.guild.id, user_id=member.id, amount=-amount) + return await interaction.response.send_message(f"Removed `{amount}` XP from {member.display_name}") + + @rewards.command(name="add") + @app_commands.describe(role="The role to reward.", level="The level to reward it at.") + async def add_role(self, interaction: core.InteractionType, role: discord.Role, level: int): + """Add a levelling reward.""" + query = """ + INSERT INTO levelling_roles (guild_id, role_id, required_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, role_id) + DO NOTHING + RETURNING *; + """ + + required_xp = utils.get_xp_for_level(level) + + await CustomRole.ensure_exists( + guild_id=role.guild.id, role_id=role.id, name=role.name, color=str(role.color.value) + ) + + record = await LevellingRole.fetchrow(query, role.guild.id, role.id, required_xp) + + if record is None: + return await interaction.response.send_message( + f"{role.mention} is already a levelling reward.", ephemeral=True + ) + + return await interaction.response.send_message( + f"{role.mention} has been added as a reward for reaching level {level}!", + allowed_mentions=discord.AllowedMentions.none(), + ) + + @rewards.command(name="remove") + async def remove_role(self, interaction: core.InteractionType, role: discord.Role): + """Remove a levelling reward.""" + query = """ + DELETE FROM levelling_roles + WHERE role_id = $1 + """ + + deleted = await LevellingRole.execute(query, role.id) + # count = int(deleted.removeprefix("DELETE ")) + + return await interaction.response.send_message(deleted, ephemeral=True) + + @rewards.command(name="list") + async def list_roles(self, interaction: core.InteractionType): + """Lists the levelling roles in this guild.""" + query = """ + SELECT * + FROM levelling_roles + WHERE guild_id = $1 + """ + + rewards = await LevellingRole.fetch(query, interaction.guild.id) + + # TODO: This command needs to be fixed, the max role name length needs to be dynamic + response = "| {:<10} | {:<5} |".format("Role Name", "Level") + response += "\n|" + "-" * 12 + "|" + "-" * 7 + "|" + + for reward in rewards: + role = interaction.guild.get_role(reward.role_id) + level = utils.get_level_for_xp(reward.required_xp) + + role_name = f"Unknown role (`{reward.role_id}`)" if role is None else role.name + + response += "\n| {:<10} | {:<5} |".format(role_name, level) - self.xp_boost = multiplier - return await interaction.response.send_message(f"XP multiplied by {multiplier}x.") + return await interaction.response.send_message(f"```\n{response}\n```\n") -async def setup(bot: commands.Bot): +async def setup(bot: core.DiscordBot): await bot.add_cog(Levelling(bot=bot)) diff --git a/bot/extensions/levelling/events.py b/bot/extensions/levelling/events.py index 80ab43bf..420f886a 100644 --- a/bot/extensions/levelling/events.py +++ b/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 + ) diff --git a/bot/extensions/levelling/utils.py b/bot/extensions/levelling/utils.py new file mode 100644 index 00000000..d27758c5 --- /dev/null +++ b/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] diff --git a/bot/extensions/persistent_roles/events.py b/bot/extensions/persistent_roles/events.py index 3e3bf07c..ca32d6e3 100644 --- a/bot/extensions/persistent_roles/events.py +++ b/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) diff --git a/bot/models/__init__.py b/bot/models/__init__.py index c80a4bfa..3fb6422e 100644 --- a/bot/models/__init__.py +++ b/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 diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py index 4c650291..73c90fba 100644 --- a/bot/models/custom_roles.py +++ b/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) diff --git a/bot/models/migrations/003_up__snowflake_function.sql b/bot/models/migrations/003_up__snowflake_function.sql index 0af7a145..bb88ac78 100644 --- a/bot/models/migrations/003_up__snowflake_function.sql +++ b/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; diff --git a/bot/models/migrations/004_up__levelling.sql b/bot/models/migrations/004_up__levelling.sql index 5f8f70bc..76d851cf 100644 --- a/bot/models/migrations/004_up__levelling.sql +++ b/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) diff --git a/bot/models/model.py b/bot/models/model.py index b469565e..28a9dce6 100644 --- a/bot/models/model.py +++ b/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