update mongo exec

controller.py

@@ -13,7 +13,6 @@ import psutil
 import ssl
 from typing import List, Optional, Dict, Any, TypeVar, Generic
 from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
-from threading import Lock
 
 # --- FastAPI & Core ---
 from fastapi import FastAPI, HTTPException, Depends
@@ -32,16 +31,13 @@ from asyncmy.cursors import DictCursor
 
 # --- Database Drivers (Sync - unused but kept for imports) ---
 from bson import ObjectId
-import mysql.connector
-import psycopg2
-from psycopg2.extras import RealDictCursor
 import uvicorn
 
 # --- Existing Services ---
 from csv_analysis_service import execute_analysis_logic
 from csv_chart_service import execute_python_code
 from csv_metadata_service import CsvDataRequest, CsvInfoRequest, CsvInfoResponse, PythonExecutionRequest, PythonExecutionResponse, execute_python_logic, get_csv_basic_info, get_robust_csv_rows
-from mongo_service import execute_mongo_operation, parse_query_input
+from mongo_service import convert_oid, execute_mongo_operation, extract_database_name, sanitize_json_input
 from pydantic_csv_analysis_model import AnalysisRequest, AnalysisResponse
 from pydantic_csv_charts_model import ChartExecutionPayload, ChartExecutionResponse
 from pydantic_mongo_executor_model import ExecutorPayload, ExecutorResponse
