ok

controller.py

@@ -5,13 +5,14 @@ import time
 import uuid
 import os
 import re
-import asyncio  # Added for parallel batching
-import multiprocessing  # Added for worker calculation
-import hashlib # Added for Coalescing
-import json # Added for Coalescing
-import psutil # Added for Dynamic RAM calculation
-from typing import List, Optional, Dict, Any, TypeVar, Generic  # Added Generic/TypeVar
+import asyncio
+import multiprocessing
+import hashlib
+import json
+import psutil
+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
@@ -19,18 +20,17 @@ from fastapi.encoders import jsonable_encoder
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.concurrency import run_in_threadpool
-from anyio import to_thread # Added for AnyIO 4.x concurrency tuning
+from anyio import to_thread
 from dotenv import load_dotenv
 from pydantic import BaseModel
 
-# --- Database Drivers & Pooling ---
+# --- Database Drivers ---
 from bson import ObjectId
 import mysql.connector
-from mysql.connector import pooling # For MySQL Pooling
+from mysql.connector import pooling
 import psycopg2
-from psycopg2 import pool as pg_pool # For Postgres Pooling
+from psycopg2 import pool as pg_pool
 from psycopg2.extras import RealDictCursor
-from threading import Lock # To handle concurrency safely
 import uvicorn
 
 # --- Existing Services ---
