joi ma

controller.py

@@ -24,12 +24,15 @@ from anyio import to_thread
 from dotenv import load_dotenv
 from pydantic import BaseModel
 
-# --- Database Drivers ---
+# --- Async Database Drivers ---
+import asyncpg
+import asyncmy
+from asyncmy.cursors import DictCursor
+
+# --- Database Drivers (Kept for compatibility/types if needed, but unused in SQL routes) ---
 from bson import ObjectId
 import mysql.connector
-from mysql.connector import pooling
 import psycopg2
-from psycopg2 import pool as pg_pool
 from psycopg2.extras import RealDictCursor
 import uvicorn
 
@@ -53,6 +56,72 @@ logging.basicConfig(
 )
 logger = logging.getLogger("API_Controller")
 
+# ==============================================================================
+# ASYNC CONNECTION POOL MANAGER
+# ==============================================================================
+class AsyncPoolManager:
+    def __init__(self):
+        self._pg_pools: Dict[str, asyncpg.Pool] = {}
+        self._mysql_pools: Dict[str, asyncmy.Pool] = {}
+        self._lock = asyncio.Lock()
+
+    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]}...")
+                try:
+                    # asyncpg handles Keepalives and SSL automatically better than psycopg2
+                    pool = await asyncpg.create_pool(
+                        dsn=db_url,
+                        min_size=1,
+                        max_size=20,
+                        max_inactive_connection_lifetime=300,
+                        command_timeout=60
+                    )
+                    self._pg_pools[db_url] = pool
+                except Exception as e:
+                    logger.error(f"Failed to create PG pool: {e}")
+                    raise e
+        return self._pg_pools[db_url]
+
+    async def get_mysql_pool(self, db_url: str) -> asyncmy.Pool:
+        async with self._lock:
+            if db_url not in self._mysql_pools:
+                logger.info(f"Creating new AsyncMy pool for: {db_url[:20]}...")
+                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
+                
+                try:
+                    pool = await asyncmy.create_pool(
+                        user=parsed.username,
+                        password=parsed.password,
+                        host=parsed.hostname,
+                        port=parsed.port or 3306,
+                        db=parsed.path.lstrip("/"),
+                        minsize=1,
+                        maxsize=20,
+                        autocommit=True,
+                        pool_recycle=280, # Recycle before default 8hr timeout to prevent "Gone Away"
+                    )
+                    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):
+        logger.info("Closing all async database pools...")
+        for url, pool in self._pg_pools.items():
+            await pool.close()
+        for url, pool in self._mysql_pools.items():
+            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:
@@ -64,7 +133,6 @@ def get_dynamic_thread_limit():
         BYTES_PER_THREAD = 8 * 1024 * 1024 
         calculated_limit = int(safe_ram_pool / BYTES_PER_THREAD)
         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:
@@ -73,25 +141,23 @@ 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}.")
+    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)
 
 # ==============================================================================
-# BATCH MODELS
+# MODELS & MIDDLEWARE
 # ==============================================================================
 T = TypeVar("T")
+class BatchRequest(BaseModel, Generic[T]): requests: List[T]
+class BatchResponse(BaseModel, Generic[T]): responses: List[T]
 
-class BatchRequest(BaseModel, Generic[T]):
-    requests: List[T]
-
-class BatchResponse(BaseModel, Generic[T]):
-    responses: List[T]
-
-# --- Directory & CORS ---
 CHART_DIR = "generated_charts"
 os.makedirs(CHART_DIR, exist_ok=True)
 
@@ -106,7 +172,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Security ---
 security = HTTPBearer()
 API_SECRET_TOKEN = os.getenv("API_BEARER_TOKEN")
 
