mongo_service.py

import logging
from typing import Dict, List, Union
from bson import ObjectId
import ast
import json
import re
from pymongo import MongoClient
from threading import Lock # Needed for thread-safe pooling

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO
)
logger = logging.getLogger("mongo_service")

# ==============================================================================
# MONGO CONNECTION POOL MANAGER
# ==============================================================================
class MongoConnectionManager:
    def __init__(self):
        self._clients: Dict[str, MongoClient] = {}
        self._lock = Lock()

    def get_client(self, uri: str) -> MongoClient:
        """
        Returns a cached MongoClient instance.
        If it doesn't exist, creates one with a connection pool.
        """
        # Double-checked locking for performance
        if uri not in self._clients:
            with self._lock:
                if uri not in self._clients:
                    logger.info(f"Initialize new MongoDB Client for URI: {uri[:20]}...")
                    # maxPoolSize=50 allows 50 concurrent ops per URI. 
                    # The rest will wait in queue automatically.
                    self._clients[uri] = MongoClient(
                        uri, 
                        serverSelectionTimeoutMS=5000, 
                        connectTimeoutMS=5000,
                        maxPoolSize=50, 
                        minPoolSize=1
                    )
        return self._clients[uri]

# Global Manager Instance
mongo_manager = MongoConnectionManager()

# ==============================================================================
# HELPER FUNCTIONS (Unchanged)
# ==============================================================================

def convert_oid(obj):
    """Recursively convert $oid format or ObjectId objects to strings."""
    if isinstance(obj, ObjectId):
        return str(obj)
    if isinstance(obj, dict):
        if set(obj.keys()) == {"$oid"}:
            return str(obj["$oid"])
        return {k: convert_oid(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_oid(item) for item in obj]
    else:
        return obj

def parse_query_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
    """Ensures the input is a valid MongoDB query object (Dict or List)."""
    if isinstance(query_input, (dict, list)):
        return convert_oid(query_input)

    query_str = str(query_input).strip()
    
    try:
        return json.loads(query_str)
    except json.JSONDecodeError:
        pass

    match = re.search(r"```(?:python|json|javascript)?\s*(.*?)\s*```", query_str, re.DOTALL)
    if match:
        query_str = match.group(1).strip()

    query_str = re.sub(r'^db\.\w+\.\w+\(', '', query_str)
    if query_str.endswith(')'):
        query_str = query_str[:-1]

    try:
        parsed = ast.literal_eval(query_str)
        return convert_oid(parsed)
    except (ValueError, SyntaxError) as e:
        logger.error(f"Failed to parse query string: {e}")
        raise ValueError(f"Could not parse query string: {str(e)}")

def get_value_ignore_case(d: Dict, keys: List[str], default=None):
    for k in keys:
        if k in d:
            return d[k]
    return default

# ==============================================================================
# MAIN EXECUTION LOGIC (Modified)
# ==============================================================================

def execute_mongo_operation(
    mongo_uri: str, 
    db_name: str, 
    collection_name: str, 
    query: Union[Dict, List],
    limited: bool = False,
    limit_rows: int = 20
):
    """
    Executes MongoDB operations using the shared connection pool.
    """
    # 1. GET CLIENT FROM MANAGER (Do NOT create new MongoClient here)
    client = mongo_manager.get_client(mongo_uri)
    
    try:
        db = client[db_name]
        collection = db[collection_name]
        
        results = []
        
        # --- Aggregation ---
        if isinstance(query, list):
            # Apply Limit if requested and not already present at the end
            if limited:
                if not (query and "$limit" in query[-1]):
                    query.append({"$limit": limit_rows})
            
            cursor = collection.aggregate(query, allowDiskUse=True)
            results = list(cursor)
            
        # --- Find / Command ---
        elif isinstance(query, dict):
            command_keys = {'filter', 'query', '$query', 'projection', 'sort', 'limit', 'skip'}
            has_command_keys = bool(set(query.keys()) & command_keys)
            
            cursor = None

            if has_command_keys and ('filter' in query or 'query' in query or '$query' in query):
                query_filter = get_value_ignore_case(query, ['filter', 'query', '$query'], {})
                projection = get_value_ignore_case(query, ['projection', 'fields'], None)
                if projection == {}: projection = None 
                
                internal_limit = int(get_value_ignore_case(query, ['limit'], 0))
                skip = int(get_value_ignore_case(query, ['skip'], 0))
                sort_val = get_value_ignore_case(query, ['sort', '$orderby'], None)
                
                cursor = collection.find(query_filter, projection)
                
                if sort_val:
                    if isinstance(sort_val, dict):
                        sort_val = list(sort_val.items())
                    cursor = cursor.sort(sort_val)
                
                if skip > 0: cursor = cursor.skip(skip)
                
                if limited:
                    cursor = cursor.limit(limit_rows)
                elif internal_limit > 0: 
                    cursor = cursor.limit(internal_limit)

            else:
                cursor = collection.find(query)
                if limited:
                    cursor = cursor.limit(limit_rows)
            
            results = list(cursor)
        else:
            raise ValueError("Query must be a Dictionary (find) or List (aggregate)")

        return results

    except Exception as e:
        logger.error(f"DB Execution Error: {e}")
        raise e
    
    # IMPORTANT: DO NOT CLOSE THE CLIENT
    # finally:
    #     if client: client.close()  <-- REMOVED