src/ctp_slack_bot/db/mongo_db.py

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)