diff --git a/bot.py b/bot.py index 0452bf30..24421ada 100644 --- a/bot.py +++ b/bot.py @@ -131,15 +131,18 @@ def restart(): async def database_updates(): - handler = schema.SchemaHandler(db_file) + handler = schema.SchemaHandler(db_file, bot) if handler.version == 0: - handler.update() + handler.zero_to_one() messages = db.fetch_all_messages() for message in messages: channel_id = message[1] channel = await getchannel(channel_id) db.add_guild(channel.id, channel.guild.id) + if handler.version == 1: + handler.one_to_two() + async def system_notification(guild_id, text): # Send a message to the system channel (if set) @@ -154,9 +157,12 @@ async def system_notification(guild_id, text): ) return + if server_channel: + server_channel = server_channel[0][0] + if server_channel: try: - target_channel = await getchannel(server_channel[0][0]) + target_channel = await getchannel(server_channel) await target_channel.send(text) except discord.Forbidden: @@ -379,7 +385,6 @@ async def on_ready(): ) await database_updates() - db.migrate_admins(bot) maintain_presence.start() cleandb.start() check_cleanup_queued_guilds.start() @@ -436,6 +441,8 @@ async def on_raw_reaction_add(payload): if user_id != bot.user.id: try: await member.add_roles(role) + if db.notify(guild_id): + await user.send(f"You now have the following role: **{role.name}**") except discord.Forbidden: await system_notification( @@ -483,6 +490,8 @@ async def on_raw_reaction_remove(payload): role = discord.utils.get(server.roles, id=role_id) try: await member.remove_roles(role) + if db.notify(guild_id): + await member.send(f"You do not have the following role anymore: **{role.name}**") except discord.Forbidden: await system_notification( @@ -1054,6 +1063,23 @@ async def set_systemchannel(ctx): else: await ctx.send("You do not have an admin role.") + +@bot.command(name="notify") +async def toggle_notify(ctx): + if isadmin(ctx.message.author, ctx.guild.id): + notify = db.toggle_notify(ctx.guild.id) + if notify: + await ctx.send( + "Notifications have been set to ON for this server.\n" + "Use this command again to turn them off." + ) + else: + await ctx.send( + "Notifications have been set to OFF for this server.\n" + "Use this command again to turn them on." + ) + + @commands.is_owner() @bot.command(name="colour") async def set_colour(ctx): diff --git a/core/database.py b/core/database.py index 7d7cb777..1bf55773 100644 --- a/core/database.py +++ b/core/database.py @@ -42,14 +42,13 @@ def initialize(database): cursor.execute("CREATE TABLE IF NOT EXISTS 'cleanup_queue_guilds' ('guild_id' INT, 'unix_timestamp' INT);") cursor.execute("CREATE TABLE IF NOT EXISTS 'dbinfo' ('version' INT);") cursor.execute( - "CREATE TABLE IF NOT EXISTS 'systemchannels' ('guild_id' INT, 'channel_id'" - " INT);" + "CREATE TABLE IF NOT EXISTS 'guild_settings' ('guild_id' INT, 'notify' INT, 'systemchannel' INT);" ) cursor.execute( - "CREATE UNIQUE INDEX IF NOT EXISTS reactionrole_idx ON messages (reactionrole_id);" + "CREATE UNIQUE INDEX IF NOT EXISTS guild_id_idx ON guild_settings (guild_id);" ) cursor.execute( - "CREATE UNIQUE INDEX IF NOT EXISTS guild_id_idx ON systemchannels (guild_id);" + "CREATE UNIQUE INDEX IF NOT EXISTS reactionrole_idx ON messages (reactionrole_id);" ) cursor.execute( "CREATE UNIQUE INDEX IF NOT EXISTS guild_id_index ON cleanup_queue_guilds (guild_id);" @@ -66,41 +65,6 @@ def __init__(self, database): self.reactionrole_creation = {} - def migrate_admins(self, client): - import discord - conn = sqlite3.connect(self.database) - cursor = conn.cursor() - cursor.execute("PRAGMA table_info(admins);") - result = cursor.fetchall() - columns = [value[1] for value in result] - if "guild_id" not in columns: - cursor.execute("SELECT role_id FROM admins") - admins = cursor.fetchall() - admins2 = [] - for admin in admins: - admins2.append(admin[0]) - admins = admins2 - guilds = {} - for guild in client.guilds: - guilds[guild.id] = [] - for admin_id in admins: - role = discord.utils.get(guild.roles, id=admin_id) - if role is not None: - guilds[guild.id].append(role.id) - - cursor.execute("ALTER TABLE admins ADD COLUMN 'guild_id' INT;") - conn.commit() - for guild in guilds: - for admin_id in guilds[guild]: - cursor.execute("UPDATE admins SET guild_id = ? WHERE role_id = ?;", (guild, admin_id)) - conn.commit() - cursor.execute("DELETE FROM admins WHERE guild_id IS NULL;") - conn.commit() - print("Successfully migrated admins.") - - cursor.close() - conn.close() - def add_reaction_role(self, rl_dict: dict): try: conn = sqlite3.connect(self.database) @@ -237,9 +201,9 @@ def remove_guild(self, guild_id): "DELETE FROM reactionroles WHERE reactionrole_id = ?;", (reactionrole_id,), ) - # Deleting the guilds systemchannels database entries + # Deleting the guilds guild_settings database entries cursor.execute( - "DELETE FROM systemchannels WHERE guild_id = ?;", + "DELETE FROM guild_settings WHERE guild_id = ?;", (guild_id,), ) # Delete the guilds admin roles @@ -337,10 +301,15 @@ def add_systemchannel(self, guild_id, channel_id): try: conn = sqlite3.connect(self.database) cursor = conn.cursor() + notify = 0 cursor.execute( - "REPLACE INTO 'systemchannels' ('guild_id', 'channel_id')" + "INSERT OR IGNORE INTO guild_settings ('guild_id', 'notify', 'systemchannel')" " values(?, ?);", - (guild_id, channel_id), + (guild_id, notify, channel_id), + ) + cursor.execute( + "UPDATE guild_settings SET systemchannel = ? WHERE guild_id = ?;", + (channel_id, guild_id), ) conn.commit() cursor.close() @@ -353,8 +322,16 @@ def remove_systemchannel(self, guild_id): try: conn = sqlite3.connect(self.database) cursor = conn.cursor() + channel_id = 0 # Set to false + notify = 0 cursor.execute( - "DELETE FROM systemchannels WHERE guild_id = ?;", (guild_id,) + "INSERT OR IGNORE INTO guild_settings ('guild_id', 'notify', 'systemchannel')" + " values(?, ?);", + (guild_id, notify, channel_id), + ) + cursor.execute( + "UPDATE guild_settings SET systemchannel = ? WHERE guild_id = ?;", + (channel_id, guild_id), ) conn.commit() cursor.close() @@ -368,7 +345,7 @@ def fetch_systemchannel(self, guild_id): conn = sqlite3.connect(self.database) cursor = conn.cursor() cursor.execute( - "SELECT channel_id FROM systemchannels WHERE guild_id = ?;", (guild_id,) + "SELECT systemchannel FROM guild_settings WHERE guild_id = ?;", (guild_id,) ) result = cursor.fetchall() cursor.close() @@ -385,7 +362,7 @@ def fetch_all_guilds(self): cursor.execute("SELECT guild_id FROM messages;") message_guilds = cursor.fetchall() - cursor.execute("SELECT guild_id FROM systemchannels;") + cursor.execute("SELECT guild_id FROM guild_settings;") systemchannel_guilds = cursor.fetchall() cursor.execute("SELECT guild_id FROM admins;") @@ -505,3 +482,55 @@ def fetch_cleanup_guilds(self, guild_ids_only=False): except sqlite3.Error as e: return e + + def toggle_notify(self, guild_id: int): + # SQLite doesn't support booleans + # INTs are used: 1 = True, 0 = False + try: + conn = sqlite3.connect(self.database) + cursor = conn.cursor() + cursor.execute("SELECT notify FROM guild_settings WHERE guild_id = ?", (guild_id,)) + results = cursor.fetchall() + if not results: + # If the guild was not in the table because the command was never used before + notify = 1 + systemchannel = 0 + cursor.execute("INSERT INTO 'guild_settings' ('guild_id', 'notify') values(?, ?, ?);", (guild_id, systemchannel, notify)) + else: + notify = results[0][0] + if notify: + notify = 0 + cursor.execute("UPDATE guild_settings SET notify = ? WHERE guild_id = ?", (notify, guild_id)) + + else: + notify = 1 + cursor.execute("UPDATE guild_settings SET notify = ? WHERE guild_id = ?", (notify, guild_id)) + conn.commit() + cursor.close() + conn.close() + return notify + + except sqlite3.Error as e: + return e + + def notify(self, guild_id: int): + # SQLite doesn't support booleans + # INTs are used: 1 = True, 0 = False + try: + conn = sqlite3.connect(self.database) + cursor = conn.cursor() + cursor.execute("SELECT notify FROM guild_settings WHERE guild_id = ?", (guild_id,)) + results = cursor.fetchall() + if not results: + # If the guild was not in the table because the command was never used before + notify = 0 + systemchannel = 0 + cursor.execute("INSERT INTO 'guild_settings' ('guild_id', `systemchannel`, 'notify') values(?, ?, ?);", (guild_id, systemchannel, notify)) + else: + notify = results[0][0] + cursor.close() + conn.close() + return notify + + except sqlite3.Error as e: + return e diff --git a/core/schema.py b/core/schema.py index 340d9d8e..1c66b895 100644 --- a/core/schema.py +++ b/core/schema.py @@ -24,11 +24,13 @@ import sqlite3 +import discord class SchemaHandler: - def __init__(self, database): + def __init__(self, database, client): self.database = database + self.client = client self.version = self.version_check() def version_check(self): @@ -45,10 +47,6 @@ def version_check(self): version = version[0][0] return version - def update(self): - if self.version == 0: - self.zero_to_one() - def set_version(self, version): conn = sqlite3.connect(self.database) cursor = conn.cursor() @@ -79,3 +77,49 @@ def zero_to_one(self): cursor.close() conn.close() self.set_version(1) + + def one_to_two(self): + conn = sqlite3.connect(self.database) + cursor = conn.cursor() + cursor.execute("PRAGMA table_info(admins);") + result = cursor.fetchall() + columns = [value[1] for value in result] + if "guild_id" not in columns: + cursor.execute("SELECT role_id FROM admins") + admins = cursor.fetchall() + admins2 = [] + for admin in admins: + admins2.append(admin[0]) + admins = admins2 + guilds = {} + for guild in self.client.guilds: + guilds[guild.id] = [] + for admin_id in admins: + role = discord.utils.get(guild.roles, id=admin_id) + if role is not None: + guilds[guild.id].append(role.id) + + cursor.execute("ALTER TABLE admins ADD COLUMN 'guild_id' INT;") + conn.commit() + for guild in guilds: + for admin_id in guilds[guild]: + cursor.execute("UPDATE admins SET guild_id = ? WHERE role_id = ?;", (guild, admin_id)) + cursor.execute("DELETE FROM admins WHERE guild_id IS NULL;") + conn.commit() + + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='systemchannels';") + systemchannels_table = cursor.fetchall() + if systemchannels_table: + cursor.execute("SELECT * FROM systemchannels;") + entries = cursor.fetchall() + for entry in entries: + guild_id = entry[0] + channel_id = entry[1] + notify = 0 # Set default to not notify + cursor.execute("INSERT INTO guild_settings ('guild_id', 'notify', 'systemchannel') values(?, ?, ?);", (guild_id, notify, channel_id)) + cursor.execute("DROP TABLE systemchannels;") + conn.commit() + + cursor.close() + conn.close() + self.set_version(2)