wag-managment-api-service-version-4/Services/MongoDb/Models/mixins.py

"""
MongoDB CRUD Operation Mixins.

This module provides mixins for common MongoDB operations:
1. Document creation (insert)
2. Document retrieval (find)
3. Document updates
4. Document deletion
5. Aggregation operations
"""

from typing import Any, Dict, List, Optional
from functools import wraps

from pymongo.collection import Collection
from pymongo.errors import (
    ConnectionFailure,
    OperationFailure,
    ServerSelectionTimeoutError,
    PyMongoError,
)

from ApiLibrary.common.line_number import get_line_number_for_error
from ErrorHandlers.ErrorHandlers.api_exc_handler import HTTPExceptionApi


def handle_mongo_errors(func):
    """Decorator to handle MongoDB operation errors.

    Catches MongoDB-specific errors and converts them to HTTPExceptionApi.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ConnectionFailure:
            raise HTTPExceptionApi(
                error_code="HTTP_503_SERVICE_UNAVAILABLE",
                lang="en",
                loc=get_line_number_for_error(),
                sys_msg="MongoDB connection failed",
            )
        except ServerSelectionTimeoutError:
            raise HTTPExceptionApi(
                error_code="HTTP_504_GATEWAY_TIMEOUT",
                lang="en",
                loc=get_line_number_for_error(),
                sys_msg="MongoDB connection timed out",
            )
        except OperationFailure as e:
            raise HTTPExceptionApi(
                error_code="HTTP_400_BAD_REQUEST",
                lang="en",
                loc=get_line_number_for_error(),
                sys_msg=str(e),
            )
        except PyMongoError as e:
            raise HTTPExceptionApi(
                error_code="HTTP_500_INTERNAL_SERVER_ERROR",
                lang="en",
                loc=get_line_number_for_error(),
                sys_msg=str(e),
            )

    return wrapper


class MongoInsertMixin:
    """Mixin for MongoDB insert operations."""

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

    @handle_mongo_errors
    def insert_many(self, collection: Collection, documents: List[Dict[str, Any]]):
        """Insert multiple documents into the collection."""
        result = collection.insert_many(documents)
        return result


class MongoFindMixin:
    """Mixin for MongoDB find operations."""

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

    @handle_mongo_errors
    def find_many(
        self,
        collection: Collection,
        filter_query: Dict[str, Any],
        projection: Optional[Dict[str, Any]] = None,
        sort: Optional[List[tuple]] = None,
        limit: Optional[int] = None,
        skip: Optional[int] = None,
    ):
        """Find multiple documents in the collection with pagination support."""
        cursor = 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:
    """Mixin for MongoDB update operations."""

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

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


class MongoDeleteMixin:
    """Mixin for MongoDB delete operations."""

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

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


class MongoAggregateMixin:
    """Mixin for MongoDB aggregation operations."""

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