@@ -115,68 +180,6 @@ async def validate_token(credentials: HTTPAuthorizationCredentials = Depends(sec
         raise HTTPException(status_code=403, detail="Invalid Authentication Token")
     return credentials.credentials
 
-# ==============================================================================
-# CONNECTION POOL MANAGER (With Keepalives)
-# ==============================================================================
-class ConnectionPoolManager:
-    def __init__(self):
-        self._mysql_pools = {}
-        self._pg_pools = {}
-        self._lock = Lock() 
-
-    def get_mysql_connection(self, db_url: str):
-        with self._lock:
-            if db_url not in self._mysql_pools:
-                logger.info(f"Creating new MySQL pool for: {db_url}")
-                parsed = urlparse(db_url)
-                db_config = {
-                    "user": parsed.username,
-                    "password": parsed.password,
-                    "host": parsed.hostname,
-                    "port": parsed.port or 3306,
-                    "database": parsed.path.lstrip("/"),
-                    "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
-                )
-        return self._mysql_pools[db_url].get_connection()
-
-    def get_postgres_connection(self, db_url: str):
-        with self._lock:
-            if db_url not in self._pg_pools:
-                logger.info(f"Creating new Postgres pool for: {db_url}")
-                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,
-                    "database": parsed.path.lstrip("/"),
-                    "user": parsed.username,
-                    "password": parsed.password,
-                    "sslmode": sslmode,
-                    "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, 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
 # ==============================================================================
@@ -195,11 +198,12 @@ class RequestCoalescer:
         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()
+            
+            future = asyncio.get_running_loop().create_future()
             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
@@ -214,7 +218,7 @@ class RequestCoalescer:
 coalescer = RequestCoalescer()
 
 # ==============================================================================
-# DB EXECUTION LOGIC (WITH RETRY & SELF-HEALING)
+# ASYNC DB LOGIC
 # ==============================================================================
 
 def is_aggregate_query(query: str) -> bool:
@@ -245,163 +249,115 @@ 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:
+# --- ASYNC MYSQL EXECUTOR ---
+async def _execute_async_mysql(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
-    
-    # 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:
+        # Get pool
+        pool = await async_pool_manager.get_mysql_pool(db_url)
+        
+        async with pool.acquire() as conn:
+            async with conn.cursor(cursor=DictCursor) as cursor:
+                clean_query = sql_query.strip()
+                query_lower = clean_query.lower()
 
-            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 query_lower.startswith("select"):
+                    await cursor.execute(clean_query)
+                    return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}
 
-            if not limited:
-                cursor.execute(clean_query)
-                results = cursor.fetchall()
+                final_query = clean_query
                 is_aggregate = is_aggregate_query(clean_query)
-                message = f"Raw query executed. Returned {len(results)} row(s)."
                 is_limited_result = False
-            else:
-                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
+                message = ""
 
-            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
-            }
+                if not limited:
+                    message = f"Raw query executed."
+                else:
+                    if is_aggregate:
+                        message = f"Aggregate query completed."
+                    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}"
+                        is_limited_result = True
+                        message = f"Showing first {max_rows} rows."
+
+                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."
+                elif is_limited_result:
+                    message = f"Showing first {len(results)} rows."
 
-        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
-            
-            # 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
+                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 Exception as e:
+        return {"success": False, "error": str(e), "executionTime": 0.0}
 
-# --- SELF-HEALING POSTGRES EXECUTION ---
-def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
+# --- ASYNC POSTGRES EXECUTOR ---
+async def _execute_async_postgres(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
-    
-    # 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)
-            
+    try:
+        pool = await async_pool_manager.get_pg_pool(db_url)
+        
+        async with pool.acquire() as conn:
             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)
+                await conn.execute(clean_query)
                 return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}
-            
-            # Read Logic
+
+            final_query = clean_query
+            is_aggregate = is_aggregate_query(clean_query)
+            is_limited_result = False
+            message = ""
+
             if not limited:
-                cursor.execute(clean_query)
-                results = cursor.fetchall()
-                is_aggregate = is_aggregate_query(clean_query)
-                message = f"Raw query executed. Returned {len(results)} row(s)."
-                is_limited_result = False
+                message = f"Raw query executed."
             else:
-                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
+                if is_aggregate:
+                    message = f"Aggregate query completed."
                 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
-            }
+                    is_limited_result = True
+                    message = f"Showing first {max_rows} rows."
 
-        except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
-            # Bad Connection. DISCARD IT.
-            if connection:
-                pool_manager.return_postgres_connection(db_url, connection, close=True)
+            # Fetch results
+            records = await conn.fetch(final_query)
+            results = [dict(r) for r in records]
             
-            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}
+            if is_limited_result and len(results) < max_rows:
+                is_limited_result = False
+                message = f"Returned {len(results)} rows."
+            elif is_limited_result:
+                message = f"Showing first {len(results)} rows."
 
-        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
+            columns = list(results[0].keys()) if results else []
+
+            return {
+                "success": True, 
+                "results": jsonable_encoder(results, custom_encoder={uuid.UUID: str, ObjectId: str}), 
+                "columns": columns, "rowCount": len(results), 
+                "executionTime": time.time() - start_time,
+                "is_aggregate": is_aggregate, "limited": is_limited_result, 
+                "message": message, "error": None
+            }
+    except Exception as e:
+        return {"success": False, "error": str(e), "executionTime": 0.0}
 
 # ==============================================================================
-# API ROUTES
+# ROUTES
 # ==============================================================================
 
 class SqlQueryRequest(BaseModel):
@@ -449,8 +405,9 @@ 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, run_in_threadpool, _run_mysql_synchronously, 
+            "mysql", unique_params, _execute_async_mysql, 
             db_url=normalized_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
         )
         final_result = result_dict.copy()
@@ -468,8 +425,9 @@ 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, run_in_threadpool, _run_postgres_synchronously, 
+            "postgres", unique_params, _execute_async_postgres, 
             db_url=clean_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited
         )
         final_result = result_dict.copy()
@@ -488,6 +446,7 @@ 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, 

requirements.txt

@@ -19,4 +19,6 @@ asyncpg==0.30.0
 aiomysql==0.2.0
 requests==2.32.3
 psutil==6.1.1
-
+asyncpg==0.30.0
+asyncmy==0.2.10
+cryptography==43.0.3