from motor.motor_asyncio import AsyncIOMotorClient from pymongo import MongoClient import os import logging # Global Motor client singleton - per event loop _motor_clients = {} # event_loop_id -> (client, database) async def get_db(): """Get database connection using singleton Motor client per event loop.""" import asyncio # Get current event loop to ensure Motor client affinity try: current_loop = asyncio.get_running_loop() loop_id = id(current_loop) except RuntimeError: raise RuntimeError("get_db() must be called from within an async context") # Return cached database for this event loop if available if loop_id in _motor_clients: client, database = _motor_clients[loop_id] return database # Try to read environment variables for MongoDB credentials mongo_user = os.environ.get('MONGO_USER') mongo_pass = os.environ.get('MONGO_PASS') mongo_host = os.environ.get('MONGO_HOST', 'localhost') mongo_port = os.environ.get('MONGO_PORT', '27017') # Try with standard credentials first standard_credentials = [ {"user": "admin", "pass": "admin", "db": "admin"}, {"user": "mongodb", "pass": "mongodb", "db": "admin"}, {"user": "root", "pass": "root", "db": "admin"}, {"user": "user", "pass": "pass", "db": "admin"} ] # Try each set of standard credentials for creds in standard_credentials: try: uri = f"mongodb://{creds['user']}:{creds['pass']}@{mongo_host}:{mongo_port}/semblance_db?authSource={creds['db']}" motor_client = AsyncIOMotorClient(uri, serverSelectionTimeoutMS=2000) database = motor_client.semblance_db # Test the connection with a simple command await database.command('ping') logging.debug(f"Successfully connected to MongoDB with standard credentials ({creds['user']})") # Cache for this event loop _motor_clients[loop_id] = (motor_client, database) return database except Exception as e: # Continue trying other credentials pass # Try to connect without authentication if standard credentials don't work try: motor_client = AsyncIOMotorClient(f'mongodb://{mongo_host}:{mongo_port}', serverSelectionTimeoutMS=5000) database = motor_client.semblance_db # Test the connection with a simple command await database.command('ping') # Try a write operation to verify we have proper access test_result = await database.test_collection.insert_one({"test": "auth_test"}) await database.test_collection.delete_one({"_id": test_result.inserted_id}) logging.debug("Successfully connected to MongoDB without authentication") # Cache for this event loop _motor_clients[loop_id] = (motor_client, database) return database except Exception as e: logging.debug(f"Could not connect without auth: {e}") # If we get here, we need authentication - try with environment vars if provided if mongo_user and mongo_pass: try: uri = f"mongodb://{mongo_user}:{mongo_pass}@{mongo_host}:{mongo_port}/semblance_db?authSource=admin" motor_client = AsyncIOMotorClient(uri, serverSelectionTimeoutMS=5000) database = motor_client.semblance_db await database.command('ping') # Test the connection logging.debug(f"Successfully connected to MongoDB with credentials for user: {mongo_user}") # Cache for this event loop _motor_clients[loop_id] = (motor_client, database) return database except Exception as e: logging.warning(f"Failed to connect with environment credentials: {e}") # Last resort - log warning and return client that will fail later if DB actually needs auth logging.warning("Could not authenticate with MongoDB. If authentication is required, operations will fail.") logging.warning("To fix this: Set MONGO_USER and MONGO_PASS environment variables.") # Return a client that will likely fail when operations are performed, but the app will start motor_client = AsyncIOMotorClient(f'mongodb://{mongo_host}:{mongo_port}', serverSelectionTimeoutMS=5000) database = motor_client.semblance_db # Cache for this event loop _motor_clients[loop_id] = (motor_client, database) return database def close_db_connections(): """Close all Motor clients and their PyMongo background threads.""" global _motor_clients closed_count = 0 for loop_id, (client, database) in _motor_clients.items(): try: client.close() closed_count += 1 except Exception as e: logging.warning(f"Error closing Motor client for loop {loop_id}: {e}") if closed_count > 0: logging.info(f"🗄️ Closed {closed_count} Motor clients - PyMongo threads should stop") _motor_clients.clear()