Spaces:
Runtime error
Runtime error
| from dependency_injector.resources import AsyncResource | |
| from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection | |
| from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError | |
| from loguru import logger | |
| from pydantic import BaseModel, PrivateAttr | |
| from typing import Any, Dict, Self | |
| from ctp_slack_bot.core.config import Settings | |
| from ctp_slack_bot.utils import sanitize_mongo_db_uri | |
| class MongoDB(BaseModel): | |
| """ | |
| MongoDB connection manager using Motor for async operations. | |
| """ | |
| settings: Settings | |
| _client: PrivateAttr = PrivateAttr() | |
| _db: PrivateAttr = PrivateAttr() | |
| class Config: | |
| frozen=True | |
| arbitrary_types_allowed = True | |
| def __init__(self: Self, **data: Dict[str, Any]) -> None: | |
| super().__init__(**data) | |
| logger.debug("Created {}", self.__class__.__name__) | |
| def connect(self: Self) -> None: | |
| """Initialize MongoDB client with settings.""" | |
| try: | |
| connection_string = self.settings.MONGODB_URI.get_secret_value() | |
| logger.debug("Connecting to MongoDB using URI: {}", sanitize_mongo_db_uri(connection_string)) | |
| # Create client with appropriate settings. | |
| self._client = AsyncIOMotorClient( | |
| connection_string, | |
| serverSelectionTimeoutMS=5000, | |
| connectTimeoutMS=10000, | |
| socketTimeoutMS=45000, | |
| maxPoolSize=100, | |
| retryWrites=True, | |
| w="majority" | |
| ) | |
| # Get the database name. | |
| db_name = self.settings.MONGODB_NAME | |
| self._db = self._client[db_name] | |
| logger.debug("MongoDB client initialized for database: {}", db_name) | |
| except Exception as e: | |
| logger.error("Failed to initialize MongoDB client: {}", e) | |
| self._client = None | |
| self._db = None | |
| raise | |
| async def ping(self: Self) -> bool: | |
| """Check if MongoDB connection is alive.""" | |
| try: | |
| await self._client.admin.command("ping") | |
| logger.debug("MongoDB connection is active!") | |
| return True | |
| except (ConnectionFailure, ServerSelectionTimeoutError) as e: | |
| logger.error("MongoDB connection failed: {}", e) | |
| except Exception as e: | |
| logger.error("Unexpected error during MongoDB ping: {}", e) | |
| return False | |
| async def get_collection(self: Self, name: str) -> AsyncIOMotorCollection: | |
| """ | |
| Get a collection by name or creates it if it doesn’t exist. | |
| """ | |
| # First ensure we can connect at all. | |
| if not await self.ping(): | |
| logger.error("Cannot get collection '{}' because a MongoDB connection is not available.", name) | |
| raise ConnectionError("MongoDB connection is not available.") | |
| try: | |
| # Get all collection names to check if this one exists. | |
| logger.debug("Checking if collection '{}' exists…", name) | |
| collection_names = await self._db.list_collection_names() | |
| if name not in collection_names: | |
| logger.info("Collection '{}' does not exist. Creating it…", name) | |
| # Create the collection. | |
| await self._db.create_collection(name) | |
| logger.debug("Successfully created collection: {}", name) | |
| else: | |
| logger.debug("Collection '{}' already exists!", name) | |
| # Get and return the collection. | |
| collection = self._db[name] | |
| return collection | |
| except Exception as e: | |
| logger.error("Error accessing collection '{}': {}", name, e) | |
| raise | |
| def close(self: Self) -> None: | |
| """Close the MongoDB connection.""" | |
| if self._client: | |
| self._client.close() | |
| logger.info("Closed MongoDB connection.") | |
| self._client = None | |
| self._db = None | |
| class MongoDBResource(AsyncResource): | |
| async def init(self: Self, settings: Settings) -> MongoDB: | |
| logger.info("Initializing MongoDB connection for database: {}", settings.MONGODB_NAME) | |
| mongo_db = MongoDB(settings=settings) | |
| mongo_db.connect() | |
| await self._test_connection(mongo_db) | |
| return mongo_db | |
| async def _test_connection(self: Self, mongo_db: MongoDB) -> None: | |
| """Test MongoDB connection and log the result.""" | |
| try: | |
| is_connected = await mongo_db.ping() | |
| if is_connected: | |
| logger.info("MongoDB connection test successful!") | |
| else: | |
| logger.error("MongoDB connection test failed!") | |
| except Exception as e: | |
| logger.error("Error testing MongoDB connection: {}", e) | |
| raise | |
| async def shutdown(self: Self, mongo_db: MongoDB) -> None: | |
| """Close MongoDB connection on shutdown.""" | |
| try: | |
| mongo_db.close() | |
| except Exception as e: | |
| logger.error("Error closing MongoDB connection: {}", e) | |