impr: Code organization and bug fixes.

MusicBot/database/base.py

@@ -1,5 +1,5 @@
 from typing import Iterable, Any, cast
-from pymongo import AsyncMongoClient, ReturnDocument
+from pymongo import AsyncMongoClient, ReturnDocument, UpdateOne
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.results import UpdateResult
 
@@ -43,6 +43,19 @@ class BaseUsersDatabase:
             upsert=True,
             projection=projection
         )
+        ops = []
+
+        for key, value in self.DEFAULT_USER.items():
+            if key not in user and (projection is None or key in projection):
+                user[key] = value
+                ops.append(UpdateOne({'_id': uid}, {'$set': {key: value}}))
+        for key, value in user.copy().items():
+            if key not in self.DEFAULT_USER and key != '_id':
+                del user[key]
+                ops.append(UpdateOne({'_id': uid}, {'$unset': {key: ''}}))
+
+        if ops:
+            await users.bulk_write(ops)
         return cast(ExplicitUser, user)
 
     async def get_ym_token(self, uid: int) -> str | None:
@@ -68,6 +81,7 @@ class BaseGuildsDatabase:
         is_stopped=True,
         allow_explicit=True,
         always_allow_menu=False,
+        allow_connect=False,
         vote_next_track=True,
         vote_add_track=True,
         vote_add_album=True,
@@ -95,6 +109,19 @@ class BaseGuildsDatabase:
             upsert=True,
             projection=projection
         )
+        ops = []
+
+        for key, value in self.DEFAULT_GUILD.items():
+            if key not in guild and (projection is None or key in projection):
+                guild[key] = value
+                ops.append(UpdateOne({'_id': gid}, {'$set': {key: value}}))
+        for key, value in guild.copy().items():
+            if key not in self.DEFAULT_GUILD and key != '_id':
+                del guild[key]
+                ops.append(UpdateOne({'_id': gid}, {'$unset': {key: ''}}))
+    
+        if ops:
+            await guilds.bulk_write(ops)
         return cast(ExplicitGuild, guild)
 
     async def update_vote(self, gid: int, mid: int, data: MessageVotes) -> UpdateResult:

MusicBot/database/extensions.py

@@ -55,7 +55,7 @@ class VoiceGuildsDatabase(BaseGuildsDatabase):
         operations = {
             'insert': {'$push': {field: {'$each': track_data, '$position': 0}}},
             'append': {'$push': {field: {'$each': track_data}}},
-            'extend': {'$push': {field: {'$each': track_data}}},
+            'extend': {'$push': {field: {'$each': track_data}}},  # Same as append for consistency with python
             'pop_start': {'$pop': {field: -1}},
             'pop_end': {'$pop': {field: 1}}
         }
@@ -139,9 +139,7 @@ class VoiceGuildsDatabase(BaseGuildsDatabase):
 
         await guilds.update_one(
             {'_id': gid},
-            {
-                '$set': {'current_track': track}
-            }
+            {'$set': {'current_track': track}}
         )
 
     async def clear_tracks(self, gid: int, list_type: Literal['next', 'previous']) -> None:

MusicBot/database/guild.py

@@ -15,6 +15,7 @@ class Guild(TypedDict, total=False):
     is_stopped: bool
     allow_explicit: bool
     always_allow_menu: bool
+    allow_connect: bool
     vote_next_track: bool
     vote_add_track: bool
     vote_add_album: bool
@@ -35,6 +36,7 @@ class ExplicitGuild(TypedDict):
     is_stopped: bool  # Prevents the `after` callback of play_track
     allow_explicit: bool
     always_allow_menu: bool
+    allow_connect: bool
     vote_next_track: bool
     vote_add_track: bool
     vote_add_album: bool