Closes #7

bot/requirements.txt

@@ -1,7 +1,7 @@
 asyncio==3.4.*
 discord.py==2.3.*
 psutil==5.9.*
-psycopg2-binary==2.9.*
+asyncpg==0.29.*
 pypandoc==1.13.*
 python-dotenv==1.0.*
 regex==2023.12.25

bot/src/cogs/galtinn.py

@@ -23,19 +23,16 @@ class Galtinn(commands.Cog):
         """
 
         self.bot = bot
-        self.cursor = self.bot.db_connection.cursor()
-
-        self.init_db()
 
         self.membership_check.start()
         self.verification_cleanup.start()
 
-    def init_db(self):
+    async def init_db(self):
         """
         Create the necessary tables for the cog to work
         """
 
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             CREATE TABLE IF NOT EXISTS galtinn_users (
                 discord_id BIGINT PRIMARY KEY,
@@ -43,7 +40,7 @@ class Galtinn(commands.Cog):
             );
             """
         )
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             CREATE TABLE IF NOT EXISTS galtinn_verification (
                 discord_id BIGINT PRIMARY KEY,
@@ -53,7 +50,7 @@ class Galtinn(commands.Cog):
             );
             """
         )
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             CREATE TABLE IF NOT EXISTS galtinn_roles (
                 discord_role_id BIGINT PRIMARY KEY,
@@ -66,7 +63,6 @@ class Galtinn(commands.Cog):
         self.bot.logger.info("Unloading cog")
         self.membership_check.cancel()
         self.verification_cleanup.cancel()
-        self.cursor.close()
 
     @tasks.loop(time=misc_utils.MIDNIGHT)
     async def membership_check(self):
@@ -76,14 +72,14 @@ class Galtinn(commands.Cog):
 
         # Fetch galtinn organization roles
         roles = {}
-        self.cursor.execute(
+        db_roles = await self.bot.db.fetch(
             """
             SELECT *
             FROM galtinn_roles;
             """
         )
-        if db_roles := self.cursor.fetchall():
-            for role_id, galtinn_org_id in db_roles:
+        if db_roles:
+            for role_id, galtinn_org_id in list(db_roles):
                 role = self.bot.get_role(role_id)
                 if not role:
                     role = await self.bot.fetch_role(role_id)
@@ -95,17 +91,17 @@ class Galtinn(commands.Cog):
                 roles[galtinn_org_id] = role
 
         # Fetch registered users
-        self.cursor.execute(
+        db_users = await self.bot.db.fetch(
             """
             SELECT *
             FROM galtinn_users;
             """
         )
-        if not (results := self.cursor.fetchall()):
+        if not db_users:
             self.bot.logger.info("No registered users found")
 
-        for result in results:
-            discord_id, galtinn_id = result
+        for db_user in db_users:
+            discord_id, galtinn_id = list(db_user)
 
             user = self.bot.get_user(discord_id)  # Try fetching discord user from cache
             if not user:
@@ -135,23 +131,31 @@ class Galtinn(commands.Cog):
         Clean up any pending verifications that have expired. This is in case the register command fails to do so
         """
 
