commit 3961e7ffe852ab629c329bf5d5ee05252eb3a4de
parent c15ffe9b4ce2879587d5524014d495f99b2445a8
Author: Asher Morgan <59518073+ashermorgan@users.noreply.github.com>
Date: Wed, 24 Apr 2024 18:26:43 -0700
Update activate, deactivate, and reload commands
Diffstat:
5 files changed, 76 insertions(+), 45 deletions(-)
diff --git a/countdown_bot/bot.py b/countdown_bot/bot.py
@@ -29,8 +29,8 @@ class CountdownBot(commands.Bot):
async def setup_hook(self):
+ await self.add_cog(utilitiesCog.Utilities(self, self.databaseSessionMaker, self.db_connection))
await self.add_cog(analyticsCog.Analytics(self, self.databaseSessionMaker))
- await self.add_cog(utilitiesCog.Utilities(self, self.databaseSessionMaker))
@@ -80,6 +80,9 @@ class CountdownBot(commands.Bot):
async def on_command_error(self, ctx, error):
+ # Rollback database transaction
+ self.db_connection.rollback()
+
# Send error embed
embed=discord.Embed(title=":warning: Error", description=str(error), color=COLORS["error"])
if (isinstance(error, commands.CommandNotFound)):
diff --git a/countdown_bot/botUtilities.py b/countdown_bot/botUtilities.py
@@ -166,7 +166,7 @@ def isCountdown(cur, id):
"""
cur.execute("CALL isCountdown(%s, null);",
- (ctx.channel.id,))
+ (id,))
return cur.fetchone()[0]
@@ -324,17 +324,17 @@ async def loadCountdown(bot, countdown):
The bot to load messages with
cur : psycopg.cursor
The database cursor
- countdown : Countdown
- The countdown to load messages for
+ countdown : int
+ The ID of the countdown to load messages for
"""
with bot.db_connection.cursor() as cur:
# Clear countdown
- cur.execute("CALL clearCountdown(%s);", (countdown.id,))
+ cur.execute("CALL clearCountdown(%s);", (countdown,))
# Get Discord messages
messages = [message async for message in
- bot.get_channel(countdown.id).history(limit=10100)]
+ bot.get_channel(countdown).history(limit=10100)]
messages.reverse()
# Add messages to countdown
diff --git a/countdown_bot/utilitiesCog.py b/countdown_bot/utilitiesCog.py
@@ -3,15 +3,16 @@ import discord
from discord.ext import commands
# Import modules
-from .botUtilities import COLORS, CommandError, getContextCountdown, getCountdown, loadCountdown
+from .botUtilities import COLORS, CommandError, isCountdown, getContextCountdown, getCountdown, loadCountdown
from .models import Countdown, Prefix, Reaction
class Utilities(commands.Cog):
- def __init__(self, bot, databaseSessionMaker):
+ def __init__(self, bot, databaseSessionMaker, db_connection):
self.bot = bot
self.databaseSessionMaker = databaseSessionMaker
+ self.db_connection = db_connection
self.bot.remove_command("help")
@@ -22,9 +23,9 @@ class Utilities(commands.Cog):
Turns a channel into a countdown
"""
- with self.databaseSessionMaker() as session:
+ with self.db_connection.cursor() as cur:
# Check if channel is already a countdown
- if (getCountdown(session, ctx.channel.id)):
+ if (isCountdown(cur, ctx.channel.id)):
raise CommandError("This channel is already a countdown")
# Check if channel is a DM
@@ -36,14 +37,8 @@ class Utilities(commands.Cog):
raise CommandError("You must be an administrator to turn a channel into a countdown")
# Create countdown
- countdown = Countdown(
- id = ctx.channel.id,
- server_id = ctx.channel.guild.id,
- timezone = 0,
- prefixes = [Prefix(countdown_id=ctx.channel.id, value=x) for x in self.bot.prefixes],
- reactions = [],
- messages = [],
- )
+ cur.execute("CALL createCountdown(%s, %s, %s);",
+ (ctx.channel.id, ctx.channel.guild.id, self.bot.prefixes[0]))
# Send initial response
self.bot.logger.info(f"Activated {self.bot.get_channel(ctx.channel.id)} (ID {ctx.channel.id}) as a countdown")
@@ -51,9 +46,8 @@ class Utilities(commands.Cog):
msg = await ctx.send(embed=embed)
# Load countdown
- await loadCountdown(self.bot, countdown)
- session.add(countdown)
- session.commit()
+ await loadCountdown(self.bot, ctx.channel.id)
+ self.db_connection.commit()
# Send final response
embed = discord.Embed(title=":white_check_mark: Countdown Activated", description="This channel is now a countdown\nYou may start counting!", color=COLORS["embed"])
@@ -136,10 +130,9 @@ class Utilities(commands.Cog):
Deactivates a countdown channel
"""
- with self.databaseSessionMaker() as session:
+ with self.db_connection.cursor() as cur:
# Check if channel isn't a countdown
- countdown = getCountdown(session, ctx.channel.id)
- if (not countdown):
+ if (not isCountdown(cur, ctx.channel.id)):
raise CommandError("This channel isn't a countdown")
# Check if user isn't authorized
@@ -147,8 +140,9 @@ class Utilities(commands.Cog):
raise CommandError("You must be an administrator to deactivate a countdown channel")
# Delete countdown
- session.delete(countdown)
- session.commit()
+ cur.execute("CALL deleteCountdown(%s);",
+ (ctx.channel.id,))
+ self.db_connection.commit()
# Send response
self.bot.logger.info(f"Deactivated {self.bot.get_channel(ctx.channel.id)} (ID {ctx.channel.id}) as a countdown")
@@ -402,20 +396,20 @@ class Utilities(commands.Cog):
Reloads the countdown cache
"""
- with self.databaseSessionMaker() as session:
- countdown = getCountdown(session, ctx.channel.id)
- if (countdown):
- # Send initial response
- embed = discord.Embed(title=":clock3: Reloading Countdown Cache", description="Please wait to continue counting", color=COLORS["embed"])
- msg = await ctx.channel.send(embed=embed)
-
- # Reload messages
- await loadCountdown(self.bot, countdown)
- session.commit()
-
- # Send final response
- self.bot.logger.info(f"Reloaded messages from {self.bot.get_channel(ctx.channel.id)} (ID {ctx.channel.id})")
- embed = discord.Embed(title=":white_check_mark: Countdown Cache Reloaded", description="Done! You may continue counting!", color=COLORS["embed"])
- await msg.edit(embed=embed)
- else:
+ with self.db_connection.cursor() as cur:
+ # Check if channel isn't a countdown
+ if (not isCountdown(cur, ctx.channel.id)):
raise CommandError("Countdown not found\nThis command must be used in a countdown channel")
+
+ # Send initial response
+ embed = discord.Embed(title=":clock3: Reloading Countdown Cache", description="Please wait to continue counting", color=COLORS["embed"])
+ msg = await ctx.channel.send(embed=embed)
+
+ # Reload messages
+ await loadCountdown(self.bot, ctx.channel.id)
+ self.db_connection.commit()
+
+ # Send final response
+ self.bot.logger.info(f"Reloaded messages from {self.bot.get_channel(ctx.channel.id)} (ID {ctx.channel.id})")
+ embed = discord.Embed(title=":white_check_mark: Countdown Cache Reloaded", description="Done! You may continue counting!", color=COLORS["embed"])
+ await msg.edit(embed=embed)
diff --git a/models/tables.sql b/models/tables.sql
@@ -18,7 +18,9 @@ CREATE TABLE messages (
userID BIGINT NOT NULL, -- The author's Discord user ID
value INT NOT NULL, -- The message's numeric value
timestamp TIMESTAMPTZ NOT NULL, -- The message timestamp
- FOREIGN KEY (countdownID) REFERENCES countdowns(countdownID)
+ FOREIGN KEY (countdownID)
+ REFERENCES countdowns(countdownID)
+ ON DELETE CASCADE
);
-- Records bot command prefixes
@@ -26,7 +28,9 @@ CREATE table prefixes (
prefixID SERIAL PRIMARY KEY, -- The prefix ID
countdownID BIGINT NOT NULL, -- The countdown ID
value VARCHAR(8) NOT NULL, -- The prefix
- FOREIGN KEY (countdownID) REFERENCES countdowns(countdownID)
+ FOREIGN KEY (countdownID)
+ REFERENCES countdowns(countdownID)
+ ON DELETE CASCADE
);
-- Records custom countdown reactions
@@ -35,5 +39,7 @@ CREATE table reactions (
countdownID BIGINT NOT NULL, -- The countdown ID
number INT NOT NULL, -- The prefix
value CHAR NOT NULL, -- The reaction
- FOREIGN KEY (countdownID) REFERENCES countdowns(countdownID)
+ FOREIGN KEY (countdownID)
+ REFERENCES countdowns(countdownID)
+ ON DELETE CASCADE
);
diff --git a/models/utilities.sql b/models/utilities.sql
@@ -3,6 +3,8 @@
DROP FUNCTION IF EXISTS getReactions;
DROP PROCEDURE IF EXISTS addMessage;
DROP TYPE IF EXISTS addMessageResults;
+DROP PROCEDURE IF EXISTS deleteCountdown;
+DROP PROCEDURE IF EXISTS createCountdown;
DROP PROCEDURE IF EXISTS clearCountdown;
DROP PROCEDURE IF EXISTS isCountdown;
DROP PROCEDURE IF EXISTS getUserContextCountdown;
@@ -118,6 +120,32 @@ BEGIN
END
$$;
+-- Create a new countdown
+CREATE PROCEDURE createCountdown (
+ _countdownID IN BIGINT, -- The countdown channel ID
+ _serverID IN BIGINT, -- The server ID
+ prefix IN VARCHAR(8) -- The initial prefix
+)
+LANGUAGE plpgsql AS $$
+BEGIN
+ INSERT INTO countdowns (countdownID, serverID)
+ VALUES (_countdownID, _serverID);
+ INSERT INTO prefixes (countdownID, value)
+ VALUES (_countdownID, prefix);
+END
+$$;
+
+-- Delete a countdown
+CREATE PROCEDURE deleteCountdown (
+ _countdownID IN BIGINT -- The countdown channel ID
+)
+LANGUAGE plpgsql AS $$
+BEGIN
+ DELETE
+ FROM countdowns
+ WHERE countdownID = _countdownID;
+END
+$$;
-- Possible results of the addMessage procedure
CREATE TYPE addMessageResults AS ENUM (
