From b4b6db683af62920628e8eafa42cb9d70ef32c44 Mon Sep 17 00:00:00 2001 From: Nebulo9 Date: Tue, 21 Nov 2023 12:06:09 +0100 Subject: [PATCH] Changed data storage to a distant database --- main.py | 124 +++++++++----- modules/cogs/birthday.py | 160 +++++++++--------- modules/cogs/productivity.py | 182 +++++++++++++-------- modules/setup/data.py | 39 ----- modules/setup/db.py | 304 +++++++++++++++++++++++++++++++++++ 5 files changed, 590 insertions(+), 219 deletions(-) delete mode 100644 modules/setup/data.py create mode 100644 modules/setup/db.py diff --git a/main.py b/main.py index d2241b2..a0df7a3 100644 --- a/main.py +++ b/main.py @@ -1,4 +1,5 @@ -import os, discord, json, sys +import os, discord, json, sys, random +import modules.setup.db as db from argparse import ArgumentParser from datetime import datetime from asyncio import create_task @@ -20,20 +21,36 @@ EXEC_ARGS = PARSER.parse_args() async def birthday_anouncements_task(): LOGGER.debug('birthday_anouncements_task started.') for guild in bot.guilds: - guild_data = get_guild_data(guild.id) - if 'birthday_announcements_channel' in guild_data['features']['birthday'].keys(): - channel = guild.get_channel(guild_data['features']['birthday']['birthday_announcements_channel']) - if 'birthdays' in guild_data.keys(): - for user_id in guild_data['features']['birthday']['birthdays'].keys(): - user = guild.get_member(int(user_id)) - if user: - if guild_data['features']['birthday']['birthdays'][user_id]['announcements']: - date = guild_data['features']['birthday']['birthdays'][user_id]['date'] - today = datetime.today().strftime('%d/%m') - if date == today: - await channel.send(f'Joyeux anniversaire {user.mention}!') + birthday_settings = db.select('birthday_settings',f'guild_id = {guild.id}') + if birthday_settings: + if birthday_settings['is_enabled']: + for user in guild.members: + user_productivity_data = db.select('guild_user_productivity',f'user_id = {user.id} AND guild_id = {guild.id}') + if user_productivity_data: + if user_productivity_data['is_enabled']: + user = guild.get_member(user.id) + if user: + birthday = user_productivity_data['birthday'] + today = datetime.today().strftime('%d/%m') + if birthday == today: + channel = guild.get_channel(user_productivity_data['channel_id']) + if channel: + message = user_productivity_data['birthday_message'] + if message: + message = message.replace('{user}',user.mention) + await channel.send(message) + else: + LOGGER.debug(f'Channel {user_productivity_data["channel_id"]} not found.') + else: + LOGGER.debug(f'No birthday today for user {user.name} ({user.id}).') + else: + LOGGER.debug(f'User {user.name} ({user.id}) not found.') + else: + LOGGER.debug(f'User {user.name} ({user.id}) has birthday announcements disabled.') + else: + LOGGER.debug(f'No birthday data found for user {user.name} ({user.id}).') else: - LOGGER.debug(f'No birthday_announcements_channel set for guild {guild.name} ({guild.id}).') + LOGGER.debug(f'No productivity settings found for guild {guild.name} ({guild.id}).') LOGGER.debug('birthday_anouncements_task ended.') @bot.slash_command(name='help',description='Displays the help message.') @@ -63,20 +80,42 @@ async def reload(ctx:discord.ApplicationContext, feature:str): LOGGER.error(f'Error reloading feature {feature}: {e}') await ctx.send_response(f'Error reloading feature {feature}: {e}',ephemeral=True) +@bot.event +async def on_member_join(member:discord.Member): + if not member.bot: + user_id = member.id + user_name = member.name + user_mention = member.mention + db_user = db.select('global_user',f'WHERE user_id = {user_id}') + if not db_user: + db.insert('global_user',user_id=user_id,user_name=user_name,user_mention=user_mention) + +@bot.event +async def on_raw_member_remove(payload:discord.RawMemberRemoveEvent): + user = payload.user + if not user.bot: + user_id = user.id + guild_id = payload.guild_id + db.delete('guild_user_birthday',f'user_id = {user_id} AND guild_id = {guild_id}') + db.delete('guild_user_productivity',f'user_id = {user_id} AND guild_id = {guild_id}') @bot.event async def on_guild_join(guild:discord.Guild): guild_id = guild.id guild_name = guild.name LOGGER.info(f'Bot joined guild {guild_name} ({guild_id}).') - create_guild_data(guild_id) + db.insert('guild',guild_id=guild_id) + db.insert('birthday_settings',guild_id=guild_id,is_enabled=False,channel_id=random.choice(guild.text_channels).id) + db.insert('productivity_settings',guild_id=guild_id,is_enabled=False) @bot.event async def on_guild_remove(guild:discord.Guild): guild_id = guild.id guild_name = guild.name LOGGER.info(f'Bot left guild {guild_name} ({guild_id}).') - delete_guild_data(guild_id) + db.delete('birthday_settings',f'guild_id = {guild_id}') + db.delete('productivity_settings',f'guild_id = {guild_id}') + db.delete('guild',f'guild_id = {guild_id}') @bot.event async def on_ready(): @@ -88,30 +127,41 @@ async def on_ready(): await channel.send(f'{bot.user.name} Bot reloaded!',silent=True) birthday_anouncements_task.start() for guild in bot.guilds: - guild_data = get_guild_data(guild.id) - if 'features' in guild_data.keys(): - if 'productivity' in guild_data['features'].keys(): - if 'reminders' in guild_data['features']['productivity'].keys(): - for user_id in guild_data['features']['productivity']['reminders'].keys(): - user_data = guild_data['features']['productivity']['reminders'][user_id] - if user_data['enabled']: - next_reminder = user_data['next_reminder'] + for user in guild.members: + user_id = user.id + db_user = db.select('global_user',f'user_id = {user_id}') + if not db_user: + db.insert('global_user',user_id=user_id,user_name=user.name,user_mention=user.mention) + + productivity_settings = db.select('productivity_settings',f'guild_id = {guild.id}') + if productivity_settings: + if productivity_settings['is_enabled']: + for user in guild.members: + user_productivity_data = db.select('guild_user_productivity',f'user_id = {user.id} AND guild_id = {guild.id}') + if user_productivity_data: + if user_productivity_data['is_enabled']: + next_reminder = user_productivity_data['next_reminder'] if next_reminder > datetime.now().timestamp(): cooldown = next_reminder - datetime.now().timestamp() - user = guild.get_member(int(user_id)) - channel = guild.get_channel(user_data['channel']) - custom_message = user_data['custom_message'] - PRODUCTIVITY_TASKS[user_id] = create_task(send_reminder(cooldown,user,channel,custom_message)) + channel = guild.get_channel(user_productivity_data['channel_id']) + reminder_message = user_productivity_data['reminder_message'] + PRODUCTIVITY_TASKS[str(user.id)] = create_task(send_reminder(cooldown,user,channel,reminder_message)) else: - user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] - save_guild_data(guild.id, guild_data) - user = guild.get_member(int(user_id)) - channel = guild.get_channel(user_data['channel']) - custom_message = user_data['custom_message'] - PRODUCTIVITY_TASKS[user_id] = create_task(send_reminder(user_data['cooldown'],user,channel,custom_message)) + user_productivity_data['next_reminder'] = datetime.now().timestamp() + user_productivity_data['cooldown'] + db.update('guild_user_productivity',f'user_id = {user.id} AND guild_id = {guild.id}',next_reminder=user_productivity_data['next_reminder']) + channel = guild.get_channel(user_productivity_data['channel_id']) + reminder_message = user_productivity_data['reminder_message'] + PRODUCTIVITY_TASKS[str(user.id)] = create_task(send_reminder(user_productivity_data['cooldown'],user,channel,reminder_message)) + else: + LOGGER.debug(f'No productivity settings found for guild {guild.name} ({guild.id}).') if __name__ == '__main__': - load_dotenv() - TOKEN = os.getenv('TOKEN') - bot.run(TOKEN) + db.check_conn() + result = db.select('token','token_name = \'discord\'','token_value') + if result: + TOKEN = result['token_value'] + bot.run(TOKEN) + else: + LOGGER.error('No token found.') + sys.exit(1) \ No newline at end of file diff --git a/modules/cogs/birthday.py b/modules/cogs/birthday.py index f3c718d..e2e42e2 100644 --- a/modules/cogs/birthday.py +++ b/modules/cogs/birthday.py @@ -1,8 +1,10 @@ import discord, re from discord import option from discord.ext import commands +from datetime import datetime from ..setup.logger import LOGGER from ..setup.data import get_guild_data, save_guild_data, is_feature_enabled +from ..setup import db class BirthdayCog(commands.Cog): def __init__(self, bot): @@ -10,96 +12,106 @@ class BirthdayCog(commands.Cog): @commands.slash_command(description='Sets birthday date. Must be in DAY/MONTH format.') @option(name='date',description='Birthday date.',required=True) + @option(name='announcements',description='Enable or disable birthday announcements.',required=True,type=bool) + @option(name='message',description='The message to send when it is the user\'s birthday.',required=False) @option(name='for_user',description='The user to set the birthday date for.',required=False,type=discord.Member) - async def birthday_set(self,ctx:discord.ApplicationContext, date:str, for_user:discord.Member): + async def birthday_set(self,ctx:discord.ApplicationContext, date:str,announcements:bool,message:str,for_user:discord.Member): """Sets birthday date. Must be in DAY/MONTH format.""" command_name = 'birthday_set' guild_id = ctx.guild.id - guild_data = get_guild_data(guild_id) author = ctx.author - if is_feature_enabled('birthday',data=guild_data): - if re.match(r'\d{2}\/\d{2}',date): # Check if date is in DAY/MONTH format - feature_data = guild_data['features']['birthday'] - if 'birthdays' not in feature_data.keys(): - feature_data['birthdays'] = dict() - if str(for_user.id) not in feature_data['birthdays'].keys(): - feature_data[str(for_user.id)] = dict() + birthday_settings = db.select('birthday_settings',f'guild_id = {guild_id}') + if birthday_settings: + if birthday_settings['is_enabled']: if for_user: - if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user - LOGGER.debug(f'{author.name} used /{command_name} {date} for {for_user.name}.') - feature_data['birthdays'][str(for_user.id)]['date'] = date - feature_data['birthdays'][str(for_user.id)]['announcements'] = True - - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Birthday for {for_user.name} has been set to {date}.',ephemeral=True) + if author.guild_permissions.administrator: + if re.match(r'^\d{1,2}\/\d{1,2}$', date): + user_data = db.select('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {for_user.id}') + if user_data: + # Update user data + if message: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} {message} for {for_user.name}.') + db.update('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {for_user.id}',birthday=date,is_enabled=announcements,birthday_message=message) + else: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} for {for_user.name}.') + db.update('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {for_user.id}',birthday=date,is_enabled=announcements) + await ctx.send_response(f'{for_user.name}\'s birthday has been set to {date}.',ephemeral=True) + else: + # Create user data + if message: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} {message} for {for_user.name}.') + db.insert('guild_user_birthday',guild_id=guild_id,user_id=for_user.id,birthday=date,is_enabled=announcements,birthday_message=message) + else: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} for {for_user.name}.') + db.insert('guild_user_birthday',guild_id=guild_id,user_id=for_user.id,birthday=date,is_enabled=announcements) + await ctx.send_response(f'{for_user.name}\'s birthday has been set to {date}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} {announcements} {message} for {for_user.name} but the date format is invalid.') + await ctx.send_response('The date format must be DAY/MONTH.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {date} for {for_user.name} but is not an administrator.') + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} {announcements} {message} for {for_user.name} but is not an administrator.') await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) else: - # Set birthday for author - LOGGER.debug(f'{author.name} used /{command_name} {date}.') - feature_data['birthdays'][str(author.id)]['date'] = date - feature_data['birthdays'][str(author.id)]['announcements'] = True - save_guild_data(guild_id, guild_data) - - # The status defines the human readable status of the announcements to be displayed in the response. - announcements_status = 'will be' if feature_data['birthdays'][str(author.id)]['announcements'] else 'will not be' - await ctx.send_response(f'Your birthday has been set to {date} and {announcements_status} announced.',ephemeral=True) + if re.match(r'^\d{1,2}\/\d{1,2}$', date): + user_data = db.select('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {author.id}') + if user_data: + # Update user data + if message: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} {message}.') + db.update('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {author.id}',birthday=date,is_enabled=announcements,birthday_message=message) + else: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements}.') + db.update('guild_user_birthday',f'guild_id = {guild_id} AND user_id = {author.id}',birthday=date,is_enabled=announcements) + await ctx.send_response(f'Your birthday has been set to {date}.',ephemeral=True) + else: + # Create user data + if message: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements} {message}.') + db.insert('guild_user_birthday',guild_id=guild_id,user_id=author.id,birthday=date,is_enabled=announcements,birthday_message=message) + else: + LOGGER.debug(f'{author.name} used /{command_name} {date} {announcements}.') + db.insert('guild_user_birthday',guild_id=guild_id,user_id=author.id,birthday=date,is_enabled=announcements) + await ctx.send_response(f'Your birthday has been set to {date}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} {announcements} {message} but the date format is invalid.') + await ctx.send_response('The date format must be DAY/MONTH.',ephemeral=True) else: - LOGGER.debug(f'{author.name} used /set_birthday {date} but the format is not correct.') - await ctx.send_response('The date must is DAY/MONTH format.',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} {announcements} {message} but the birthday feature is not enabled.') + await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /set_birthday {date} but the birthday feature is not enabled.') + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} {announcements} {message} but the birthday feature is not enabled.') await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) - @commands.slash_command(description='Enable or disable birthday announcements.') - @option(name='enable',description='Enable birthday announcements.',required=True) - @option(name='for_user',description='The user to enable or disable birthday announcements for.',required=False,type=discord.Member) - async def birthday_announcements(self,ctx:discord.ApplicationContext, enable:bool, for_user:discord.Member): - """Enable or disable birthday announcements.""" - command_name = 'birthday_announcements' - guild_id = ctx.guild.id - guild_data = get_guild_data(guild_id) or dict() - author = ctx.author - if is_feature_enabled('birthday',data=guild_data): - feature_data = guild_data['features']['birthday'] - if for_user: - if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user - LOGGER.debug(f'{author.name} used /{command_name} {enable} for {for_user.name}.') - feature_data[str(for_user.id)]['announcements'] = enable - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Announcements for {for_user.name} are set to {enable}',ephemeral=True) - else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} for {for_user.name} but is not an administrator.') - await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) - else: - LOGGER.debug(f'{author.name} used /{command_name} {enable}.') - feature_data['birthdays'][str(author.id)]['announcements'] = enable - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Your birthday announcements have been set to {enable}',ephemeral=True) - else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} but the birthday feature is not enabled.') - await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) - - @commands.slash_command(description='Sets the channel to send birthday announcements to.') + @commands.slash_command(description='Configurate birthday announcements.') @option(name='channel',description='The channel to send birthday announcements to.',required=True,type=discord.TextChannel) - async def birthday_announcements_channel(self,ctx:discord.ApplicationContext, channel:discord.TextChannel): - """Sets the channel to send birthday announcements to.""" - command_name = 'birthday_announcements_channel' + @option(name='enable',description='Enable or disable birthday announcements.',required=True,type=bool) + async def birthday_config(self,ctx:discord.ApplicationContext, channel:discord.TextChannel,enable:bool): + """Configurate birthday announcements.""" + command_name = 'birthday_config' guild_id = ctx.guild.id - guild_data = get_guild_data(guild_id) or dict() author = ctx.author - if is_feature_enabled('birthday',data=guild_data): - feature_data = guild_data['features']['birthday'] - if author.guild_permissions.administrator: - LOGGER.debug(f'{author.name} used /{command_name} {channel.name}.') - feature_data['birthday_announcements_channel'] = channel.id - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Birthday announcements will be sent to {channel.mention}.',ephemeral=True) + birthday_settings = db.select('birthday_settings',f'guild_id = {guild_id}') + if birthday_settings: + # Update guild birthday settings + if birthday_settings['is_enabled']: + if author.guild_permissions.administrator: + LOGGER.debug(f'{author.name} used /{command_name} {channel.name} {enable}.') + db.update('birthday_settings',f'guild_id = {guild_id}',channel_id=channel.id,is_enabled=enable) + await ctx.send_response(f'Birthday announcements have been {"enabled" if enable else "disabled"} in {channel.mention}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.name} {enable} but is not an administrator.') + await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but is not an administrator.') - await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.name} {enable} but the birthday feature is not enabled.') + await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but the birthday feature is not enabled.') - await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) + # Create guild birthday settings + if author.guild_permissions.administrator: + LOGGER.debug(f'{author.name} used /{command_name} {channel.name} {enable}.') + db.insert('birthday_settings',guild_id=guild_id,channel_id=channel.id,is_enabled=enable) + await ctx.send_response(f'Birthday announcements have been {"enabled" if enable else "disabled"} in {channel.mention}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.name} {enable} but is not an administrator.') + await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) + diff --git a/modules/cogs/productivity.py b/modules/cogs/productivity.py index d705060..813ee4f 100644 --- a/modules/cogs/productivity.py +++ b/modules/cogs/productivity.py @@ -2,21 +2,20 @@ import discord from datetime import datetime from discord.ext import commands, tasks from discord import option -from typing import Union +from typing import Union, Optional from asyncio import sleep, create_task from ..setup.logger import LOGGER +from ..setup import db from ..setup.data import get_guild_data, save_guild_data, is_feature_enabled TASKS = {} -async def send_reminder(cooldown:int,author:discord.Member,channel:discord.TextChannel,reminder_message:str=None): +async def send_reminder(cooldown:int,author:discord.Member,channel:discord.TextChannel,reminder_message:str): while True: await sleep(cooldown) LOGGER.debug(f'Sending productivity reminder to {author.name} in {channel.name} of {channel.guild.name}.') - if reminder_message: - await channel.send(f'{author.mention} {reminder_message}') - else: - await channel.send(f'{author.mention} You have not been active in this channel for a while. Please consider being more productive.') + message = reminder_message.replace('{user}',author.mention) + await channel.send(message) class ProductivityCog(commands.Cog): @@ -29,79 +28,124 @@ class ProductivityCog(commands.Cog): @option(name='days',description='Number of days to wait before sending a reminder if the user has not been active in the channel.',required=True,type=int) @option(name='custom_message',description='Custom message to send with the reminder.',required=False,type=str) @option(name='for_user',description='The user to configurate productivity reminders for.',required=False,type=discord.Member) - async def productivity_config(self,ctx:discord.ApplicationContext, channel:Union[discord.TextChannel,discord.Thread], enable:bool, days:int, custom_message=None, for_user:discord.Member=None): + async def productivity_set(self,ctx:discord.ApplicationContext, channel:Union[discord.TextChannel,discord.Thread], enable:bool, days:int, custom_message=None, for_user:discord.Member=None): """Configurate productivity reminder.""" + command_name = 'productivity_set' + guild_id = ctx.guild.id + author = ctx.author + productivity_settings = db.select('productivity_settings',f'guild_id = {guild_id}') + if productivity_settings: + if productivity_settings['is_enabled']: + if for_user: + if author.guild_permissions.administrator: + user_data = db.select('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {for_user.id}') + cooldown = days*60*60*24 + next_reminder = int(datetime.now().timestamp() + cooldown) + if user_data: + # Update user data + if custom_message: + db.update('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {for_user.id}',channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,reminder_message=custom_message,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} {custom_message} for {for_user.name}.') + await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) + else: + db.update('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {for_user.id}',channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} for {for_user.name}.') + await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) + else: + # Create user data + if custom_message: + db.insert('guild_user_productivity',guild_id=guild_id,user_id=for_user.id,channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,reminder_message=custom_message,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} {custom_message} for {for_user.name}.') + await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) + else: + db.insert('guild_user_productivity',guild_id=guild_id,user_id=for_user.id,channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} for {for_user.name}.') + await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) + user_data = db.select('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {for_user.id}') + if str(for_user.id) in TASKS.keys(): + TASKS[str(for_user.id)].cancel() + TASKS[str(for_user.id)] = create_task(send_reminder(cooldown,for_user,channel,user_data['reminder_message'])) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} {enable} {days} {custom_message} for {for_user.name} but is not an administrator.') + await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) + else: + user_data = db.select('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {author.id}') + cooldown = days*60*60*24 + next_reminder = int(datetime.now().timestamp() + cooldown) + if user_data: + # Update user data + if custom_message: + db.update('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {author.id}',channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,reminder_message=custom_message,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} {custom_message}.') + await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + else: + db.update('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {author.id}',channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days}.') + await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + else: + # Create user data + if custom_message: + db.insert('guild_user_productivity',guild_id=guild_id,user_id=author.id,channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,reminder_message=custom_message,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days} {custom_message}.') + await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + else: + db.insert('guild_user_productivity',guild_id=guild_id,user_id=author.id,channel_id=channel.id,cooldown=cooldown,next_reminder=next_reminder,is_enabled=enable) + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {enable} {days}.') + await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + user_data = db.select('guild_user_productivity',f'guild_id = {guild_id} AND user_id = {author.id}') + if str(author.id) in TASKS.keys(): + TASKS[str(author.id)].cancel() + TASKS[str(author.id)] = create_task(send_reminder(cooldown,author,channel,user_data['reminder_message'])) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} {enable} {days} {custom_message} but the productivity feature is not enabled.') + await ctx.send_response('The productivity feature is not enabled.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} {enable} {days} {custom_message} but no productivity settings were found.') + await ctx.send_response('No productivity settings were found.',ephemeral=True) + + @commands.slash_command(description='Enable or disable productivity reminders on the server.') + @option(name='enable',description='Enable or disable productivity reminders.',required=True) + async def productivity_config(self,ctx:discord.ApplicationContext, enable:bool): + """Enable or disable productivity reminders on the server.""" command_name = 'productivity_config' guild_id = ctx.guild.id - guild_data = get_guild_data(guild_id) author = ctx.author - if is_feature_enabled('productivity',data=guild_data): - feature_data = guild_data['features']['productivity'] - if 'reminders' not in feature_data.keys(): - feature_data['reminders'] = dict() - if str(author.id) not in feature_data['reminders'].keys(): - feature_data['reminders'][str(author.id)] = dict() - if for_user: - if author.guild_permissions.administrator: - user_data = feature_data['reminders'][str(for_user.id)] - LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {for_user.name}.') - user_data['channel'] = channel.id - user_data['enabled'] = enable - user_data['cooldown'] = days#*60*60*24 - user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] - user_data['custom_message'] = custom_message - save_guild_data(guild_id, guild_data) - # create a task for the user - if for_user.id in TASKS.keys(): - TASKS[str(for_user.id)].cancel() - TASKS[str(for_user.id)] = create_task(send_reminder(user_data['cooldown'],for_user,channel,user_data['custom_message'])) - LOGGER.debug(f'{author.name} set up productivity reminders for {for_user.name}.') - await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) - else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} {for_user.name} but is not an administrator.') - await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) + productivity_settings = db.select('productivity_settings',f'guild_id = {guild_id}') + if productivity_settings: + # Update guild productivity settings + if author.guild_permissions.administrator: + LOGGER.debug(f'{author.name} used /{command_name} {enable}.') + db.update('productivity_settings',f'guild_id = {guild_id}',is_enabled=enable) + await ctx.send_response(f'Productivity reminders have been {"enabled" if enable else "disabled"}.',ephemeral=True) else: - LOGGER.debug(f'{author.name} used /{command_name} {channel.mention}.') - user_data = feature_data['reminders'][str(author.id)] - user_data['channel'] = channel.id - user_data['enabled'] = enable - user_data['cooldown'] = days#*60*60*24 - user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] - user_data['custom_message'] = custom_message - save_guild_data(guild_id, guild_data) - # create a task for the user - if str(author.id) in TASKS.keys(): - TASKS[str(author.id)].cancel() - TASKS[str(author.id)] = create_task(send_reminder(user_data['cooldown'],author,channel,user_data['custom_message'])) - LOGGER.debug(f'{author.name} set up productivity reminders for {author.name}.') - await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} but is not an administrator.') + await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but the productivity feature is not enabled.') - await ctx.send_response('The productivity feature is not enabled.',ephemeral=True) - + # Create guild productivity settings + if author.guild_permissions.administrator: + LOGGER.debug(f'{author.name} used /{command_name} {enable}.') + db.insert('productivity_settings',guild_id=guild_id,is_enabled=enable) + await ctx.send_response(f'Productivity reminders have been {"enabled" if enable else "disabled"}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} but is not an administrator.') + await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) + @commands.Cog.listener() async def on_message(self,message:discord.Message): author = message.author - if author.id != self.bot.user.id: + if author.id != self.bot.user.id: # Ignore messages from the bot guild_id = message.guild.id - guild_data = get_guild_data(guild_id) - if is_feature_enabled('productivity',data=guild_data): - channel = message.channel - feature_data = guild_data['features']['productivity'] - if 'reminders' in feature_data.keys(): - if str(author.id) in feature_data['reminders'].keys(): # If a user sat up its reminders - user_data = feature_data['reminders'][str(author.id)] - if user_data['enabled']: # If reminders are enabled for this user - reminder_channel_id = user_data['channel'] - if channel.id == reminder_channel_id: - user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] - save_guild_data(guild_id, guild_data) - if author.id in TASKS.keys(): - TASKS[str(author.id)].cancel() - TASKS[str(author.id)] = create_task(send_reminder(user_data['cooldown'],author,channel,user_data['custom_message'])) - LOGGER.debug(f'{author.name} sent a message in {channel.name} and reminders are set up in {reminder_channel_id}.') - else: - LOGGER.debug(f'{author.name} sent a message in {channel.name} but reminders are set up in {reminder_channel_id}.') + productivity_settings = db.select('productivity_settings',f'guild_id = {guild_id}') + if productivity_settings: + if productivity_settings['is_enabled']: + user_productivity_data = db.select('guild_user_productivity',f'user_id = {author.id} AND guild_id = {guild_id}') + if user_productivity_data: + if user_productivity_data['is_enabled']: + next_reminder = int(datetime.now().timestamp() + user_productivity_data['cooldown']) + db.update('guild_user_productivity',f'user_id = {author.id} AND guild_id = {guild_id}',next_reminder=next_reminder) + if str(author.id) in TASKS.keys(): + TASKS[str(author.id)].cancel() + TASKS[str(author.id)] = create_task(send_reminder(user_productivity_data['cooldown'],author,message.channel,user_productivity_data['reminder_message'])) \ No newline at end of file diff --git a/modules/setup/data.py b/modules/setup/data.py deleted file mode 100644 index 03b5913..0000000 --- a/modules/setup/data.py +++ /dev/null @@ -1,39 +0,0 @@ -import os, json -from .logger import LOGGER - -FEATURES = ['birthday','productivity'] - -DATA_DIR = 'data/' - -def get_guild_data(guild_id:int): - """Reads and returns the dictionary of guild data from the guild's JSON data file.""" - path = os.path.join(DATA_DIR, f'{guild_id}.json') - with open(path, 'r') as f: - return json.load(f) - -def save_guild_data(guild_id:int, data): - """Saves the dictionary of guild data to the guild's JSON data file.""" - path = os.path.join(DATA_DIR, f'{guild_id}.json') - with open(path, 'w') as f: - json.dump(data, f, indent=2) - LOGGER.debug(f'Saved data file for guild {guild_id}.') - -def create_guild_data(guild_id:int): - """Creates a new guild data file.""" - path = os.path.join(DATA_DIR, f'{guild_id}.json') - with open(path,'w') as f: - json.dump({'features': {key:{'enabled':True} for key in FEATURES}}, f, indent=2) - LOGGER.debug(f'Created data file for guild {guild_id}.') - -def delete_guild_data(guild_id:int): - """Deletes the guild data file.""" - path = os.path.join(DATA_DIR, f'{guild_id}.json') - os.remove(path) - LOGGER.debug(f'Deleted data file for guild {guild_id}.') - -def is_feature_enabled(feature:str,data=None,guild_id=0): - """Returns True if the feature is enabled in the guild, False otherwise.""" - if data: - return data['features'][feature]['enabled'] - else: - return get_guild_data(guild_id)['features'][feature]['enabled'] \ No newline at end of file diff --git a/modules/setup/db.py b/modules/setup/db.py new file mode 100644 index 0000000..842a17b --- /dev/null +++ b/modules/setup/db.py @@ -0,0 +1,304 @@ +import os, psycopg2 +from dotenv import load_dotenv +from typing import Optional, Callable, Any, Union +from .logger import LOGGER + +load_dotenv() +DB_NAME = os.getenv('DB_NAME') +DB_USER = os.getenv('DB_USER') +DB_PASSWORD = os.getenv('DB_PASSWORD') +DB_HOST = os.getenv('DB_HOST') + +def check_conn(): + """Tests the database connection.""" + connection = psycopg2.connect( + database=DB_NAME, + user=DB_USER, + password=DB_PASSWORD, + host=DB_HOST, + port="5432" + ) + if connection.status: + LOGGER.debug('Established connection to database.') + LOGGER.debug(f'Connection status: {connection.status}') + else: + LOGGER.error('Failed to connect to database.') + connection.close() + LOGGER.debug('Closed connection to database.') + +def insert(table:str,**kwargs): + """ + Inserts the given data into the given table. + + Global_User + user_id: int + user_mention: str + user_mention: str + + Guild + guild_id: int + + Birthday_Settings + guild_id: int + is_enabled: bool + channel_id: int + + Productivity_Settings + guild_id: int + is_enabled: bool + + Guild_User_Birthday + user_id: int + guild_id: int + birthday: int + birthday_message: str | None + is_enabled: bool + + Guild_User_Productivity + user_id: int + guild_id: int + is_enabled: bool + channel_id: int + cooldown: int + next_reminder: int + reminder_message: str | None + """ + connection = psycopg2.connect( + database=DB_NAME, + user=DB_USER, + password=DB_PASSWORD, + host=DB_HOST, + port="5432", + ) + cursor = connection.cursor() + + query = f"INSERT INTO {table} (" + params = [] + for key,value in kwargs.items(): + query += f"{key}," + params.append(value) + query = query.rstrip(",") # Remove trailing comma + query += ") VALUES (" + for _ in range(len(params)): + query += "%s," + query = query.rstrip(",") # Remove trailing comma + query += ")" + + query = cursor.mogrify(query, params) + cursor.execute(query) + connection.commit() + cursor.close() + connection.close() + LOGGER.debug(f'Inserted data into {table}.') + +def update(table:str,where:str,**kwargs): + """ + Updates the given data in the given table. + + Global_User + user_id: int + user_mention: str + user_mention: str + + Guild + guild_id: int + + Birthday_Settings + guild_id: int + is_enabled: bool + channel_id: int + + Productivity_Settings + guild_id: int + is_enabled: bool + + Guild_User_Birthday + user_id: int + guild_id: int + birthday: int + birthday_message: str | None + is_enabled: bool + + Guild_User_Productivity + user_id: int + guild_id: int + is_enabled: bool + channel_id: int + cooldown: int + next_reminder: int + reminder_message: str | None + """ + connection = psycopg2.connect( + database=DB_NAME, + user=DB_USER, + password=DB_PASSWORD, + host=DB_HOST, + port="5432", + ) + cursor = connection.cursor() + + query = f"UPDATE {table} SET " + params = [] + for key,value in kwargs.items(): + query += f"{key} = %s," + params.append(value) + query = query.rstrip(",") # Remove trailing comma + query += f" WHERE {where}" + + query = cursor.mogrify(query, params) + cursor.execute(query) + connection.commit() + cursor.close() + connection.close() + LOGGER.debug(f'Updated data in {table}.') + +def delete(table:str,where:str): + """ + Deletes the given data from the given table. + + Global_User + user_id: int + user_mention: str + user_mention: str + + Guild + guild_id: int + + Birthday_Settings + guild_id: int + is_enabled: bool + channel_id: int + + Productivity_Settings + guild_id: int + is_enabled: bool + + Guild_User_Birthday + user_id: int + guild_id: int + birthday: int + birthday_message: str | None + is_enabled: bool + + Guild_User_Productivity + user_id: int + guild_id: int + is_enabled: bool + channel_id: int + cooldown: int + next_reminder: int + reminder_message: str | None + """ + connection = psycopg2.connect( + database=DB_NAME, + user=DB_USER, + password=DB_PASSWORD, + host=DB_HOST, + port="5432", + ) + cursor = connection.cursor() + + query = f"DELETE FROM {table} WHERE {where}" + cursor.execute(query) + connection.commit() + cursor.close() + connection.close() + LOGGER.debug(f'Deleted data from {table}.') + +def select(table:str,where:str,*columns:str) -> dict[str,Any] | None: + """ + Selects the given data from the given table, and returns it as a dictionary. + + Global_User + user_id: int + user_mention: str + user_mention: str + + Guild + guild_id: int + + Birthday_Settings + guild_id: int + is_enabled: bool + channel_id: int + + Productivity_Settings + guild_id: int + is_enabled: bool + + Guild_User_Birthday + user_id: int + guild_id: int + birthday: int + birthday_message: str | None + is_enabled: bool + + Guild_User_Productivity + user_id: int + guild_id: int + is_enabled: bool + channel_id: int + cooldown: int + next_reminder: int + reminder_message: str | None + """ + connection = psycopg2.connect( + database=DB_NAME, + user=DB_USER, + password=DB_PASSWORD, + host=DB_HOST, + port="5432" + ) + connection.readonly = True + cursor = connection.cursor() + + query = f"SELECT " + if columns: + for column in columns: + query += f"{column}," + query = query.rstrip(",") + else: + query += "*" + query += f" FROM {table} WHERE {where}" + cursor.execute(query) + record = cursor.fetchone() + result = None + if record is None: + LOGGER.debug(f'No data found in {table} for {where}.') + else: + result = {} + if len(columns) > 0: + for i in range(len(columns)): + result[columns[i]] = record[i] + else: + if table.lower() == 'global_user': + result['user_id'] = record[0] + result['user_name'] = record[1] + result['user_mention'] = record[2] + elif table.lower() == 'guild': + result['guild_id'] = record[0] + elif table.lower() == 'birthday_settings': + result['guild_id'] = record[0] + result['is_enabled'] = record[1] + result['channel_id'] = record[2] + elif table.lower() == 'productivity_settings': + result['guild_id'] = record[0] + result['is_enabled'] = record[1] + elif table.lower() == 'guild_user_birthday': + result['user_id'] = record[0] + result['guild_id'] = record[1] + result['birthday'] = record[2] + result['is_enabled'] = record[3] + elif table.lower() == 'guild_user_productivity': + result['user_id'] = record[0] + result['guild_id'] = record[1] + result['is_enabled'] = record[2] + result['channel_id'] = record[3] + result['cooldown'] = record[4] + result['next_reminder'] = record[5] + result['reminder_message'] = record[6] + else: + LOGGER.error(f'Unknown table {table}.') + cursor.close() + connection.close() + return result \ No newline at end of file