374 lines
14 KiB
Python
374 lines
14 KiB
Python
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()
|