do

controller.py

@@ -10,6 +10,7 @@ import multiprocessing
 import hashlib
 import json
 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
@@ -29,7 +30,7 @@ import asyncpg
 import asyncmy
 from asyncmy.cursors import DictCursor
 
-# --- Database Drivers (Kept for compatibility/types if needed, but unused in SQL routes) ---
+# --- Database Drivers (Sync - unused but kept for imports) ---
 from bson import ObjectId
 import mysql.connector
 import psycopg2
@@ -68,7 +69,7 @@ class AsyncPoolManager:
     async def get_pg_pool(self, db_url: str) -> asyncpg.Pool:
         async with self._lock:
             if db_url not in self._pg_pools:
-                logger.info(f"Creating new AsyncPG pool for: {db_url[:20]}...")
+                logger.info(f"Creating new AsyncPG pool for: {db_url[:25]}...")
                 try:
                     # asyncpg handles Keepalives and SSL automatically better than psycopg2
                     pool = await asyncpg.create_pool(
@@ -85,13 +86,27 @@ class AsyncPoolManager:
         return self._pg_pools[db_url]
 
     async def get_mysql_pool(self, db_url: str) -> asyncmy.Pool:
+        # Check first without lock for speed
+        if db_url in self._mysql_pools:
+            return self._mysql_pools[db_url]
+
         async with self._lock:
+            # Check again inside lock
             if db_url not in self._mysql_pools:
-                logger.info(f"Creating new AsyncMy pool for: {db_url[:20]}...")
+                logger.info(f"Creating new AsyncMy pool for: {db_url[:25]}...")
                 parsed = urlparse(db_url)
-                # Check for SSL requirement in query params
                 qs = parse_qs(parsed.query)
-                # AsyncMy often auto-negotiates SSL, but we can enforce it if needed based on qs
+                
+                # Determine SSL settings
+                # asyncmy needs an SSL context or ssl=True/False
+                ssl_ctx = None
+                if 'ssl-mode' in qs:
+                    mode = qs['ssl-mode'][0].upper()
+                    if mode != 'DISABLED':
+                        # Create a default SSL context that allows self-signed certs (common in cloud DBs)
+                        ssl_ctx = ssl.create_default_context()
+                        ssl_ctx.check_hostname = False
+                        ssl_ctx.verify_mode = ssl.CERT_NONE
                 
                 try:
                     pool = await asyncmy.create_pool(
@@ -103,12 +118,14 @@ class AsyncPoolManager:
                         minsize=1,
                         maxsize=20,
                         autocommit=True,
-                        pool_recycle=280, # Recycle before default 8hr timeout to prevent "Gone Away"
+                        pool_recycle=280,
+                        ssl=ssl_ctx # Pass the SSL context explicitly
                     )
                     self._mysql_pools[db_url] = pool
                 except Exception as e:
                     logger.error(f"Failed to create MySQL pool: {e}")
                     raise e
+                    
         return self._mysql_pools[db_url]
 
     async def close_all(self):
@@ -119,11 +136,9 @@ class AsyncPoolManager:
             pool.close()
             await pool.wait_closed()
 
-# Initialize Global Async Manager
 async_pool_manager = AsyncPoolManager()
 
 def get_dynamic_thread_limit():
-    """Calculates a safe thread limit based on available RAM."""
     try:
         total_ram_bytes = psutil.virtual_memory().total
         total_cores = multiprocessing.cpu_count()
@@ -141,12 +156,10 @@ def get_dynamic_thread_limit():
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup
     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}. Async Drivers Ready.")
     yield
-    # Shutdown
     await async_pool_manager.close_all()
 
 app = FastAPI(title="Unified Data Executor API", lifespan=lifespan)
@@ -203,7 +216,6 @@ class RequestCoalescer:
             self._active_requests[key] = future
 
         try:
-            # Works for both async functions and run_in_threadpool
             result = await func(*args, **kwargs)
             if not future.done(): future.set_result(result)
             return result
@@ -230,10 +242,15 @@ def is_aggregate_query(query: str) -> bool:
     return False
 
 def normalize_mysql_uri(uri: str) -> str:
+    """
+    Normalizes the URI to ensure consistent cache keys.
+    DOES NOT REMOVE parameters (like ssl-mode) blindly, but sorts them.
+    """
     try:
         parsed_uri = urlparse(uri)
         query_params = parse_qs(parsed_uri.query)
-        query_params.pop('ssl-mode', None)
+        # Re-encode with sorting to ensure ?a=1&b=2 is same as ?b=2&a=1
+        # We do NOT remove ssl-mode here anymore, so the pool manager knows about it
         new_query = urlencode(query_params, doseq=True)
         parsed_uri = parsed_uri._replace(query=new_query)
         return urlunparse(parsed_uri)
@@ -253,7 +270,6 @@ def normalize_postgres_uri(uri: str) -> str:
 async def _execute_async_mysql(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
     try:
-        # Get pool
         pool = await async_pool_manager.get_mysql_pool(db_url)
         
         async with pool.acquire() as conn:
@@ -285,7 +301,6 @@ async def _execute_async_mysql(db_url: str, sql_query: str, max_rows: int = 20,
                 await cursor.execute(final_query)
                 results = await cursor.fetchall()
                 
-                # Check actual length
                 if is_limited_result and len(results) < max_rows:
                     is_limited_result = False
                     message = f"Returned {len(results)} rows."
@@ -333,7 +348,6 @@ async def _execute_async_postgres(db_url: str, sql_query: str, max_rows: int = 2
                     is_limited_result = True
                     message = f"Showing first {max_rows} rows."
 
-            # Fetch results
             records = await conn.fetch(final_query)
             results = [dict(r) for r in records]
             
@@ -405,7 +419,6 @@ async def execute_mysql_endpoint(query: SqlQueryRequest, token: str = Depends(va
         unique_params = {
             "db": normalized_url, "q": query.sql_query, "l": limit_val, "lim": query.limited
         }
-        # Direct Async Call (No Threadpool)
         result_dict = await coalescer.execute(
             "mysql", unique_params, _execute_async_mysql, 
             db_url=normalized_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
@@ -425,7 +438,6 @@ async def execute_postgres_endpoint(query: PgQueryRequest, token: str = Depends(
         unique_params = {
             "db": clean_url, "q": query.sql_query, "l": limit_val, "lim": query.limited
         }
-        # Direct Async Call (No Threadpool)
         result_dict = await coalescer.execute(
             "postgres", unique_params, _execute_async_postgres, 
             db_url=clean_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
@@ -446,7 +458,6 @@ async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(
             "uri": payload.mongo_uri, "db": payload.db_name, "col": payload.collection_name, 
             "q": parsed_query, "lim": payload.limited, "lrows": payload.limit_rows
         }
-        # Mongo is Sync via PyMongo -> Uses run_in_threadpool
         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, 
@@ -558,7 +569,7 @@ async def test_parallel_calc_endpoint(payload: TestCalcRequest):
 
 @app.get("/")
 async def root():
-    return {"message": "Code Execution Server is running"}
+    return {"message": "Python Code Execution Server is running"}
 
 @app.get("/ping")
 async def ping():