diff --git a/cogs/commands/set.py b/cogs/commands/set.py index a9b55b3..0b5925d 100644 --- a/cogs/commands/set.py +++ b/cogs/commands/set.py @@ -6,7 +6,8 @@ from datetime import datetime import nextcord from nextcord.ext import commands -from sqlalchemy import delete +from sqlalchemy import delete, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from dakiny_bot import birthday_list @@ -45,22 +46,40 @@ class setCog(commands.Cog, name="Set Commands"): day (int, optional): The day a user was born. Is SlashOption. """ + # Must be deferred, apparently datetime takes a little bit on invalid dates. + await interaction.response.defer(ephemeral=True) + if month == 0 and day == 0: stmt = delete(UserBirthday).where(UserBirthday.UserId == interaction.user.id) with Session(engine) as session: - session.scalar(stmt) + session.execute(stmt) session.commit() + await interaction.followup.send("Removed birthday.", ephemeral=True) return # check if valid day today = datetime.today() - birthday = datetime(today.year, month, day) + try: + birthday = datetime(today.year, month, day) + except ValueError: + await interaction.followup.send(f"Invalid date: {month}/{day}.", ephemeral=True) + return # If the birthday is today, add to birthday_list if today.date() == birthday.date(): birthday_list.append(interaction.user.id) - new_birthday = UserBirthday(interaction.user.id, f"{month}/{day}") + new_birthday = UserBirthday(UserId=interaction.user.id, Birthday=f"{month}/{day}") with Session(engine) as session: - session.add(new_birthday) - session.commit() + try: + session.add(new_birthday) + session.commit() + except IntegrityError: + # User already has birthday, change it instead. + session.rollback() + stmt = select(UserBirthday).where(UserBirthday.UserId == interaction.user.id) + new_birthday = session.scalar(stmt) + new_birthday.Birthday = f"{month}/{day}" + session.commit() + + await interaction.followup.send(f"Birthday set to {month}/{day}.", ephemeral=True)