@@ -54,56 +54,34 @@ logging.basicConfig(
 logger = logging.getLogger("API_Controller")
 
 def get_dynamic_thread_limit():
-    """
-    Calculates a safe thread limit based on available RAM and CPU cores.
-    Returns an integer safe for to_thread.current_default_thread_limiter().total_tokens
-    """
+    """Calculates a safe thread limit based on available RAM."""
     try:
-        # 1. Get System Resources
         total_ram_bytes = psutil.virtual_memory().total
         total_cores = multiprocessing.cpu_count()
-        
-        # 2. Identify how many Uvicorn Workers are running
-        # Assuming production setup uses: max(1, total_cores - 2)
         num_workers = max(1, total_cores - 2)
-        
-        # 3. Calculate RAM available PER WORKER
         ram_per_worker = total_ram_bytes / num_workers
-        
-        # 4. Reserve Buffer (Keep 30% for OS/Python overhead, use 70% for threads)
         safe_ram_pool = ram_per_worker * 0.70
-        
-        # 5. Estimate Thread Cost (~8MB conservative safety margin)
         BYTES_PER_THREAD = 8 * 1024 * 1024 
-        
         calculated_limit = int(safe_ram_pool / BYTES_PER_THREAD)
-        
-        # 6. Apply Reasonable Hard Caps (Min 100, Max 3000)
         final_limit = max(100, min(calculated_limit, 3000))
         
         logger.info(f"Dynamic Limit Config: {total_ram_bytes/(1024**3):.2f}GB RAM / {num_workers} Workers. Limit: {final_limit}")
         return final_limit
-
     except Exception as e:
         logger.warning(f"Failed to calculate dynamic threads ({e}). Fallback to 1000.")
         return 1000
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup logic: Calculate and set dynamic thread limit
     safe_limit = get_dynamic_thread_limit()
     to_thread.current_default_thread_limiter().total_tokens = safe_limit
     logger.info(f"Worker Process Started: Thread pool capacity set to {safe_limit}.")
     yield
-    # Shutdown logic (if any) goes here
 
-app = FastAPI(
-    title="Unified Data Executor API (Mongo, SQL, CSV)", 
-    lifespan=lifespan
-)
+app = FastAPI(title="Unified Data Executor API", lifespan=lifespan)
 
 # ==============================================================================
-# HIGH-CONCURRENCY BATCH MODELS & STARTUP
+# BATCH MODELS
 # ==============================================================================
 T = TypeVar("T")
 
@@ -113,11 +91,10 @@ class BatchRequest(BaseModel, Generic[T]):
 class BatchResponse(BaseModel, Generic[T]):
     responses: List[T]
 
-# --- Directory Setup ---
+# --- Directory & CORS ---
 CHART_DIR = "generated_charts"
 os.makedirs(CHART_DIR, exist_ok=True)
 
-# --- CORS ---
 origins_env = os.getenv("ALLOWED_ORIGINS", "*")
 ORIGINS = [origin.strip() for origin in origins_env.split(",")]
 
@@ -133,57 +110,13 @@ app.add_middleware(
 security = HTTPBearer()
 API_SECRET_TOKEN = os.getenv("API_BEARER_TOKEN")
 
-if not API_SECRET_TOKEN:
-    logger.warning("WARNING: API_BEARER_TOKEN not set in .env file! Security is compromised.")
-
 async def validate_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
     if credentials.credentials != API_SECRET_TOKEN:
         raise HTTPException(status_code=403, detail="Invalid Authentication Token")
     return credentials.credentials
 
-
-# ==============================================================================
-# PYDANTIC MODELS
-# ==============================================================================
-
-class SqlQueryRequest(BaseModel):
-    database_url: str
-    sql_query: str
-    limit_rows: Optional[int] = 20
-    limited: bool = False
-
-class SqlQueryResponse(BaseModel):
-    success: bool
-    results: Optional[List[Dict[str, Any]]] = None
-    columns: Optional[List[str]] = None
-    rowCount: Optional[int] = 0
-    executionTime: Optional[float] = 0.0
-    error: Optional[str] = None
-    request_id: str
-    is_aggregate: bool = False
-    limited: bool = False
-    message: Optional[str] = None
-
-class PgQueryRequest(BaseModel):
-    database_url: str
-    sql_query: str
-    limit_rows: Optional[int] = 20
-    limited: bool = False 
-
-class PgQueryResponse(BaseModel):
-    success: bool
-    results: Optional[List[Dict[str, Any]]] = None
-    columns: Optional[List[str]] = None
-    rowCount: Optional[int] = 0
-    executionTime: Optional[float] = 0.0
-    error: Optional[str] = None
-    request_id: str
-    is_aggregate: bool = False
-    limited: bool = False
-    message: Optional[str] = None
-    
 # ==============================================================================
-# CONNECTION POOL MANAGER
+# CONNECTION POOL MANAGER (With Keepalives)
 # ==============================================================================
 class ConnectionPoolManager:
     def __init__(self):
@@ -202,9 +135,14 @@ class ConnectionPoolManager:
                     "host": parsed.hostname,
                     "port": parsed.port or 3306,
                     "database": parsed.path.lstrip("/"),
-                    "connect_timeout": 5
+                    "connect_timeout": 5,
+                    "pool_reset_session": True  # Validates connection on checkout
                 }
-                self._mysql_pools[db_url] = pooling.MySQLConnectionPool(pool_name=str(uuid.uuid4()), pool_size=10, **db_config)
+                self._mysql_pools[db_url] = pooling.MySQLConnectionPool(
+                    pool_name=str(uuid.uuid4()), 
+                    pool_size=10, 
+                    **db_config
+                )
         return self._mysql_pools[db_url].get_connection()
 
     def get_postgres_connection(self, db_url: str):
@@ -214,6 +152,7 @@ class ConnectionPoolManager:
                 parsed = urlparse(db_url)
                 qs = parse_qs(parsed.query)
                 sslmode = qs.get('sslmode', ['require'])[0] if 'sslmode' in qs else 'prefer'
+                
                 db_config = {
                     "host": parsed.hostname,
                     "port": parsed.port or 5432,
@@ -221,24 +160,25 @@ class ConnectionPoolManager:
                     "user": parsed.username,
                     "password": parsed.password,
                     "sslmode": sslmode,
-                    "connect_timeout": 5
+                    "connect_timeout": 5,
+                    # --- TCP Keepalives (Prevents SSL Closed Unexpectedly) ---
+                    "keepalives": 1,
+                    "keepalives_idle": 30,
+                    "keepalives_interval": 10,
+                    "keepalives_count": 5
                 }
                 self._pg_pools[db_url] = pg_pool.ThreadedConnectionPool(1, 10, **db_config)
         return self._pg_pools[db_url].getconn()
 
     def return_postgres_connection(self, db_url, conn, close=False):
-        """
-        Returns connection to pool. 
-        If close=True, the connection is discarded (used for dead connections).
-        """
+        """Returns connection to pool. If close=True, discards it."""
         if db_url in self._pg_pools and conn:
             self._pg_pools[db_url].putconn(conn, close=close)
 
 pool_manager = ConnectionPoolManager()
 
-
 # ==============================================================================
-# REQUEST COALESCER (SINGLEFLIGHT)
+# REQUEST COALESCER
 # ==============================================================================
 class RequestCoalescer:
     def __init__(self):
@@ -252,26 +192,19 @@ class RequestCoalescer:
 
     async def execute(self, prefix: str, unique_params: dict, func, *args, **kwargs):
         key = self._generate_key(prefix, unique_params)
-
         async with self._lock:
             if key in self._active_requests:
                 return await self._active_requests[key]
-            
             loop = asyncio.get_running_loop()
             future = loop.create_future()
             self._active_requests[key] = future
 
         try:
-            # Execute the function (run_in_threadpool) with args/kwargs
             result = await func(*args, **kwargs)
-            
-            if not future.done():
-                future.set_result(result)
+            if not future.done(): future.set_result(result)
             return result
-            
         except Exception as e:
-            if not future.done():
-                future.set_exception(e)
+            if not future.done(): future.set_exception(e)
             raise e
         finally:
             async with self._lock:
@@ -280,18 +213,13 @@ class RequestCoalescer:
 
 coalescer = RequestCoalescer()
 
-
 # ==============================================================================
-# DB EXECUTION LOGIC
+# DB EXECUTION LOGIC (WITH RETRY & SELF-HEALING)
 # ==============================================================================
 
 def is_aggregate_query(query: str) -> bool:
     query_lower = query.lower()
-    aggregate_patterns = [
-        r'\bcount\s*\(', r'\bsum\s*\(', r'\bavg\s*\(', 
-        r'\bmin\s*\(', r'\bmax\s*\(', r'\bgroup\s+by\b',
-        r'\bdistinct\b', r'\bhaving\b'
-    ]
+    aggregate_patterns = [r'\bcount\s*\(', r'\bsum\s*\(', r'\bavg\s*\(', r'\bmin\s*\(', r'\bmax\s*\(', r'\bgroup\s+by\b', r'\bdistinct\b', r'\bhaving\b']
     for pattern in aggregate_patterns:
         if re.search(pattern, query_lower):
             return True
@@ -317,218 +245,200 @@ def normalize_postgres_uri(uri: str) -> str:
     except Exception:
         return uri
 
+# --- SELF-HEALING MYSQL EXECUTION ---
 def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
-    connection = None
-    cursor = None
-    response = {"success": False, "results": None, "columns": None, "rowCount": 0, "executionTime": 0.0, "error": None, "is_aggregate": False, "limited": False, "message": ""}
     
-    # Flag for bad connections
-    connection_broken = False
+    # Retry Loop: If connection is dead, we retry once
+    for attempt in range(2):
+        connection = None
+        cursor = None
+        try:
+            connection = pool_manager.get_mysql_connection(db_url)
+            cursor = connection.cursor(dictionary=True)
+            
+            clean_query = sql_query.strip()
+            query_lower = clean_query.lower()
 
-    try:
-        connection = pool_manager.get_mysql_connection(db_url)
-        cursor = connection.cursor(dictionary=True)
-        
-        clean_query = sql_query.strip()
-        query_lower = clean_query.lower()
-
-        if not query_lower.startswith("select"):
-            cursor.execute(clean_query)
-            connection.commit()
-            response.update({"success": True, "message": "Query executed successfully (Non-SELECT)."})
-            return response
-
-        if not limited:
-            cursor.execute(clean_query)
-            results = cursor.fetchall()
-            response["message"] = f"Raw query executed. Returned {len(results)} row(s)."
-            response["limited"] = False
-            response["is_aggregate"] = is_aggregate_query(clean_query)
-        else:
-            if is_aggregate_query(clean_query):
+            if not query_lower.startswith("select"):
+                cursor.execute(clean_query)
+                connection.commit()
+                return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}
+
+            if not limited:
                 cursor.execute(clean_query)
                 results = cursor.fetchall()
