fix: Improve logic for removing temporary roles

bot/cogs/ext/temprole/time.py

@@ -1,7 +1,7 @@
 import re
 
 
-def parse_time_string(time_string):
+def parse_time_string(time_string: str) -> int:
     time_units = {'d': 86400, 'h': 3600, 'm': 60, 's': 1}
     pattern = re.compile(r'(\d+d|\d+h|\d+m|\d+s)')
     matches = pattern.findall(time_string)

bot/cogs/temporary.py

@@ -66,27 +66,13 @@ async def add_temporary_role(
         total_duration = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds)
 
         async with self.db_pool.acquire() as conn:
-            in_table = await conn.fetchval(
-                'SELECT 1 FROM temp_role WHERE user_id = $1 AND role_id = $2',
-                user.id,
-                role.id,
-            )
-
-        async with self.db_pool.acquire() as conn:
-            if in_table:
-                await conn.execute(
-                    'UPDATE temp_role SET end_time = $1 WHERE user_id = $2 AND role_id = $3',
-                    total_duration,
-                    user.id,
-                    role.id,
-                )
-            else:
-                await conn.execute(
-                    'INSERT INTO temp_role (user_id, role_id, end_time) VALUES ($1, $2, $3)',
-                    user.id,
-                    role.id,
-                    total_duration,
-                )
+            query = '''
+                INSERT INTO temp_role (user_id, role_id, end_time)
+                VALUES ($1, $2, $3)
+                ON CONFLICT (user_id, role_id)
+                DO UPDATE SET end_time = EXCLUDED.end_time
+            '''
+            await conn.execute(query, user.id, role.id, total_duration)
 
         embed = discord.Embed(
             title="Role Added",
@@ -99,11 +85,11 @@ async def add_temporary_role(
     async def _check_temprole(self) -> None:
         current_time = datetime.now(timezone.utc)
         guild = self.bot.get_guild(GUILD_ID)
-        user_success = []
+        id_success = []
 
         async with self.db_pool.acquire() as conn:
             records = await conn.fetch(
-                'SELECT user_id, role_id FROM temp_role WHERE end_time <= $1',
+                'SELECT id, user_id, role_id FROM temp_role WHERE end_time <= $1',
                 current_time,
             )
 
@@ -118,20 +104,20 @@ async def _check_temprole(self) -> None:
                 continue
 
             if user.get_role(role.id) is None or not role:
-                user_success.append(user.id)
+                id_success.append(record['id'])
                 continue
 
             try:
                 await user.remove_roles(role)
-                user_success.append(user.id)
+                id_success.append(record['id'])
             except discord.HTTPException:
                 logger.error(f'Failed to remove role {role.id} from user {user.id}')
 
         async with self.db_pool.acquire() as conn:
-            for user in user_success:
+            for id in id_success:
                 await conn.execute(
-                    'DELETE FROM temp_role WHERE user_id = $1',
-                    user,
+                    'DELETE FROM temp_role WHERE id = $1',
+                    id,
                 )
             logger.info(f'Removed role {role.id} from {user} users')
 

bot/data/db.sql

@@ -65,12 +65,12 @@ CREATE INDEX IF NOT EXISTS custom_role_role_id_idx ON custom_role (role_id);
 
 ----- TEMPORARY ROLE FEATURE ------
 -----------------------------------
-CREATE TABLE temp_role(
-	id SERIAL NOT NULL,
-	user_id BIGINT NOT NULL,
-	role_id BIGINT NOT NULL,
-	end_time TIMESTAMP WITH TIME ZONE NOT NULL,
-	PRIMARY KEY(id),
-	UNIQUE(id)
+CREATE TABLE temp_role (
+    id SERIAL NOT NULL,
+    user_id BIGINT NOT NULL,
+    role_id BIGINT NOT NULL,
+    end_time TIMESTAMP WITH TIME ZONE NOT NULL,
+    PRIMARY KEY(id),
+    UNIQUE(user_id, role_id)
 );
-CREATE INDEX IF NOT EXISTS temp_role_role_id_idx ON temp_role (id);
+CREATE INDEX IF NOT EXISTS temp_role_user_role_idx ON temp_role (user_id, role_id);