135 lines
3.5 KiB
Python
135 lines
3.5 KiB
Python
"""
|
|
MongoDB database connection management.
|
|
|
|
This module handles MongoDB connection with fail-fast approach.
|
|
The application will terminate if MongoDB is not accessible at startup.
|
|
"""
|
|
|
|
import sys
|
|
from typing import Optional
|
|
from pymongo import MongoClient
|
|
from pymongo.database import Database
|
|
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
|
|
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
|
from app.config.settings import get_mongodb_url, get_mongodb_database_name
|
|
|
|
# Global variables for singleton pattern
|
|
_client: Optional[MongoClient] = None
|
|
_database: Optional[Database] = None
|
|
|
|
|
|
def create_mongodb_client() -> AsyncIOMotorClient:
|
|
"""
|
|
Create MongoDB client with connection validation.
|
|
|
|
Returns:
|
|
MongoClient: Connected MongoDB client
|
|
|
|
Raises:
|
|
SystemExit: If connection fails (fail-fast approach)
|
|
"""
|
|
mongodb_url = get_mongodb_url()
|
|
|
|
try:
|
|
# Create client with short timeout for fail-fast behavior
|
|
client = AsyncIOMotorClient(
|
|
mongodb_url,
|
|
serverSelectionTimeoutMS=5000, # 5 seconds timeout
|
|
connectTimeoutMS=5000,
|
|
socketTimeoutMS=5000
|
|
)
|
|
|
|
# Test connection by running admin command
|
|
client.admin.command('ping')
|
|
|
|
print(f"Successfully connected to MongoDB at {mongodb_url}")
|
|
return client
|
|
|
|
except (ConnectionFailure, ServerSelectionTimeoutError) as e:
|
|
print(f"ERROR: Failed to connect to MongoDB at {mongodb_url}")
|
|
print(f"Connection error: {str(e)}")
|
|
print("MongoDB is required for this application. Please ensure MongoDB is running and accessible.")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"ERROR: Unexpected error connecting to MongoDB: {str(e)}")
|
|
sys.exit(1)
|
|
|
|
|
|
def get_database() -> Database:
|
|
"""
|
|
Get MongoDB database instance (singleton pattern).
|
|
|
|
Returns:
|
|
Database: MongoDB database instance
|
|
|
|
This function implements singleton pattern to ensure only one
|
|
database connection is created throughout the application lifecycle.
|
|
"""
|
|
global _client, _database
|
|
|
|
if _database is None:
|
|
if _client is None:
|
|
_client = create_mongodb_client()
|
|
|
|
database_name = get_mongodb_database_name()
|
|
_database = _client[database_name]
|
|
print(f"Connected to database: {database_name}")
|
|
|
|
return _database
|
|
|
|
|
|
def close_database_connection():
|
|
"""
|
|
Close MongoDB database connection.
|
|
|
|
This function should be called during application shutdown
|
|
to properly close the database connection.
|
|
"""
|
|
global _client, _database
|
|
|
|
if _client is not None:
|
|
_client.close()
|
|
_client = None
|
|
_database = None
|
|
print("MongoDB connection closed")
|
|
|
|
|
|
def get_mongodb_client() -> Optional[MongoClient]:
|
|
"""
|
|
Get the raw MongoDB client instance.
|
|
|
|
Returns:
|
|
MongoClient or None: MongoDB client instance if connected
|
|
|
|
This is primarily for testing purposes or advanced use cases
|
|
where direct client access is needed.
|
|
"""
|
|
return _client
|
|
|
|
|
|
def get_extra_args(session):
|
|
# Build kwargs only if session is provided
|
|
kwargs = {}
|
|
if session is not None:
|
|
kwargs["session"] = session
|
|
|
|
return kwargs
|
|
|
|
|
|
def test_database_connection() -> bool:
|
|
"""
|
|
Test if database connection is working.
|
|
|
|
Returns:
|
|
bool: True if connection is working, False otherwise
|
|
|
|
This function can be used for health checks.
|
|
"""
|
|
try:
|
|
db = get_database()
|
|
# Simple operation to test connection
|
|
db.command('ping')
|
|
return True
|
|
except Exception:
|
|
return False
|