-                response["message"] = f"Aggregate query completed. Returned {len(results)} row(s)."
-                response["is_aggregate"] = True
+                is_aggregate = is_aggregate_query(clean_query)
+                message = f"Raw query executed. Returned {len(results)} row(s)."
+                is_limited_result = False
             else:
-                final_query = clean_query.rstrip(';').strip()
-                if not re.search(r'\blimit\s+\d+', query_lower):
-                    final_query = f"{final_query} LIMIT {max_rows}"
-                
-                cursor.execute(final_query)
-                results = cursor.fetchall()
-                is_limited_result = (len(results) == max_rows)
-                response["message"] = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
-                response["limited"] = is_limited_result
-                response["is_aggregate"] = False
-
-        columns = [col[0] for col in cursor.description] if cursor.description else []
-        response.update({"success": True, "results": jsonable_encoder(results), "columns": columns, "rowCount": len(results), "executionTime": time.time() - start_time})
-        return response
-
-    except Exception as e:
-        # Detect broken pipe or connection lost errors in MySQL
-        err_str = str(e).lower()
-        if "lost connection" in err_str or "gone away" in err_str or isinstance(e, mysql.connector.errors.OperationalError):
-            connection_broken = True
-        
-        response["error"] = str(e)
-        return response
-    finally:
-        try:
-            if cursor: cursor.close()
-        except: pass
-        
-        if connection:
-            if connection_broken:
-                # Discard dead connection (do NOT return to pool)
-                try: connection.close() 
+                if is_aggregate_query(clean_query):
+                    cursor.execute(clean_query)
+                    results = cursor.fetchall()
+                    is_aggregate = True
+                    message = f"Aggregate query completed. Returned {len(results)} row(s)."
+                    is_limited_result = False
+                else:
+                    final_query = clean_query.rstrip(';').strip()
+                    if not re.search(r'\blimit\s+\d+', query_lower):
+                        final_query = f"{final_query} LIMIT {max_rows}"
+                    cursor.execute(final_query)
+                    results = cursor.fetchall()
+                    is_limited_result = (len(results) == max_rows)
+                    message = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
+                    is_aggregate = False
+
+            columns = [col[0] for col in cursor.description] if cursor.description else []
+            return {
+                "success": True, "results": jsonable_encoder(results), "columns": columns, 
+                "rowCount": len(results), "executionTime": time.time() - start_time,
+                "is_aggregate": is_aggregate, "limited": is_limited_result, "message": message, "error": None
+            }
+
+        except (mysql.connector.errors.OperationalError, mysql.connector.errors.DatabaseError) as e:
+            # Fatal connection error
+            if connection:
+                try: connection.close() # Close properly so pool knows it's bad (or discards it)
                 except: pass
