115 lines
No EOL
5 KiB
Python
Executable file
115 lines
No EOL
5 KiB
Python
Executable file
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() |