From ae3d1f5dd055d1ad37a7cdefd54013954cbdca1d Mon Sep 17 00:00:00 2001 From: "Nicklas H." Date: Wed, 24 Apr 2024 11:13:00 +0200 Subject: [PATCH 1/2] Ignore new migration scripts --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index c391513..cbfd0b1 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ __pycache__ ff* *.log *.db -rebooted \ No newline at end of file +rebooted +alembic/versions/* \ No newline at end of file From 2b7bca638dfb5de86c27768e561ea323ccfc84ba Mon Sep 17 00:00:00 2001 From: "Nicklas H." Date: Wed, 24 Apr 2024 11:14:34 +0200 Subject: [PATCH 2/2] Moved reboot scheduling and checking to the database --- QuantumKat.py | 38 ++++++++++------------------- cogs/Entanglement.py | 15 +++++++----- sql/crud.py | 58 ++++++++++++++++++++++++++++++++++++++++++++ sql/models.py | 10 +++++++- sql/schemas.py | 13 ++++++++++ 5 files changed, 102 insertions(+), 32 deletions(-) diff --git a/QuantumKat.py b/QuantumKat.py index b1dd689..f70ea31 100644 --- a/QuantumKat.py +++ b/QuantumKat.py @@ -1,11 +1,10 @@ from asyncio import run -from datetime import datetime +from datetime import datetime, timedelta from os import environ from random import choice, randint from sys import exit from helpers import LogHelper, MiscHelper, DiscordHelper -from pathlib import Path import discord from discord.ext import commands from dotenv import load_dotenv @@ -208,30 +207,19 @@ async def on_ready(): except Exception as e: print(e) # Check if the bot was rebooted and edit the message to indicate it was successful - # TODO: Use the database instead? Also, it edits the message even if the bot was not rebooted. - # Maybe add a timestamp to the database and check if the bot was rebooted within the last x minutes? - if Path("rebooted").exists(): - with open("rebooted", "r") as f: - IDs = f.read() - # Order: message ID, channel ID, guild ID - IDs = IDs.split("\n") - try: - # Try and get the message from cache first - channel = bot.get_channel(int(IDs[1])) - if channel is None: - channel = await bot.fetch_channel(int(IDs[1])) - message = await channel.fetch_message(int(IDs[0])) - except Exception: - logger.error("Error fetching message to edit", exc_info=True) - if message: - try: - await message.edit(content=f"{message.content} Rebooted successfully!") - except Exception: - logger.error("Error editing message", exc_info=True) + reboot = await crud.get_reboot_status(AsyncSessionLocal) + if reboot and reboot.is_reboot_scheduled: + msg_id, channel_id, guild_id = reboot.message_location + channel = bot.get_channel(channel_id) + if channel is None: + channel = await bot.fetch_channel(channel_id) + message: discord.Message = await channel.fetch_message(msg_id) + # if reboot was more than 5 minutes ago, assume it was not successful + if reboot.reboot_time < datetime.now() - timedelta(minutes=5): + await message.edit(content=f"{message.content} Reboot was not successful.") else: - # if the message is not found, instead send a message to the bot owner - await bot.get_user(int(OWNER_ID)).send("Rebooted successfully!") - Path("rebooted").unlink() + await message.edit(content=f"{message.content} Rebooted successfully!") + await crud.unset_reboot(AsyncSessionLocal) bot.appinfo = await bot.application_info() quantum = ["reality", "universe", "dimension", "timeline"] diff --git a/cogs/Entanglement.py b/cogs/Entanglement.py index bb4d226..1306881 100644 --- a/cogs/Entanglement.py +++ b/cogs/Entanglement.py @@ -31,7 +31,7 @@ from helpers import LogHelper from sql.database import AsyncSessionLocal -from sql import crud +from sql import crud, schemas class Entanglements(commands.Cog): @@ -1145,11 +1145,14 @@ async def reboot(self, ctx: commands.Context): self.bot.reboot_scheduled = True msg = await ctx.send("Shutting down extensions and rebooting...") await asyncsleep(10) - with open("rebooted", "w") as f: - if self.bot.discord_helper.is_dm(ctx): - f.write(f"{msg.id}\n{msg.channel.id}\nNone") - else: - f.write(f"{msg.id}\n{msg.channel.id}\n{msg.guild.id}") + await crud.set_reboot( + AsyncSessionLocal, + schemas.Bot.SetReboot( + is_reboot_scheduled=True, + reboot_time=datetime.now(), + message_location=(msg.id, msg.channel.id, msg.guild.id if msg.guild else 0), + ), + ) for cog in listdir("./cogs"): if cog.endswith(".py"): try: diff --git a/sql/crud.py b/sql/crud.py index 36617c8..81b75e0 100644 --- a/sql/crud.py +++ b/sql/crud.py @@ -432,3 +432,61 @@ async def edit_server_ban(db: AsyncSession, server: schemas.Server.SetBan): result = result.scalar_one_or_none() result.is_banned = server.is_banned await db.commit() + + +@timeit +async def set_reboot(db: AsyncSession, bot: schemas.Bot.SetReboot): + """ + Set the reboot status for the bot. + + Args: + db (AsyncSession): The database session. + bot (schemas.Bot.SetReboot): The bot object with the new reboot status. + + Returns: + None + """ + async with db() as db: + result = await db.execute(select(models.Bot)) + result = result.scalar_one_or_none() + if result is None: + db.add(models.Bot(**bot.model_dump())) + else: + result.is_reboot_scheduled = bot.is_reboot_scheduled + result.reboot_time = bot.reboot_time + result.message_location = bot.message_location + await db.commit() + + +@timeit +async def unset_reboot(db: AsyncSession): + """ + Unset the reboot status for the bot. + + Args: + db (AsyncSession): The database session. + + Returns: + None + """ + async with db() as db: + result = await db.execute(select(models.Bot)) + result = result.scalar_one_or_none() + await db.delete(result) + await db.commit() + + +@timeit +async def get_reboot_status(db: AsyncSession): + """ + Get the reboot status for the bot. + + Args: + db (AsyncSession): The database session. + + Returns: + models.Bot: The bot object with the reboot status. + """ + async with db() as db: + result = await db.execute(select(models.Bot)) + return result.scalar_one_or_none() diff --git a/sql/models.py b/sql/models.py index 687b6d7..ee2f3e6 100644 --- a/sql/models.py +++ b/sql/models.py @@ -1,4 +1,4 @@ -from sqlalchemy import Column, Integer, String, ForeignKey +from sqlalchemy import Column, Integer, String, ForeignKey, DateTime, PickleType from sqlalchemy.orm import relationship from .database import Base @@ -42,6 +42,14 @@ class Chat(Base): server = relationship("Server", back_populates="chats") +class Bot(Base): + __tablename__ = "bot" + + is_reboot_scheduled = Column(Integer, primary_key=True, nullable=False, default=0) + reboot_time = Column(DateTime, nullable=False, default=0) + message_location = Column(PickleType, nullable=False, default=()) + + class AlembicVersion(Base): __tablename__ = "alembic_version" diff --git a/sql/schemas.py b/sql/schemas.py index d7e3f87..bf7ab53 100644 --- a/sql/schemas.py +++ b/sql/schemas.py @@ -1,5 +1,6 @@ from pydantic import BaseModel, Field from typing import Optional +from datetime import datetime class User: @@ -70,3 +71,15 @@ class Add(_Base): class Delete(Get): pass + + +class Bot: + class _Base(BaseModel): + is_reboot_scheduled: bool + reboot_time: datetime + + class unsetReboot(_Base): + pass + + class SetReboot(unsetReboot): + message_location: tuple[int, int, int]