-            else:
-                # Return healthy connection to pool
-                connection.close() 
+            
+            # If this was the first attempt, try again with a fresh connection
+            if attempt == 0:
+                logger.warning(f"MySQL connection lost. Retrying... Error: {e}")
+                continue 
+            
+            return {"success": False, "error": f"MySQL Error: {str(e)}", "executionTime": 0.0}
+            
+        except Exception as e:
+            # Logic error (Syntax, etc). Do not retry.
+            if connection: connection.close()
+            return {"success": False, "error": str(e), "executionTime": 0.0}
+            
+        finally:
+            if cursor: 
+                try: cursor.close()
+                except: pass
+            # Connection close handled in blocks above to handle "poisoned" logic correctly
 
+# --- SELF-HEALING POSTGRES EXECUTION ---
 def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
-    connection = None
-    cursor = None
-    response = {"success": False, "results": None, "columns": None, "rowCount": 0, "executionTime": 0.0, "error": None, "is_aggregate": False, "limited": False, "message": ""}
     
-    # Flag to determine if connection is dead
-    connection_broken = False 
-
-    try:
-        connection = pool_manager.get_postgres_connection(db_url)
-        cursor = connection.cursor(cursor_factory=RealDictCursor)
-        clean_query = sql_query.strip()
-        query_lower = clean_query.lower()
-        
-        if not query_lower.startswith(("select", "show", "explain", "with")):
-            cursor.execute(clean_query)
-            connection.commit()
-            response.update({"success": True, "message": "Query executed successfully (Non-SELECT)."})
-            return response
+    # Retry Loop for Dead Connections
+    for attempt in range(2):
+        connection = None
+        cursor = None
+        try:
+            connection = pool_manager.get_postgres_connection(db_url)
+            cursor = connection.cursor(cursor_factory=RealDictCursor)
             
