wag-services-and-backend-latest/Services/MongoService/provider.py

from typing import Optional, Dict, Any, List, TypeVar, Iterator
from contextlib import contextmanager

from pymongo import MongoClient
from pymongo.collection import Collection

from Configs.mongo import MongoConfig
from Services.MongoService.handlers import mongo_error_wrapper


class MongoBase:
    """Base class for MongoDB connection and operations."""

    collection: Collection = None


class MongoErrorHandler:
    """Error handler for MongoDB operations."""

    ...


class MongoInsertMixin(MongoBase):
    """Mixin for MongoDB insert operations."""

    def insert_one(self, document: Dict[str, Any]):
        """Insert a single document into the collection."""
        return self.collection.insert_one(document)

    def insert_many(self, documents: List[Dict[str, Any]]):
        """Insert multiple documents."""
        return self.collection.insert_many(documents)


class MongoFindMixin(MongoBase):
    """Mixin for MongoDB find operations."""

    @mongo_error_wrapper
    def find_one(
        self,
        filter_query: Dict[str, Any],
        projection: Optional[Dict[str, Any]] = None,
    ):
        """Find a single document in the collection."""
        return self.collection.find_one(filter_query, projection)

    @mongo_error_wrapper
    def find_many(
        self,
        filter_query: Dict[str, Any],
        projection: Optional[Dict[str, Any]] = None,
        sort: Optional[List[tuple[str, int]]] = None,
        limit: Optional[int] = None,
        skip: Optional[int] = None,
    ):
        """Find multiple documents in the collection with pagination support."""
        cursor = self.collection.find(filter_query, projection)
        if sort:
            cursor = cursor.sort(sort)
        if skip:
            cursor = cursor.skip(skip)
        if limit:
            cursor = cursor.limit(limit)
        return list(cursor)


class MongoUpdateMixin(MongoBase):
    """Mixin for MongoDB update operations."""

    @mongo_error_wrapper
    def update_one(
        self,
        filter_query: Dict[str, Any],
        update_data: Dict[str, Any],
        upsert: bool = False,
    ):
        """Update a single document in the collection."""
        return self.collection.update_one(filter_query, update_data, upsert=upsert)

    @mongo_error_wrapper
    def update_many(
        self,
        filter_query: Dict[str, Any],
        update_data: Dict[str, Any],
        upsert: bool = False,
    ):
        """Update multiple documents in the collection."""
        return self.collection.update_many(filter_query, update_data, upsert=upsert)


class MongoDeleteMixin(MongoBase):
    """Mixin for MongoDB delete operations."""

    @mongo_error_wrapper
    def delete_one(self, filter_query: Dict[str, Any]):
        """Delete a single document from the collection."""
        return self.collection.delete_one(filter_query)

    @mongo_error_wrapper
    def delete_many(self, filter_query: Dict[str, Any]):
        """Delete multiple documents from the collection."""
        return self.collection.delete_many(filter_query)


class MongoAggregateMixin(MongoBase):
    """Mixin for MongoDB aggregation operations."""

    @mongo_error_wrapper
    def aggregate(self, collection: Collection, pipeline: List[Dict[str, Any]]):
        """Execute an aggregation pipeline on the collection."""
        result = collection.aggregate(pipeline)
        return result


class MongoProvider(
    MongoUpdateMixin,
    MongoInsertMixin,
    MongoFindMixin,
    MongoDeleteMixin,
    MongoAggregateMixin,
):
    """Main MongoDB actions class that inherits all CRUD operation mixins.

    This class provides a unified interface for all MongoDB operations while
    managing collections based on company UUID and storage reason.
    """

    def __init__(self, client: MongoClient, database: str, storage_reason: list[str]):
        """Initialize MongoDB actions with client and collection info.

        Args:
            client: MongoDB client
            database: Database name to use
            storage_reason: Storage reason for collection naming
        """
        self.delimiter = "|"
        self._client = client
        self._database = database
        self._storage_reason: list[str] = storage_reason
        self._collection = None
        self.use_collection(storage_reason)

    @staticmethod
    @contextmanager
    def mongo_client() -> Iterator[MongoClient]:
        """
        Context provider for MongoDB test client.
        # Example Usage
        with mongo_client() as client:
            db = client["your_database"]
            print(db.list_collection_names())
        """
        client = MongoClient(MongoConfig.URL)
        try:
            client.admin.command("ping")  # Test connection
            yield client
        finally:
            client.close()  # Ensure proper cleanup

    @property
    def collection(self) -> Collection:
        """Get current MongoDB collection."""
        return self._collection

    def use_collection(self, storage_name_list: list[str]) -> None:
        """Switch to a different collection.

        Args:
            storage_name_list: New storage reason for collection naming
        """
        collection_name = ""
        for each_storage_reason in storage_name_list:
            if self.delimiter in str(each_storage_reason):
                raise ValueError(
                    f"Storage reason cannot contain delimiter : {self.delimiter}"
                )
            collection_name += f"{self.delimiter}{each_storage_reason}"
        collection_name = collection_name[1:]
        self._collection = self._client[self._database][collection_name]