mongo_service.py

import logging
from typing import Dict, List, Union, Any
from bson import ObjectId
import ast
import json
import re
from pymongo import MongoClient
from threading import Lock
from urllib.parse import urlparse
from pymongo.uri_parser import parse_uri

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

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

    def get_client(self, uri: str) -> MongoClient:
        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]}...")
                    self._clients[uri] = MongoClient(
                        uri, 
                        serverSelectionTimeoutMS=5000, 
                        connectTimeoutMS=5000, 
                        maxPoolSize=50
                    )
        return self._clients[uri]

mongo_manager = MongoConnectionManager()

# ==============================================================================
# HELPER FUNCTIONS
# ==============================================================================
def convert_oid(obj):
    """Recursively convert $oid or ObjectId 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]
    return obj

def sanitize_json_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
    """Parses input string/dict/list into Python objects, handling ObjectIds."""
    if isinstance(query_input, (dict, list)): return query_input

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

    # Sanitize JS specific wrappers
    query_str = re.sub(r"ObjectId\(['\"]([a-fA-F0-9]+)['\"]\)", r"'\1'", query_str)
    query_str = re.sub(r"ISODate\(['\"](.*?)['\"]\)", r"'\1'", query_str)
    
    try:
        return json.loads(query_str)
    except:
        try:
            return ast.literal_eval(query_str)
        except Exception as e:
            logger.error(f"Failed to parse query: {e}")
            raise ValueError(f"Parsing error: {str(e)}")

# ==============================================================================
# MAIN EXECUTION LOGIC
# ==============================================================================
def execute_mongo_operation(
    mongo_uri: str, 
    db_name: str, 
    collection_name: str, 
    query: Union[Dict, List],
    limited: bool = False,
    limit_rows: int = 20
):
    client = mongo_manager.get_client(mongo_uri)
    
    try:
        db = client[db_name]
        collection = db[collection_name]
        results = []
        
        # --- Pipeline (List) ---
        if isinstance(query, list):
            pipeline = query
            if limited:
                if not (pipeline and "$limit" in pipeline[-1]):
                    pipeline.append({"$limit": limit_rows})
            cursor = collection.aggregate(pipeline, allowDiskUse=True)
            results = list(cursor)
            
        # --- Find (Dict) ---
        elif isinstance(query, dict):
            cursor = collection.find(query)
            if limited: cursor = cursor.limit(limit_rows)
            results = list(cursor)
            
        else:
            raise ValueError("Query must be a Pipeline (List) or Filter (Dict)")

        return results
    except Exception as e:
        logger.error(f"DB Execution Error: {e}")
        raise e
    
def extract_database_name(uri: str) -> str:
    """
    Robustly extracts database name from a MongoDB URI.
    1. Tries PyMongo's strict parser.
    2. Falls back to standard URL path extraction (like your TS code).
    """
    if not uri:
        return None

    # Method 1: PyMongo Strict Parser
    try:
        parsed = parse_uri(uri)
        if parsed.get('database'):
            return parsed['database']
    except Exception:
        pass # Fallback if PyMongo complains about the URI format

    # Method 2: Generic URL Parsing (Mimics your TypeScript logic)
    try:
        # Standard URL parsing
        result = urlparse(uri)
        
        # The database name is usually the 'path' part (e.g., /api_platform)
        path = result.path
        
        # Remove leading slash if present
        if path.startswith('/'):
            path = path[1:]
            
        # If empty, return None
        return path if path else None
    except Exception:
        return None