211 lines
6.8 KiB
Python
211 lines
6.8 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 MongoDBConfig:
|
|
"""
|
|
Configuration class for MongoDB connection settings.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
uri: str = "mongodb://localhost:27017/",
|
|
max_pool_size: int = 20,
|
|
min_pool_size: int = 10,
|
|
max_idle_time_ms: int = 30000,
|
|
wait_queue_timeout_ms: int = 2000,
|
|
server_selection_timeout_ms: int = 5000,
|
|
**additional_options,
|
|
):
|
|
"""
|
|
Initialize MongoDB configuration.
|
|
"""
|
|
self.uri = uri
|
|
self.client_options = {
|
|
"maxPoolSize": max_pool_size,
|
|
"minPoolSize": min_pool_size,
|
|
"maxIdleTimeMS": max_idle_time_ms,
|
|
"waitQueueTimeoutMS": wait_queue_timeout_ms,
|
|
"serverSelectionTimeoutMS": server_selection_timeout_ms,
|
|
**additional_options,
|
|
}
|
|
|
|
|
|
class MongoDBHandler(MongoDBConfig):
|
|
"""
|
|
A MongoDB handler that provides context manager access to specific collections
|
|
with automatic retry capability.
|
|
"""
|
|
|
|
_instance = None
|
|
|
|
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,
|
|
uri: str,
|
|
max_pool_size: int = 20,
|
|
min_pool_size: int = 5,
|
|
max_idle_time_ms: int = 10000,
|
|
wait_queue_timeout_ms: int = 1000,
|
|
server_selection_timeout_ms: int = 3000,
|
|
**additional_options,
|
|
):
|
|
"""
|
|
Initialize the MongoDB handler (only happens once due to singleton).
|
|
"""
|
|
# Only initialize once
|
|
if not hasattr(self, "_initialized") or not self._initialized:
|
|
super().__init__(
|
|
uri=uri,
|
|
max_pool_size=max_pool_size,
|
|
min_pool_size=min_pool_size,
|
|
max_idle_time_ms=max_idle_time_ms,
|
|
wait_queue_timeout_ms=wait_queue_timeout_ms,
|
|
server_selection_timeout_ms=server_selection_timeout_ms,
|
|
**additional_options,
|
|
)
|
|
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
|
|
"""
|
|
try:
|
|
# Create a new client connection
|
|
self.client = MongoClient(
|
|
self.db_handler.uri, **self.db_handler.client_options
|
|
)
|
|
# Get database from URI
|
|
db_name = self.client.get_database().name
|
|
self.collection = self.client[db_name][self.collection_name]
|
|
|
|
# Enhance collection methods with retry capabilities
|
|
self._add_retry_capabilities()
|
|
|
|
return self.collection
|
|
except Exception as e:
|
|
if self.client:
|
|
self.client.close()
|
|
raise
|
|
|
|
def _add_retry_capabilities(self):
|
|
"""
|
|
Add retry capabilities to collection methods.
|
|
"""
|
|
# Store original methods
|
|
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 __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
|
|
|
|
|
|
mongo_handler = MongoDBHandler(uri=mongo_configs.url)
|
|
"""
|
|
max_pool_size: int = 20,
|
|
min_pool_size: int = 10,
|
|
max_idle_time_ms: int = 30000,
|
|
wait_queue_timeout_ms: int = 2000,
|
|
server_selection_timeout_ms: int = 5000,
|
|
""" |