@@ -452,19 +448,79 @@ async def execute_postgres_endpoint(query: PgQueryRequest, token: str = Depends(
 async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(validate_token)):
     request_id = str(uuid.uuid4())[:8]
     start_time = time.time()
+    
     try:
-        parsed_query = parse_query_input(payload.generated_query)
+        # 1. Determine Database Name
+        final_db_name = payload.db_name
+        
+        # If DB Name is missing in payload, try to extract it from URI
+        if not final_db_name:
+            final_db_name = extract_database_name(payload.mongo_uri)
+
+        # Final check
+        if not final_db_name:
+            raise HTTPException(
+                status_code=400, 
+                detail="Database name not provided in payload and could not be extracted from the URI path."
+            )
+
+        # 2. Parse Query & Determine Collection Name
+        parsed_input = sanitize_json_input(payload.generated_query)
+        final_query = parsed_input
+        final_collection = payload.collection_name
+
+        # Check for structured AI response: { "data": [ { "collectionName": "...", "pipeline": [...] } ] }
+        if isinstance(parsed_input, dict) and "data" in parsed_input and isinstance(parsed_input["data"], list):
+            if len(parsed_input["data"]) > 0:
+                extracted_data = parsed_input["data"][0]
+                
+                # Extract Pipeline
+                if "pipeline" in extracted_data:
+                    final_query = extracted_data["pipeline"]
+                
+                # Extract Collection Name (Overrides payload if present)
+                if "collectionName" in extracted_data and extracted_data["collectionName"]:
+                    final_collection = extracted_data["collectionName"]
+
+        if not final_collection:
+             raise HTTPException(status_code=400, detail="Collection name could not be determined from payload or query.")
+
+        # 3. Execute
         unique_params = {
-            "uri": payload.mongo_uri, "db": payload.db_name, "col": payload.collection_name, 
-            "q": parsed_query, "lim": payload.limited, "lrows": payload.limit_rows
+            "uri": payload.mongo_uri, 
+            "db": final_db_name, 
+            "col": final_collection,
+            "q": str(final_query),
+            "lim": payload.limited
         }
+        
+        final_query = convert_oid(final_query)
+
         result_data = await coalescer.execute(
-            "mongo", unique_params, run_in_threadpool, execute_mongo_operation, 
-            mongo_uri=payload.mongo_uri, db_name=payload.db_name, collection_name=payload.collection_name, 
-            query=parsed_query, limited=payload.limited, limit_rows=payload.limit_rows
+            "mongo", 
+            unique_params, 
+            run_in_threadpool, 
+            execute_mongo_operation, 
+            mongo_uri=payload.mongo_uri, 
+            db_name=final_db_name, 
+            collection_name=final_collection, 
+            query=final_query, 
+            limited=payload.limited, 
+            limit_rows=payload.limit_rows
         )
-        return ExecutorResponse(status="success", count=len(result_data), data=jsonable_encoder(result_data, custom_encoder={ObjectId: str}), duration_seconds=round(time.time() - start_time, 4), request_id=request_id)
+
+        return ExecutorResponse(
+            status="success", 
+            count=len(result_data), 
+            data=jsonable_encoder(result_data, custom_encoder={ObjectId: str}), 
+            duration_seconds=round(time.time() - start_time, 4), 
+            request_id=request_id
+        )
+
+    except HTTPException as he:
+        raise he
     except Exception as e:
+        logger.error(f"Endpoint Error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/api/execute_chart", response_model=ChartExecutionResponse)

mongo_service.py

@@ -1,11 +1,13 @@
 import logging
-from typing import Dict, List, Union
+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 # Needed for thread-safe pooling
+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",
@@ -14,7 +16,7 @@ logging.basicConfig(
 logger = logging.getLogger("mongo_service")
 
 # ==============================================================================
-# MONGO CONNECTION POOL MANAGER
+# MONGO CONNECTION MANAGER
 # ==============================================================================
 class MongoConnectionManager:
     def __init__(self):
@@ -22,83 +24,57 @@ class MongoConnectionManager:
         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
+                        connectTimeoutMS=5000, 
+                        maxPoolSize=50
                     )
         return self._clients[uri]
 
-# Global Manager Instance
 mongo_manager = MongoConnectionManager()
 
 # ==============================================================================
-# HELPER FUNCTIONS (Unchanged)
+# HELPER FUNCTIONS
 # ==============================================================================
-
 def convert_oid(obj):
-    """Recursively convert $oid format or ObjectId objects to strings."""
-    if isinstance(obj, ObjectId):
-        return str(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"])
+        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
+    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)
+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 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
+    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 (Modified)
+# MAIN EXECUTION LOGIC
 # ==============================================================================
-
 def execute_mongo_operation(
     mongo_uri: str, 
     db_name: str, 
@@ -107,73 +83,66 @@ def execute_mongo_operation(
     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 ---
+        # --- Pipeline (List) ---
         if isinstance(query, list):
-            # Apply Limit if requested and not already present at the end
+            pipeline = query
             if limited:
-                if not (query and "$limit" in query[-1]):
-                    query.append({"$limit": limit_rows})
-            
-            cursor = collection.aggregate(query, allowDiskUse=True)
+                if not (pipeline and "$limit" in pipeline[-1]):
+                    pipeline.append({"$limit": limit_rows})
+            cursor = collection.aggregate(pipeline, allowDiskUse=True)
             results = list(cursor)
             
-        # --- Find / Command ---
+        # --- Find (Dict) ---
         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)
-            
+            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)")
+            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
     
-    # IMPORTANT: DO NOT CLOSE THE CLIENT
-    # finally:
-    #     if client: client.close()  <-- REMOVED
+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

pydantic_mongo_executor_model.py

@@ -1,15 +1,12 @@
-
 # --- Mongo Models ---
 from typing import Any, Dict, List, Optional, Union
 from pydantic import BaseModel, Field
 
-
 class ExecutorPayload(BaseModel):
     mongo_uri: str = Field(..., description="MongoDB connection string")
-    db_name: str = Field(..., description="Target Database Name")
-    collection_name: str = Field(..., description="Target Collection Name")
-    generated_query: Union[Dict, List, str] = Field(..., description="The query to execute (Find dict or Aggregation list)")
-    user_query: Optional[str] = Field(None, description="Original user question for logging context")
+    db_name: Optional[str] = Field(None, description="Target Database Name") 
+    collection_name: Optional[str] = Field(None, description="Target Collection Name")
+    generated_query: Union[Dict, List, str] = Field(..., description="Query Payload")
     limited: bool = False
     limit_rows: int = 20