-        if not limited:
-            cursor.execute(clean_query)
-            results = cursor.fetchall()
-            response["message"] = f"Raw query executed. Returned {len(results)} row(s)."
-            response["limited"] = False
-            response["is_aggregate"] = is_aggregate_query(clean_query)
-        else:
-            if is_aggregate_query(clean_query):
+            clean_query = sql_query.strip()
+            query_lower = clean_query.lower()
+            
+            if not query_lower.startswith(("select", "show", "explain", "with")):
+                cursor.execute(clean_query)
+                connection.commit()
+                # Return Success
+                pool_manager.return_postgres_connection(db_url, connection)
+                return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}
+            
+            # Read Logic
+            if not limited:
                 cursor.execute(clean_query)
                 results = cursor.fetchall()
-                response["message"] = f"Aggregate query completed. Returned {len(results)} row(s)."
-                response["is_aggregate"] = True
-                response["limited"] = False
+                is_aggregate = is_aggregate_query(clean_query)
+                message = f"Raw query executed. Returned {len(results)} row(s)."
+                is_limited_result = False
             else:
-                final_query = clean_query.rstrip(';').strip()
-                if not re.search(r'\blimit\s+\d+', query_lower):
-                    final_query = f"{final_query} LIMIT {max_rows}"
-                cursor.execute(final_query)
-                results = cursor.fetchall()
-                is_limited_result = (len(results) == max_rows)
-                response["message"] = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
-                response["limited"] = is_limited_result
-                response["is_aggregate"] = False
-
-        columns = [desc[0] for desc in cursor.description] if cursor.description else []
-        clean_results = jsonable_encoder(results, custom_encoder={uuid.UUID: str, ObjectId: str})
-        response.update({"success": True, "results": clean_results, "columns": columns, "rowCount": len(results), "executionTime": time.time() - start_time})
-        return response
-    except Exception as e:
-        # Detect fatal connection errors
-        err_msg = str(e).lower()
-        if "closed" in err_msg or "terminat" in err_msg or isinstance(e, (psycopg2.InterfaceError, psycopg2.OperationalError)):
-            connection_broken = True 
-        
-        if connection and not connection_broken:
-            try:
-                connection.rollback()
-            except Exception:
-                connection_broken = True 
-
-        response["error"] = str(e)
-        return response
-    finally:
-        try: 
-            if cursor: cursor.close()
-        except: pass
-        
-        if connection: 
-            # If broken, set close=True to discard it from pool
-            pool_manager.return_postgres_connection(db_url, connection, close=connection_broken)
+                if is_aggregate_query(clean_query):
+                    cursor.execute(clean_query)
+                    results = cursor.fetchall()
+                    is_aggregate = True
+                    message = f"Aggregate query completed. Returned {len(results)} row(s)."
+                    is_limited_result = False
+                else:
+                    final_query = clean_query.rstrip(';').strip()
+                    if not re.search(r'\blimit\s+\d+', query_lower):
+                        final_query = f"{final_query} LIMIT {max_rows}"
+                    cursor.execute(final_query)
+                    results = cursor.fetchall()
+                    is_limited_result = (len(results) == max_rows)
+                    message = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
+                    is_aggregate = False
+
+            columns = [desc[0] for desc in cursor.description] if cursor.description else []
+            clean_results = jsonable_encoder(results, custom_encoder={uuid.UUID: str, ObjectId: str})
+            
+            # Success! Return conn
+            pool_manager.return_postgres_connection(db_url, connection)
+            return {
+                "success": True, "results": clean_results, "columns": columns, 
+                "rowCount": len(results), "executionTime": time.time() - start_time,
+                "is_aggregate": is_aggregate, "limited": is_limited_result, "message": message, "error": None
+            }
+
+        except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
+            # Bad Connection. DISCARD IT.
+            if connection:
+                pool_manager.return_postgres_connection(db_url, connection, close=True)
+            
+            if attempt == 0:
+                logger.warning(f"Postgres connection dead ({e}). Retrying with fresh connection...")
+                continue
+            return {"success": False, "error": f"DB Connection Failed: {str(e)}", "executionTime": 0.0}
 
