updated postgres and mongo updated
This commit is contained in:
@@ -35,42 +35,14 @@ def retry_operation(max_attempts=3, delay=1.0, backoff=2.0, exceptions=(PyMongoE
|
||||
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):
|
||||
class MongoDBHandler:
|
||||
"""
|
||||
A MongoDB handler that provides context manager access to specific collections
|
||||
with automatic retry capability.
|
||||
with automatic retry capability. Implements singleton pattern.
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_debug_mode = False # Set to True to enable debug mode
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
"""
|
||||
@@ -81,30 +53,42 @@ class MongoDBHandler(MongoDBConfig):
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str,
|
||||
max_pool_size: int = 5,
|
||||
min_pool_size: int = 2,
|
||||
max_idle_time_ms: int = 10000,
|
||||
wait_queue_timeout_ms: int = 1000,
|
||||
server_selection_timeout_ms: int = 3000,
|
||||
**additional_options,
|
||||
):
|
||||
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
|
||||
"""
|
||||
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._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):
|
||||
@@ -145,29 +129,172 @@ class CollectionContext:
|
||||
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
|
||||
)
|
||||
# Get database from URI
|
||||
db_name = self.client.get_database().name
|
||||
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()
|
||||
raise
|
||||
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 collection methods.
|
||||
Add retry capabilities to all collection methods.
|
||||
"""
|
||||
# Store original 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
|
||||
@@ -191,6 +318,31 @@ class CollectionContext:
|
||||
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.
|
||||
@@ -201,11 +353,5 @@ class CollectionContext:
|
||||
self.collection = None
|
||||
|
||||
|
||||
mongo_handler = MongoDBHandler(
|
||||
uri=mongo_configs.url,
|
||||
max_pool_size=5,
|
||||
min_pool_size=2,
|
||||
max_idle_time_ms=30000,
|
||||
wait_queue_timeout_ms=2000,
|
||||
server_selection_timeout_ms=5000,
|
||||
)
|
||||
# Create a singleton instance of the MongoDB handler
|
||||
mongo_handler = MongoDBHandler()
|
||||
|
||||
Reference in New Issue
Block a user