Migrated cron collection to use ODM

bot.py

@@ -20,6 +20,7 @@
 from discord.ext import commands
 from rich.logging import RichHandler
 
+import src.mongo.models
 from env import env
 from src.discord.globals import (
     CHANNEL_BOTSPAM,
@@ -164,6 +165,7 @@ async def setup_hook(self) -> None:
         await init_beanie(
             database=self.mongo_database.client["data"],
             document_models=[
+                src.mongo.models.Cron,
                 # TODO
             ],
         )

src/discord/staffcommands.py

@@ -39,6 +39,7 @@
     ROLE_WM,
 )
 from src.discord.invitationals import update_invitational_list
+from src.mongo.models import Cron
 from src.wiki.mosteditstable import run_table
 
 if TYPE_CHECKING:
@@ -887,7 +888,7 @@ async def cron(self, interaction: discord.Interaction):
         """
         commandchecks.is_staff_from_ctx(interaction)
 
-        cron_list = await self.bot.mongo_database.get_cron()
+        cron_list = await Cron.find_all().to_list()
         if not len(cron_list):
             return await interaction.response.send_message(
                 "Unfortunately, there are no items in the CRON list to manage.",

src/discord/tasks.py

@@ -13,6 +13,7 @@
 from env import env
 from src.discord.invitationals import update_invitational_list
 from src.discord.views import UnselfmuteView
+from src.mongo.models import Cron
 
 if TYPE_CHECKING:
     from bot import PiBot
@@ -204,19 +205,19 @@ async def cron(self) -> None:
         """
         logger.debug("Executing CRON...")
         # Get the relevant tasks
-        cron_list = await self.bot.mongo_database.get_cron()
+        cron_list = await Cron.find_all().to_list()
 
         for task in cron_list:
             # If the date has passed, execute task
-            if discord.utils.utcnow() > task["time"]:
+            if discord.utils.utcnow() > task.time:
                 try:
-                    if task["type"] == "UNBAN":
+                    if task.cron_type == "UNBAN":
                         await self.cron_handle_unban(task)
-                    elif task["type"] == "UNMUTE":
+                    elif task.cron_type == "UNMUTE":
                         await self.cron_handle_unmute(task)
-                    elif task["type"] == "UNSELFMUTE":
+                    elif task.cron_type == "UNSELFMUTE":
                         await self.cron_handle_unselfmute(task)
-                    elif task["type"] == "REMOVE_STATUS":
+                    elif task.cron_type == "REMOVE_STATUS":
                         await self.cron_handle_remove_status(task)
                     else:
                         logger.error("ERROR:")
@@ -227,7 +228,7 @@ async def cron(self) -> None:
                     reporter_cog: commands.Cog | Reporter = self.bot.get_cog("Reporter")
                     await reporter_cog.create_cron_task_report(task)
 
-    async def cron_handle_unban(self, task: dict):
+    async def cron_handle_unban(self, task: Cron):
         """
         Handles serving CRON tasks with the type of 'UNBAN'.
         """
@@ -239,7 +240,7 @@ async def cron_handle_unban(self, task: dict):
         assert isinstance(server, discord.Guild)
 
         # Attempt to unban user
-        member = await self.bot.fetch_user(task["user"])
+        member = await self.bot.fetch_user(task.user)
         if member in server.members:
             # User is still in server, thus already unbanned
             await reporter_cog.create_cron_unban_auto_notice(member, is_present=True)
@@ -258,9 +259,9 @@ async def cron_handle_unban(self, task: dict):
             )
 
         # Remove cron task.
-        await self.delete_from_cron(task["_id"])
+        await task.delete()
 
-    async def cron_handle_unmute(self, task: dict):
+    async def cron_handle_unmute(self, task: Cron):
         """
         Handles serving CRON tasks with the type of 'UNMUTE'.
         """