+        except Exception as e:
+            # Logic/SQL Error. Rollback and return to pool.
+            if connection:
+                try: connection.rollback()
+                except: pass # If rollback fails, pool.putconn might fail next, but we try
+                pool_manager.return_postgres_connection(db_url, connection, close=False)
+                
+            return {"success": False, "error": str(e), "executionTime": 0.0}
+            
+        finally:
+            if cursor: 
+                try: cursor.close()
+                except: pass
 
 # ==============================================================================
 # API ROUTES
 # ==============================================================================
 
-@app.post("/api/execute_mongo", response_model=ExecutorResponse)
-async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(validate_token)):
-    request_id = str(uuid.uuid4())[:8]
-    start_time = time.time()
-    try:
-        # Parse logic is fast, can stay outside threadpool
-        parsed_query = parse_query_input(payload.generated_query)
-        
-        # --- COALESCING MAGIC ---
-        # Create a unique signature for this request
-        unique_params = {
-            "uri": payload.mongo_uri,
-            "db": payload.db_name,
-            "col": payload.collection_name,
-            "q": parsed_query, # parsed_query is a Dict or List, json.dumps handles it
-            "lim": payload.limited,
-            "lrows": payload.limit_rows
-        }
+class SqlQueryRequest(BaseModel):
+    database_url: str
+    sql_query: str
+    limit_rows: Optional[int] = 20
+    limited: bool = False
 
-        # Use the coalescer to prevent duplicate simultaneous DB hits
-        result_data = await coalescer.execute(
-            "mongo",                 # Prefix
-            unique_params,           # Unique params dict
-            run_in_threadpool,       # Runner
-            execute_mongo_operation, # The function to run
-            # Arguments for 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
-        )
-        # ------------------------
-
-        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 Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+class SqlQueryResponse(BaseModel):
+    success: bool
+    results: Optional[List[Dict[str, Any]]] = None
+    columns: Optional[List[str]] = None
+    rowCount: Optional[int] = 0
+    executionTime: Optional[float] = 0.0
+    error: Optional[str] = None
+    request_id: str
+    is_aggregate: bool = False
+    limited: bool = False
+    message: Optional[str] = None
 
