Update database/connection.py

database/connection.py

@@ -1,4 +1,4 @@
-# filename: database/connection.py
+# filename: database/connection.py (Corrected Version)
 
 import logging
 from motor.motor_asyncio import AsyncIOMotorClient
@@ -8,51 +8,41 @@ logger = logging.getLogger(__name__)
 
 class Database:
     """
-    Manages connections to one or more MongoDB databases.
+    Manages connections to one or more MongoDB clusters.
     """
-    def __init__(self, uris: list[str]):
+    def __init__(self, uris: list[str], db_name: str):
         if not uris:
             raise ValueError("At least one MongoDB URI is required.")
+        if not db_name:
+            raise ValueError("A database name is required.")
             
         self.clients = [AsyncIOMotorClient(uri) for uri in uris]
-        self.databases = [client.get_default_database() for client in self.clients]
         
-        # For data that should not be split (like cache), we use the first DB as the primary.
-        self.primary_db = self.databases[0]
+        # --- THE FIX IS HERE ---
+        # Instead of getting the "default" database, we explicitly get the one
+        # you defined in the config.
+        self.databases = [client[db_name] for client in self.clients]
         
-        # For distributing new users, we'll cycle through the available databases.
+        self.primary_db = self.databases[0]
         self._user_db_round_robin_counter = 0
         
-        logger.info(f"Successfully connected to {len(self.clients)} MongoDB database(s).")
+        logger.info(f"Successfully connected to {len(self.clients)} MongoDB cluster(s). Using database: '{db_name}'")
 
     async def get_user_db(self, user_id: int):
-        """
-        Gets the database assigned to a specific user.
-        This ensures a user's data always stays in the same database.
-        We use a simple hashing method to distribute users.
-        """
         db_index = user_id % len(self.databases)
         return self.databases[db_index]
         
     async def find_user_db(self, user_id: int):
-        """
-        Searches across all databases to find which one contains the user.
-        Returns the database object if found, otherwise None.
-        """
         for db in self.databases:
             if await db.users.find_one({"user_id": user_id}):
                 return db
         return None
 
     def get_next_db_for_new_user(self):
-        """
-        Picks a database for a new user using a round-robin strategy.
-        This helps distribute new users evenly.
-        """
         db = self.databases[self._user_db_round_robin_counter]
         self._user_db_round_robin_counter = (self._user_db_round_robin_counter + 1) % len(self.databases)
         return db
 
 # --- Initialize the Database Connection ---
-# This single `db` object will be imported and used by the rest of the application.
-db = Database(config.MONGODB_URIS)
+# ** We now pass the DATABASE_NAME from the config **
+db = Database(config.MONGODB_URIS, db_name=config.DATABASE_NAME)