add notification toggles and checks

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)
@@ -379,7 +382,6 @@ async def on_ready():
         )
 
     await database_updates()
-    db.migrate_admins(bot)
     maintain_presence.start()
     cleandb.start()
     check_cleanup_queued_guilds.start()
@@ -436,6 +438,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.mention}")
 
                 except discord.Forbidden:
                     await system_notification(
@@ -483,6 +487,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 user.send(f"You do not have the following role anymore: {role.mention}")
 
             except discord.Forbidden:
                 await system_notification(
@@ -1054,6 +1060,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:
+            ctx.send(
+                "Notifications have been set to ON for this server.\n"
+                "Use this command again to turn them off."
+            )
+        else:
+            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):

core/database.py

@@ -45,6 +45,12 @@ def initialize(database):
         "CREATE TABLE IF NOT EXISTS 'systemchannels' ('guild_id' INT, 'channel_id'"
         " INT);"
     )
+    cursor.execute(
+        "CREATE TABLE IF NOT EXISTS 'guild_settings' ('guild_id' INT, 'notify' INT);"
+    )
+    cursor.execute(
+        "CREATE UNIQUE INDEX IF NOT EXISTS guild_id_idx ON guild_settings (guild_id);"
+    )
     cursor.execute(
         "CREATE UNIQUE INDEX IF NOT EXISTS reactionrole_idx ON messages (reactionrole_id);"
     )
@@ -66,41 +72,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)
@@ -505,3 +476,53 @@ 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
+                cursor.execute("INSERT INTO 'guild_settings' ('guild_id', 'notify') values(?,?);", (guild_id, 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
+                cursor.execute("INSERT INTO 'guild_settings' ('guild_id', 'notify') values(?,?);", (guild_id, notify))
+            else:
+                notify = results[0][0]
+            cursor.close()
+            conn.close()
+            return notify
+
+        except sqlite3.Error as e:
+            return e

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,36 @@ 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.close()
+        conn.close()
+        self.set_version(2)