-@app.post("/api/execute_chart", response_model=ChartExecutionResponse)
-async def execute_chart_endpoint(payload: ChartExecutionPayload, token: str = Depends(validate_token)):
-    request_id = str(uuid.uuid4())[:8]
-    try:
-        image_bytes, error_msg, logs = await run_in_threadpool(execute_python_code, code=payload.code, csv_url=payload.csv_url)
-        
-        if error_msg: 
-            return ChartExecutionResponse(status="error", error=error_msg, output_log=logs, request_id=request_id)
+class PgQueryRequest(BaseModel):
+    database_url: str
+    sql_query: str
+    limit_rows: Optional[int] = 20
+    limited: bool = False 
 
-        if payload.return_base64:
-            base64_str = base64.b64encode(image_bytes).decode('utf-8')
-            return ChartExecutionResponse(status="success", base64_image=base64_str, output_log=logs, request_id=request_id)
-        else:
-            unique_name = f"{uuid.uuid4()}.png"
-            public_url = await run_in_threadpool(upload_bytes_to_supabase, image_bytes=image_bytes, file_name=unique_name, chat_id=payload.chat_id)
-            return ChartExecutionResponse(status="success", image_url=public_url, output_log=logs, request_id=request_id)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+class PgQueryResponse(BaseModel):
+    success: bool
+    results: Optional[List[Dict[str, Any]]] = None
+    columns: Optional[List[str]] = None
+    rowCount: Optional[int] = 0
+    executionTime: Optional[float] = 0.0
+    error: Optional[str] = None
+    request_id: str
+    is_aggregate: bool = False
+    limited: bool = False
+    message: Optional[str] = None
 
 @app.post("/api/execute_sql_query", response_model=SqlQueryResponse)
 async def execute_mysql_endpoint(query: SqlQueryRequest, token: str = Depends(validate_token)):
@@ -536,34 +446,17 @@ async def execute_mysql_endpoint(query: SqlQueryRequest, token: str = Depends(va
     try:
         normalized_url = normalize_mysql_uri(query.database_url)
         limit_val = query.limit_rows if query.limit_rows is not None else 20
-        
-        # Unique ID for Coalescing
         unique_params = {
-            "db": normalized_url,
-            "q": query.sql_query,
-            "l": limit_val,
-            "lim": query.limited
+            "db": normalized_url, "q": query.sql_query, "l": limit_val, "lim": query.limited
         }
-
-        # FIXED: Pass _run_mysql_synchronously as the first POSITIONAL argument after run_in_threadpool
         result_dict = await coalescer.execute(
-            "mysql",                  # Prefix
-            unique_params,            # Unique Params
-            run_in_threadpool,        # Runner
-            _run_mysql_synchronously, # Arg 1 (The Function)
-            db_url=normalized_url,    # Kwargs
-            sql_query=query.sql_query,
-            max_rows=limit_val,
-            limited=query.limited
+            "mysql", unique_params, run_in_threadpool, _run_mysql_synchronously, 
+            db_url=normalized_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
         )
-
         final_result = result_dict.copy()
         final_result["request_id"] = request_id
         return SqlQueryResponse(**final_result)
-
     except Exception as e:
-        # If run_in_threadpool fails (e.g. TypeError) it ends up here
-        logger.error(f"MySQL Endpoint Error: {str(e)}")
         raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})
 
 @app.post("/api/execute_postgres_query", response_model=PgQueryResponse)
@@ -572,34 +465,54 @@ async def execute_postgres_endpoint(query: PgQueryRequest, token: str = Depends(
     try:
         clean_url = normalize_postgres_uri(query.database_url)
         limit_val = query.limit_rows if query.limit_rows is not None else 20
-        
         unique_params = {
-            "db": clean_url,
-            "q": query.sql_query,
-            "l": limit_val,
-            "lim": query.limited
+            "db": clean_url, "q": query.sql_query, "l": limit_val, "lim": query.limited
         }
-
-        # FIXED: Pass _run_postgres_synchronously as the first POSITIONAL argument
         result_dict = await coalescer.execute(
-            "postgres",
-            unique_params,
-            run_in_threadpool,
-            _run_postgres_synchronously, # Arg 1 (The Function)
-            db_url=clean_url,
-            sql_query=query.sql_query,
-            max_rows=limit_val,
-            limited=query.limited
+            "postgres", unique_params, run_in_threadpool, _run_postgres_synchronously, 
+            db_url=clean_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
         )
-
         final_result = result_dict.copy()
         final_result["request_id"] = request_id
         return PgQueryResponse(**final_result)
     except Exception as e:
-        logger.error(f"Postgres Endpoint Error: {str(e)}")
         raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})
 
