Merge branch 'dev'

.gitignore

@@ -6,4 +6,5 @@ __pycache__
 ff*
 *.log
 *.db
-rebooted
+rebooted
+alembic/versions/*

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"]

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:

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()

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"
 

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]