Bot will now exit after the migration

classes/pycord_bot.py

@@ -27,7 +27,7 @@ logger: Logger = get_logger(__name__)
 
 
 class PycordBot(LibPycordBot):
-    __version__ = "1.0.0"
+    __version__ = "1.0.1"
 
     started: datetime
     cache: CacheMemcached | CacheRedis | None = None

classes/pycord_event.py

@@ -11,7 +11,6 @@ from libbot.cache.classes import Cache
 from pymongo import DESCENDING
 from pymongo.results import InsertOneResult
 
-from classes.abstract import Cacheable
 from classes.base.base_cacheable import BaseCacheable
 from classes.errors import EventNotFoundError
 from modules.database import col_events
@@ -86,7 +85,7 @@ class PycordEvent(BaseCacheable):
         if cached_entry is not None:
             return cls(**cls._entry_from_cache(cached_entry))
 
-        db_entry = await cls.__collection__.find_one(
+        db_entry: Dict[str, Any] | None = await cls.__collection__.find_one(
             {"_id": event_id if isinstance(event_id, ObjectId) else ObjectId(event_id)}
         )
 

classes/pycord_event_stage.py

@@ -64,7 +64,7 @@ class PycordEventStage(BaseCacheable):
         if cached_entry is not None:
             return cls(**cls._entry_from_cache(cached_entry))
 
-        db_entry = await cls.__collection__.find_one(
+        db_entry: Dict[str, Any] | None = await cls.__collection__.find_one(
             {"_id": stage_id if isinstance(stage_id, ObjectId) else ObjectId(stage_id)}
         )
 
@@ -155,7 +155,11 @@ class PycordEventStage(BaseCacheable):
         """
         return {
             "_id": self._id if not json_compatible else str(self._id),
-            "event_id": self.event_id if not json_compatible else str(self.event_id),
+            "event_id": (
+                self.event_id
+                if not json_compatible
+                else (None if self.event_id is None else str(self.event_id))
+            ),
             "guild_id": self.guild_id,
             "sequence": self.sequence,
             "created": self.created if not json_compatible else self.created.isoformat(),

classes/pycord_guild.py

@@ -60,7 +60,7 @@ class PycordGuild(BaseCacheable):
         if cached_entry is not None:
             return cls(**cls._entry_from_cache(cached_entry))
 
-        db_entry = await cls.__collection__.find_one({"id": guild_id})
+        db_entry: Dict[str, Any] | None = await cls.__collection__.find_one({"id": guild_id})
 
         if db_entry is None:
             if not allow_creation:

classes/pycord_user.py

@@ -86,7 +86,9 @@ class PycordUser(BaseCacheable):
         if cached_entry is not None:
             return cls(**cls._entry_from_cache(cached_entry))
 
-        db_entry = await cls.__collection__.find_one({"id": user_id, "guild_id": guild_id})
+        db_entry: Dict[str, Any] | None = await cls.__collection__.find_one(
+            {"id": user_id, "guild_id": guild_id}
+        )
 
         if db_entry is None:
             if not allow_creation:
@@ -119,10 +121,14 @@ class PycordUser(BaseCacheable):
             "event_channels": self.event_channels,
             "is_jailed": self.is_jailed,
             "current_event_id": (
-                self.current_event_id if not json_compatible else str(self.current_event_id)
+                self.current_event_id
+                if not json_compatible
+                else (None if self.current_event_id is None else str(self.current_event_id))
             ),
             "current_stage_id": (
-                self.current_stage_id if not json_compatible else str(self.current_stage_id)
+                self.current_stage_id
+                if not json_compatible
+                else (None if self.current_stage_id is None else str(self.current_stage_id))
             ),
             "registered_event_ids": (
                 self.registered_event_ids

main.py

@@ -45,6 +45,7 @@ def main() -> None:
     # Perform migration if command line argument was provided
     if args.migrate:
         migrate_database()
+        exit()
 
     # if args.downgrade:
     #     if not args.confirm:

migrations/202505192040.py

@@ -0,0 +1,59 @@
+from typing import Dict, Any
+
+from mongodb_migrations.base import BaseMigration
+
+
+class Migration(BaseMigration):
+    def upgrade(self):
+        index_information_users: Dict[str, Any] = self.db.users.index_information()
+        index_information_events: Dict[str, Any] = self.db.events.index_information()
+        index_information_stages: Dict[str, Any] = self.db.stages.index_information()
+
+        # Update users collection
+        if "user_id" in index_information_users:
+            self.db.users.drop_index("user_id")
+
+        if "user_id-guild_id" not in index_information_users:
+            self.db.users.create_index(["id", "guild_id"], name="user_id-guild_id", unique=True)
+
+        # Update events collection
+        if "guild_id" in index_information_events:
+            self.db.events.drop_index("guild_id")
+
+        if "event_name-guild_id" not in index_information_events:
+            self.db.events.create_index(["name", "guild_id"], name="event_name-guild_id", unique=False)
+
+        # Update stages collection
+        if "event_id-and-guild_id" in index_information_stages:
+            self.db.stages.drop_index("event_id-and-guild_id")
+
+        if "event_id-guild_id" not in index_information_stages:
+            self.db.stages.create_index(["event_id", "guild_id"], name="event_id-guild_id", unique=False)
+
+    def downgrade(self):
+        index_information_users: Dict[str, Any] = self.db.users.index_information()
+        index_information_events: Dict[str, Any] = self.db.events.index_information()
+        index_information_stages: Dict[str, Any] = self.db.stages.index_information()
+
+        # Update users collection
+        if "user_id-guild_id" in index_information_users:
+            self.db.users.drop_index("user_id-guild_id")
+
+        if "user_id" not in index_information_users:
+            self.db.users.create_index("id", name="user_id", unique=True)
+
+        # Update events collection
+        if "event_name-guild_id" in index_information_events:
+            self.db.events.drop_index("event_name-guild_id")
+
+        if "guild_id" not in index_information_events:
+            self.db.events.create_index("guild_id", name="guild_id", unique=False)
+
+        # Update stages collection
+        if "event_id-guild_id" in index_information_stages:
+            self.db.stages.drop_index("event_id-guild_id")
+
+        if "event_id-and-guild_id" not in index_information_stages:
+            self.db.stages.create_index(
+                ["event_id", "guild_id"], name="event_id-and-guild_id", unique=False
+            )

modules/database.py

@@ -32,7 +32,7 @@ col_stages: AsyncCollection = db.get_collection("stages")
 
 # Update indexes
 async def _update_database_indexes() -> None:
-    await col_users.create_index("id", name="user_id", unique=True)
+    await col_users.create_index(["id", "guild_id"], name="user_id-guild_id", unique=True)
     await col_guilds.create_index("id", name="guild_id", unique=True)
-    await col_events.create_index("guild_id", name="guild_id", unique=False)
-    await col_stages.create_index(["event_id", "guild_id"], name="event_id-and-guild_id", unique=False)
+    await col_events.create_index(["name", "guild_id"], name="event_name-guild_id", unique=False)
+    await col_stages.create_index(["event_id", "guild_id"], name="event_id-guild_id", unique=False)