-# ... (CSV and Report endpoints remain unchanged) ...
+@app.post("/api/execute_mongo", response_model=ExecutorResponse)
+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)
+        unique_params = {
+            "uri": payload.mongo_uri, "db": payload.db_name, "col": payload.collection_name, 
+            "q": parsed_query, "lim": payload.limited, "lrows": payload.limit_rows
+        }
+        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
+        )
+        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 Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/execute_chart", response_model=ChartExecutionResponse)
+async def execute_chart_endpoint(payload: ChartExecutionPayload, token: str = Depends(validate_token)):
+    request_id = str(uuid.uuid4())[:8]
+    try:
+        image_bytes, error_msg, logs = await run_in_threadpool(execute_python_code, code=payload.code, csv_url=payload.csv_url)
+        if error_msg: return ChartExecutionResponse(status="error", error=error_msg, output_log=logs, request_id=request_id)
+        if payload.return_base64:
+            base64_str = base64.b64encode(image_bytes).decode('utf-8')
+            return ChartExecutionResponse(status="success", base64_image=base64_str, output_log=logs, request_id=request_id)
+        else:
+            unique_name = f"{uuid.uuid4()}.png"
+            public_url = await run_in_threadpool(upload_bytes_to_supabase, image_bytes=image_bytes, file_name=unique_name, chat_id=payload.chat_id)
+            return ChartExecutionResponse(status="success", image_url=public_url, output_log=logs, request_id=request_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.post("/api/execute_csv_analysis", response_model=AnalysisResponse)
 async def execute_analysis_endpoint(payload: AnalysisRequest, token: str = Depends(validate_token)):
     request_id = str(uuid.uuid4())[:8]
@@ -669,10 +582,21 @@ async def batch_execute_mongo(payload: BatchRequest[ExecutorPayload], token: str
     return BatchResponse(responses=responses)
 
 # --- Test Endpoint ---
+class TestCalcRequest(BaseModel):
+    value: int
+
+def _heavy_calculation_task(value: int) -> int:
+    time.sleep(0.1)
+    return value * 2
+
+@app.post("/api/test/parallel_calc_no_auth")
+async def test_parallel_calc_endpoint(payload: TestCalcRequest):
+    try:
+        result = await run_in_threadpool(_heavy_calculation_task, value=payload.value)
+        return {"success": True, "result": result, "process_id": os.getpid()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-# ==============================================================================
-# KEEP-ALIVE ROUTE
-# ==============================================================================
 @app.get("/")
 async def root():
     return {"message": "Python Code Execution Server is running"}
@@ -681,26 +605,9 @@ async def root():
 async def ping():
     return {"message": "I am alive!"}
 
-
-# ==============================================================================
-# HIGH PERFORMANCE SERVER EXECUTION
-# ==============================================================================
-
 if __name__ == "__main__":
     host = os.getenv("HOST", "0.0.0.0")
     port = int(os.getenv("PORT", 7860))
-    
-    # Calculate workers (save 2 cores for system overhead)
-    # Ensure at least 1 worker exists
     num_workers = max(1, multiprocessing.cpu_count() - 2)
-    
     print(f"Starting production server on {host}:{port} with {num_workers} workers...")
-    print("Using 'asyncio' loop to prevent Pandas/Numpy segfaults.")
-
-    uvicorn.run(
-        "controller:app",  
-        host=host, 
-        port=port, 
-        workers=num_workers,
-        loop="asyncio",
-    )
+    uvicorn.run("controller:app", host=host, port=port, workers=num_workers, loop="asyncio")