import time import functools from pymongo import MongoClient from pymongo.errors import PyMongoError from Controllers.Mongo.config import mongo_configs def retry_operation(max_attempts=3, delay=1.0, backoff=2.0, exceptions=(PyMongoError,)): """ Decorator for retrying MongoDB operations with exponential backoff. Args: max_attempts: Maximum number of retry attempts delay: Initial delay between retries in seconds backoff: Multiplier for delay after each retry exceptions: Tuple of exceptions to catch and retry """ def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): mtries, mdelay = max_attempts, delay while mtries > 1: try: return func(*args, **kwargs) except exceptions as e: time.sleep(mdelay) mtries -= 1 mdelay *= backoff return func(*args, **kwargs) return wrapper return decorator class MongoDBHandler: """ A MongoDB handler that provides context manager access to specific collections with automatic retry capability. Implements singleton pattern. """ _instance = None _debug_mode = False # Set to True to enable debug mode def __new__(cls, *args, **kwargs): """ Implement singleton pattern for the handler. """ if cls._instance is None: cls._instance = super(MongoDBHandler, cls).__new__(cls) cls._instance._initialized = False return cls._instance def __init__(self, debug_mode=False, mock_mode=False): """Initialize the MongoDB handler. Args: debug_mode: If True, use a simplified connection for debugging mock_mode: If True, use mock collections instead of real MongoDB connections """ if not hasattr(self, "_initialized") or not self._initialized: self._debug_mode = debug_mode self._mock_mode = mock_mode if mock_mode: # In mock mode, we don't need a real connection string self.uri = "mongodb://mock:27017/mockdb" print("MOCK MODE: Using simulated MongoDB connections") elif debug_mode: # Use a direct connection without authentication for testing self.uri = f"mongodb://{mongo_configs.HOST}:{mongo_configs.PORT}/{mongo_configs.DB}" print(f"DEBUG MODE: Using direct connection: {self.uri}") else: # Use the configured connection string with authentication self.uri = mongo_configs.url print(f"Connecting to MongoDB: {self.uri}") # Define MongoDB client options with increased timeouts for better reliability self.client_options = { "maxPoolSize": 5, "minPoolSize": 1, "maxIdleTimeMS": 60000, "waitQueueTimeoutMS": 5000, "serverSelectionTimeoutMS": 10000, "connectTimeoutMS": 30000, "socketTimeoutMS": 45000, "retryWrites": True, "retryReads": True, } self._initialized = True def collection(self, collection_name: str): """ Get a context manager for a specific collection. Args: collection_name: Name of the collection to access Returns: A context manager for the specified collection """ return CollectionContext(self, collection_name) class CollectionContext: """ Context manager for MongoDB collections with automatic retry capability. """ def __init__(self, db_handler: MongoDBHandler, collection_name: str): """ Initialize collection context. Args: db_handler: Reference to the MongoDB handler collection_name: Name of the collection to access """ self.db_handler = db_handler self.collection_name = collection_name self.client = None self.collection = None def __enter__(self): """ Enter context, establishing a new connection. Returns: The MongoDB collection object with retry capabilities """ # If we're in mock mode, return a mock collection immediately if self.db_handler._mock_mode: return self._create_mock_collection() try: # Create a new client connection self.client = MongoClient(self.db_handler.uri, **self.db_handler.client_options) if self.db_handler._debug_mode: # In debug mode, we explicitly use the configured DB db_name = mongo_configs.DB print(f"DEBUG MODE: Using database '{db_name}'") else: # In normal mode, extract database name from the URI try: db_name = self.client.get_database().name except Exception: db_name = mongo_configs.DB print(f"Using fallback database '{db_name}'") self.collection = self.client[db_name][self.collection_name] # Enhance collection methods with retry capabilities self._add_retry_capabilities() return self.collection except pymongo.errors.OperationFailure as e: if "Authentication failed" in str(e): print(f"MongoDB authentication error: {e}") print("Attempting to reconnect with direct connection...") try: # Try a direct connection without authentication for testing direct_uri = f"mongodb://{mongo_configs.HOST}:{mongo_configs.PORT}/{mongo_configs.DB}" print(f"Trying direct connection: {direct_uri}") self.client = MongoClient(direct_uri, **self.db_handler.client_options) self.collection = self.client[mongo_configs.DB][self.collection_name] self._add_retry_capabilities() return self.collection except Exception as inner_e: print(f"Direct connection also failed: {inner_e}") # Fall through to mock collection creation else: print(f"MongoDB operation error: {e}") if self.client: self.client.close() self.client = None except Exception as e: print(f"MongoDB connection error: {e}") if self.client: self.client.close() self.client = None return self._create_mock_collection() def _create_mock_collection(self): """ Create a mock collection for testing or graceful degradation. This prevents the application from crashing when MongoDB is unavailable. Returns: A mock MongoDB collection with simulated behaviors """ from unittest.mock import MagicMock if self.db_handler._mock_mode: print(f"MOCK MODE: Using mock collection '{self.collection_name}'") else: print(f"Using mock MongoDB collection '{self.collection_name}' for graceful degradation") # Create in-memory storage for this mock collection if not hasattr(self.db_handler, '_mock_storage'): self.db_handler._mock_storage = {} if self.collection_name not in self.db_handler._mock_storage: self.db_handler._mock_storage[self.collection_name] = [] mock_collection = MagicMock() mock_data = self.db_handler._mock_storage[self.collection_name] # Define behavior for find operations def mock_find(query=None, *args, **kwargs): # Simple implementation that returns all documents return mock_data def mock_find_one(query=None, *args, **kwargs): # Simple implementation that returns the first matching document if not mock_data: return None return mock_data[0] def mock_insert_one(document, *args, **kwargs): # Add _id if not present if '_id' not in document: document['_id'] = f"mock_id_{len(mock_data)}" mock_data.append(document) result = MagicMock() result.inserted_id = document['_id'] return result def mock_insert_many(documents, *args, **kwargs): inserted_ids = [] for doc in documents: result = mock_insert_one(doc) inserted_ids.append(result.inserted_id) result = MagicMock() result.inserted_ids = inserted_ids return result def mock_update_one(query, update, *args, **kwargs): result = MagicMock() result.modified_count = 1 return result def mock_update_many(query, update, *args, **kwargs): result = MagicMock() result.modified_count = len(mock_data) return result def mock_delete_one(query, *args, **kwargs): result = MagicMock() result.deleted_count = 1 if mock_data: mock_data.pop(0) # Just remove the first item for simplicity return result def mock_delete_many(query, *args, **kwargs): count = len(mock_data) mock_data.clear() result = MagicMock() result.deleted_count = count return result def mock_count_documents(query, *args, **kwargs): return len(mock_data) def mock_aggregate(pipeline, *args, **kwargs): return [] def mock_create_index(keys, **kwargs): return f"mock_index_{keys}" # Assign the mock implementations mock_collection.find.side_effect = mock_find mock_collection.find_one.side_effect = mock_find_one mock_collection.insert_one.side_effect = mock_insert_one mock_collection.insert_many.side_effect = mock_insert_many mock_collection.update_one.side_effect = mock_update_one mock_collection.update_many.side_effect = mock_update_many mock_collection.delete_one.side_effect = mock_delete_one mock_collection.delete_many.side_effect = mock_delete_many mock_collection.count_documents.side_effect = mock_count_documents mock_collection.aggregate.side_effect = mock_aggregate mock_collection.create_index.side_effect = mock_create_index # Add retry capabilities to the mock collection self._add_retry_capabilities_to_mock(mock_collection) self.collection = mock_collection return self.collection def _add_retry_capabilities(self): """ Add retry capabilities to all collection methods. """ # Store original methods for common operations original_insert_one = self.collection.insert_one original_insert_many = self.collection.insert_many original_find_one = self.collection.find_one original_find = self.collection.find original_update_one = self.collection.update_one original_update_many = self.collection.update_many original_delete_one = self.collection.delete_one original_delete_many = self.collection.delete_many original_replace_one = self.collection.replace_one original_count_documents = self.collection.count_documents # Add retry capabilities to methods self.collection.insert_one = retry_operation()(original_insert_one) self.collection.insert_many = retry_operation()(original_insert_many) self.collection.find_one = retry_operation()(original_find_one) self.collection.find = retry_operation()(original_find) self.collection.update_one = retry_operation()(original_update_one) self.collection.update_many = retry_operation()(original_update_many) self.collection.delete_one = retry_operation()(original_delete_one) self.collection.delete_many = retry_operation()(original_delete_many) self.collection.replace_one = retry_operation()(original_replace_one) self.collection.count_documents = retry_operation()(original_count_documents) def _add_retry_capabilities_to_mock(self, mock_collection): """ Add retry capabilities to mock collection methods. This is a simplified version that just wraps the mock methods. Args: mock_collection: The mock collection to enhance """ # List of common MongoDB collection methods to add retry capabilities to methods = [ 'insert_one', 'insert_many', 'find_one', 'find', 'update_one', 'update_many', 'delete_one', 'delete_many', 'replace_one', 'count_documents', 'aggregate' ] # Add retry decorator to each method for method_name in methods: if hasattr(mock_collection, method_name): original_method = getattr(mock_collection, method_name) setattr( mock_collection, method_name, retry_operation(max_retries=1, retry_interval=0)(original_method) ) def __exit__(self, exc_type, exc_val, exc_tb): """ Exit context, closing the connection. """ if self.client: self.client.close() self.client = None self.collection = None # Create a singleton instance of the MongoDB handler mongo_handler = MongoDBHandler()