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]