-        self.cursor.execute(
+        results = await self.bot.db.fetch(
             """
             SELECT discord_id, expires
             FROM galtinn_verification;
             """
         )
-        for result in self.cursor.fetchall():
-            discord_id, expires = result
+        for result in results:
+            discord_id, expires = list(result)
             if expires < datetime.utcnow():
-                self.cursor.execute(
+                await self.bot.db.execute(
                     """
                     DELETE FROM galtinn_verification
-                    WHERE discord_id = %s;
+                    WHERE discord_id = $1;
                     """,
-                    (discord_id,),
+                    discord_id,
                 )
 
+    @verification_cleanup.before_loop
+    async def before_verification_cleanup(self):
+        """
+        Make sure bot is ready before starting the verification cleanup loop
+        """
+
+        await self.bot.wait_until_ready()
+
     @app_commands.checks.bot_has_permissions(embed_links=True)
     @galtinn_group.command(name="registrer", description="Koble Galtinnbrukeren din til Discord")
     async def register(self, interaction: discord.Interaction):
@@ -166,29 +170,29 @@ class Galtinn(commands.Cog):
         await interaction.response.defer()
 
         # Check if user is already registered
-        self.cursor.execute(
+        db_user = await self.bot.db.fetchrow(
             """
             SELECT *
             FROM galtinn_users
-            WHERE discord_id = %s
+            WHERE discord_id = $1;
             """,
-            (interaction.user.id,),
+            interaction.user.id,
         )
-        if self.cursor.fetchone():
+        if db_user:
             embed = embed_templates.error_warning("Du er allerede registrert med en Galtinnbruker!")
             await interaction.followup.send(embed=embed)
             return
 
         # Check if user is already pending verification
-        self.cursor.execute(
+        verification = await self.bot.db.fetchrow(
             """
             SELECT discord_id, challenge, state
             FROM galtinn_verification
-            WHERE discord_id = %s
+            WHERE discord_id = $1;
             """,
-            (interaction.user.id,),
+            interaction.user.id,
         )
-        if self.cursor.fetchone():
+        if verification:
             embed = embed_templates.error_warning("Du har allerede en pågående verifikasjon!")
             await interaction.followup.send(embed=embed)
             return
@@ -208,12 +212,14 @@ class Galtinn(commands.Cog):
         url = f"{base_url}?{urllib.parse.urlencode(params)}"
 
         # Insert into verfications table
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             INSERT INTO galtinn_verification
-            VALUES (%s, %s, %s)
+            VALUES ($1, $2, $3);
             """,
-            (interaction.user.id, challenge, state),
+            interaction.user.id,
+            challenge,
+            state,
         )
 
         self.bot.logger.info(f"Generated challenge for user {interaction.user.id}")
@@ -226,12 +232,12 @@ class Galtinn(commands.Cog):
         await interaction.followup.send(embed=embed, ephemeral=True)
 
         await asyncio.sleep(120)
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             DELETE FROM galtinn_verification
-            WHERE discord_id = %s;
+            WHERE discord_id = $1;
             """,
-            (interaction.user.id,),
+            interaction.user.id,
         )
 
     @app_commands.checks.bot_has_permissions(embed_links=True)
@@ -247,14 +253,14 @@ class Galtinn(commands.Cog):
 
         await interaction.response.defer()
 
-        self.cursor.execute(
+        db_user = await self.bot.db.fetchrow(
             """
             SELECT * FROM galtinn_users
-            WHERE discord_id = %s
+            WHERE discord_id = $1;
             """,
-            (interaction.user.id,),
+            interaction.user.id,
         )
-        if not self.cursor.fetchone():
+        if not db_user:
             embed = embed_templates.error_warning("Du er ikke registrert med en Galtinnbruker!")
             await interaction.followup.send(embed=embed)
             return
@@ -264,7 +270,7 @@ class Galtinn(commands.Cog):
             "Er du sikker på at du vil slette koblingen mellom Galtinnbrukeren din og Discord?\n\n"
             + "Merk at roller knyttet til Galtinn vil bli fjernet."
         )
-        await interaction.followup.send(embed=embed, view=DeleteView(self.bot, self.cursor), ephemeral=True)
+        await interaction.followup.send(embed=embed, view=DeleteView(self.bot), ephemeral=True)
 
     @commands.has_permissions(manage_guild=True)
     @commands.bot_has_permissions(embed_links=True)
@@ -293,12 +299,13 @@ class Galtinn(commands.Cog):
         role (discord.Role): Discord role object
         """
 
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             INSERT INTO galtinn_roles
-            VALUES (%s, %s)
+            VALUES ($1, $2);
             """,
-            (role.id, galtinn_org_id),
+            role.id,
+            galtinn_org_id,
         )
         embed = embed_templates.success(
             f"Rollen {role.mention} ble lagt til for organisasjonen med id {galtinn_org_id}"
@@ -316,12 +323,12 @@ class Galtinn(commands.Cog):
         galtinn_org_id (str): Galtinn organization ID
         """
 
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             DELETE FROM galtinn_roles
-            WHERE discord_role_id = %s
+            WHERE discord_role_id = $1;
             """,
-            (role.id,),
+            role.id,
         )
         embed = embed_templates.success("Rollenkobling fjernet!")
         await ctx.send(embed=embed)
@@ -336,20 +343,19 @@ class Galtinn(commands.Cog):
         ctx (commands.Context): Context object
         """
 
-        self.cursor.execute(
+        db_roles = await self.bot.db.fetch(
             """
             SELECT *
             FROM galtinn_roles;
             """
         )
-        results = self.cursor.fetchall()
-        if not results:
+        if not db_roles:
             embed = embed_templates.error_warning("Ingen roller tilknyttet Galtinn")
             await ctx.send(embed=embed)
             return
 
         roles_formatted = [
-            f"<@&{role_id}> - `{galtinn_org_id}`" for role_id, galtinn_org_id in results
+            f"<@&{row.r['discord_role_id']}> - `{row.r['galtinn_org_id']}`" for row in db_roles
         ]  # TODO: add pagination? requires slash command, hence why I decided against it
 
         embed = discord.Embed(title="Roller tilknyttet Galtinn")
@@ -358,10 +364,9 @@ class Galtinn(commands.Cog):
 
 
 class DeleteView(discord.ui.View):
-    def __init__(self, bot: commands.Bot, cursor):
+    def __init__(self, bot: commands.Bot):
         super().__init__()
         self.bot = bot
-        self.cursor = cursor
 
     @discord.ui.button(label="Ja, slett", style=discord.ButtonStyle.danger, custom_id="delete_yes")
     async def delete_yes(self, interaction: discord.Interaction, button: discord.ui.Button):
@@ -370,12 +375,12 @@ class DeleteView(discord.ui.View):
         # Remove Roles
 
         # Delete user from database
-        self.cursor.execute(
+        await self.bot.db.execute(
             """
             DELETE FROM galtinn_users
-            WHERE discord_id = %s
+            WHERE discord_id = $1;
             """,
-            (interaction.user.id,),
+            interaction.user.id,
         )
         embed = embed_templates.success("Du har slettet tilkoblingen til Galtinnbrukeren din!")
         await interaction.message.edit(embed=embed, view=None)
@@ -395,4 +400,6 @@ async def setup(bot: commands.Bot):
     bot (commands.Bot): Bot instance
     """
 
-    await bot.add_cog(Galtinn(bot))
+    cog = Galtinn(bot)
+    await cog.init_db()
+    await bot.add_cog(cog)

bot/src/run.py

@@ -1,8 +1,8 @@
 import os
 from time import time
 
+import asyncpg
 import discord
-import psycopg2
 from discord.ext import commands
 from dotenv import load_dotenv
 
@@ -25,14 +25,6 @@ class Bot(commands.Bot):
 
         self.cog_files = set(os.listdir("./src/cogs"))
 
-        self.db_connection = psycopg2.connect(
-            host=os.environ["DATABASE_HOST"],
-            dbname=os.environ["DATABASE_NAME"],
-            user=os.environ["DATABASE_USER"],
-            password=os.environ["DATABASE_PASSWORD"],
-        )
-        self.db_connection.autocommit = True  # Scary
-
         self.galtinn = {
             "api_url": os.environ.get("GALTINN_API_URL", "http://127.0.0.1:8000"),
             "client_id": os.environ.get("GALTINN_CLIENT_ID"),
@@ -42,8 +34,18 @@ class Bot(commands.Bot):
         self.guild_id = os.environ.get("BOT_DEV_GUILD", 1162158668079444199)
 
     async def setup_hook(self):
-        cogs = os.listdir("./src/cogs")
+        # DB needs to be setup here because contructor is not async
+        # Not ideal but does the job
+        credentials = {
+            "host": os.environ["DATABASE_HOST"],
+            "database": os.environ["DATABASE_NAME"],
+            "user": os.environ["DATABASE_USER"],
+            "password": os.environ["DATABASE_PASSWORD"],
+        }
+        self.db = await asyncpg.create_pool(**credentials)
+
         # Load cogs
+        cogs = os.listdir("./src/cogs")
         for file in cogs:
             if file.endswith(".py"):
                 name = file[:-3]

verification_server/requirements.txt

@@ -1,6 +1,6 @@
 aiohttp==3.9.*
 fastapi==0.110.*
 jinja2==3.1.*
-psycopg2-binary==2.9.*
+asyncpg==0.29.*
 python-dotenv==1.0.*
 uvicorn==0.29.*

verification_server/src/server.py

@@ -1,7 +1,7 @@
 import os
 
+import asyncpg
 import aiohttp
-import psycopg2
 from dotenv import load_dotenv
 from fastapi import FastAPI
 from fastapi import Request
@@ -23,15 +23,20 @@ GALTIINN_CLIENT_ID = os.environ["GALTINN_CLIENT_ID"]
 GALTINN_REDIRECT_URI = os.environ["GALTINN_REDIRECT_URI"]
 
 
-# DB
-db_connection = psycopg2.connect(
-    host=os.environ["DATABASE_HOST"],
-    dbname=os.environ["DATABASE_NAME"],
-    user=os.environ["DATABASE_USER"],
-    password=os.environ["DATABASE_PASSWORD"],
-)
-db_connection.autocommit = True  # Scary
-db_cursor = db_connection.cursor()
+@app.on_event("startup")
+async def startup():
+    credentials = {
+        "host": os.environ["DATABASE_HOST"],
+        "database": os.environ["DATABASE_NAME"],
+        "user": os.environ["DATABASE_USER"],
+        "password": os.environ["DATABASE_PASSWORD"],
+    }
+    app.state.pool = await asyncpg.create_pool(**credentials)
+
+
+@app.on_event("shutdown")
+async def shutdown():
+    await app.state.pool.close()
 
 
 @app.get("/")
@@ -43,14 +48,14 @@ async def index():
 async def callback(request: Request, code: str, state: str):
 
     # Check if user is pending verification
-    db_cursor.execute(
+    verification = await app.state.pool.fetchrow(
         """
         SELECT discord_id, challenge, state
-        FROM galtinn_verification WHERE state = %s
+        FROM galtinn_verification WHERE state = $1
         """,
-        (state,),
+        state,
     )
-    if not (db_result := db_cursor.fetchone()):
+    if not verification:
         return templates.TemplateResponse(
             "error.html",
             {
@@ -59,7 +64,7 @@ async def callback(request: Request, code: str, state: str):
             },
         )
 
-    discord_id, code_challenge, state = db_result
+    discord_id, code_challenge, state = list(verification.values())
 
     # Get auth_token from code
     payload = {
@@ -99,21 +104,22 @@ async def callback(request: Request, code: str, state: str):
             user = await r.json()
 
     # Enter user into database
-    db_cursor.execute(
+    await app.state.pool.execute(
         """
         INSERT INTO galtinn_users (discord_id, galtinn_id)
-        VALUES (%s, %s)
+        VALUES ($1, $2)
         """,
-        (discord_id, user["sub"]),  # TODO: use UUID?
+        discord_id,
+        user["sub"],  # TODO: use UUID?
     )
 
     # Delete verification entry
-    db_cursor.execute(
+    await app.state.pool.execute(
         """
         DELETE FROM galtinn_verification
-        WHERE discord_id = %s
+        WHERE discord_id = $1
         """,
-        (discord_id,),
+        discord_id,
     )
 
     return RedirectResponse(f"/success/{user['preferred_username']}", status_code=303)