@@ -277,7 +278,7 @@ async def cron_handle_unmute(self, task: dict):
         assert isinstance(muted_role, discord.Role)
 
         # Attempt to unmute user
-        member = server.get_member(task["user"])
+        member = server.get_member(task.user)
         if member in server.members:
             # User is still in server, thus can be unmuted
             await member.remove_roles(muted_role)
@@ -287,9 +288,9 @@ async def cron_handle_unmute(self, task: dict):
             await reporter_cog.create_cron_unmute_auto_notice(member, is_present=False)
 
         # Remove cron task.
-        await self.delete_from_cron(task["_id"])
+        await task.delete()
 
-    async def cron_handle_unselfmute(self, task: dict):
+    async def cron_handle_unselfmute(self, task: Cron):
         """
         Handles serving CRON tasks with the type of 'UNSELFMUTE'.
         """
@@ -306,18 +307,19 @@ async def cron_handle_unselfmute(self, task: dict):
         assert isinstance(muted_role, discord.Role)
 
         # Attempt to unmute user
-        member = server.get_member(task["user"])
+        member = server.get_member(task.user)
         if member in server.members:
             # User is still in server, thus can be unmuted
             await member.remove_roles(muted_role)
 
         # Remove cron task.
-        await self.delete_from_cron(task["_id"])
+        await task.delete()
 
-    async def cron_handle_remove_status(self, task: dict):
+    async def cron_handle_remove_status(self, task: Cron):
         """
         Handles serving CRON tasks with the type of 'REMOVE_STATUS'.
         """
+        # FIXME: Subject to premature removal of status if two /status commands are run with different expirations
         # Attempt to remove status
         self.bot.settings["custom_bot_status_type"] = None  # reset local settings
         self.bot.settings["custom_bot_status_text"] = None  # reset local settings
@@ -330,7 +332,7 @@ async def cron_handle_remove_status(self, task: dict):
         self.change_bot_status.restart()  # update bot now
 
         # Remove cron task.
-        await self.delete_from_cron(task["_id"])
+        await task.delete()
 
     @tasks.loop(hours=1)
     async def change_bot_status(self):

src/discord/views.py

@@ -5,6 +5,7 @@
 import discord
 
 from src.discord.globals import ROLE_SELFMUTE
+from src.mongo.models import Cron
 
 if TYPE_CHECKING:
     from bot import PiBot
@@ -47,10 +48,10 @@ async def unselfmute_button(
         try:
             item = next(
                 x
-                for x in await self.bot.mongo_database.get_cron()
-                if (x["type"] == "UNSELFMUTE" and x["user"] == interaction.user.id)
+                for x in await Cron.find_all().to_list()
+                if (x.cron_type == "UNSELFMUTE" and x.user == interaction.user.id)
             )
-            await self.bot.mongo_database.remove_doc("data", "cron", item["_id"])
+            await item.delete()
         except Exception:  # not in the database - maybe was removed!
             pass
 

src/mongo/models.py

@@ -1,3 +1,21 @@
 """
 Contains all database models
 """
+from datetime import datetime
+from typing import Annotated, Literal
+
+from beanie import Document, Indexed
+from pydantic import Field
+
+
+class Cron(Document):
+    cron_type: Literal["UNMUTE", "UNBAN", "UNSELFMUTE", "REMOVE_STATUS"] = Field(
+        alias="type",
+    )
+    time: Annotated[datetime, Indexed()]
+    user: Annotated[int, Indexed()]
+    tag: str
+
+    class Settings:
+        name = "cron"
+        use_cache = False

src/mongo/mongo.py

@@ -89,12 +89,6 @@ async def get_invitationals(self):
         """
         return await self.get_entire_collection("data", "invitationals")
 
-    async def get_cron(self):
-        """
-        Gets all documents in the CRON collection.
-        """
-        return await self.get_entire_collection("data", "cron")
-
     async def get_censor(self):
         """
         Gets the document containing censor information from the censor collection.