import base64 from contextlib import asynccontextmanager import logging import time import uuid import os import re 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 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 from dotenv import load_dotenv from pydantic import BaseModel # --- Database Drivers --- 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 # --- 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 pydantic_csv_analysis_model import AnalysisRequest, AnalysisResponse from pydantic_csv_charts_model import ChartExecutionPayload, ChartExecutionResponse from pydantic_mongo_executor_model import ExecutorPayload, ExecutorResponse from report_service import FileBoxProps, ReportRequest, execute_report_generation from supabase_service import upload_bytes_to_supabase # --- Configuration & Setup --- load_dotenv() logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO ) logger = logging.getLogger("API_Controller") 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() num_workers = max(1, total_cores - 2) ram_per_worker = total_ram_bytes / num_workers safe_ram_pool = ram_per_worker * 0.70 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: logger.warning(f"Failed to calculate dynamic threads ({e}). Fallback to 1000.") return 1000 @asynccontextmanager async def lifespan(app: FastAPI): 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 app = FastAPI(title="Unified Data Executor API", lifespan=lifespan) # ============================================================================== # BATCH MODELS # ============================================================================== T = TypeVar("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) origins_env = os.getenv("ALLOWED_ORIGINS", "*") ORIGINS = [origin.strip() for origin in origins_env.split(",")] app.add_middleware( CORSMiddleware, allow_origins=ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --- Security --- security = HTTPBearer() API_SECRET_TOKEN = os.getenv("API_BEARER_TOKEN") 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 # ============================================================================== # 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 # ============================================================================== class RequestCoalescer: def __init__(self): self._active_requests: Dict[str, asyncio.Future] = {} self._lock = asyncio.Lock() def _generate_key(self, prefix: str, data: dict) -> str: json_str = json.dumps(data, sort_keys=True, default=str) raw_str = f"{prefix}:{json_str}" return hashlib.md5(raw_str.encode()).hexdigest() 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: result = await func(*args, **kwargs) if not future.done(): future.set_result(result) return result except Exception as e: if not future.done(): future.set_exception(e) raise e finally: async with self._lock: if key in self._active_requests: del self._active_requests[key] coalescer = RequestCoalescer() # ============================================================================== # 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'] for pattern in aggregate_patterns: if re.search(pattern, query_lower): return True return False def normalize_mysql_uri(uri: str) -> str: try: parsed_uri = urlparse(uri) query_params = parse_qs(parsed_uri.query) query_params.pop('ssl-mode', None) new_query = urlencode(query_params, doseq=True) parsed_uri = parsed_uri._replace(query=new_query) return urlunparse(parsed_uri) except Exception: return uri def normalize_postgres_uri(uri: str) -> str: try: parsed_uri = urlparse(uri) if parsed_uri.scheme == 'postgres': parsed_uri = parsed_uri._replace(scheme='postgresql') return urlunparse(parsed_uri) 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() # 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() 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() 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 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 # 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() # 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) 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() 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 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 # ============================================================================== 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 @app.post("/api/execute_sql_query", response_model=SqlQueryResponse) async def execute_mysql_endpoint(query: SqlQueryRequest, token: str = Depends(validate_token)): request_id = str(uuid.uuid4())[:8] try: normalized_url = normalize_mysql_uri(query.database_url) limit_val = query.limit_rows if query.limit_rows is not None else 20 unique_params = { "db": normalized_url, "q": query.sql_query, "l": limit_val, "lim": query.limited } result_dict = await coalescer.execute( "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: raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id}) @app.post("/api/execute_postgres_query", response_model=PgQueryResponse) async def execute_postgres_endpoint(query: PgQueryRequest, token: str = Depends(validate_token)): request_id = str(uuid.uuid4())[:8] 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 } result_dict = await coalescer.execute( "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: raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id}) @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] try: result = await run_in_threadpool(execute_analysis_logic, code=payload.code, csv_url=payload.csv_url) return AnalysisResponse(success=result["success"], output_log=result["output_log"], results=result["results"], error=result["error"], request_id=request_id) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/generate_report", response_model=FileBoxProps) async def generate_report_endpoint(payload: ReportRequest, token: str = Depends(validate_token)): try: result = await execute_report_generation(code=payload.code, csv_url=payload.csv_url, chat_id=payload.chat_id) return result except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/get_csv_info", response_model=CsvInfoResponse) async def get_csv_info_endpoint(payload: CsvInfoRequest, token: str = Depends(validate_token)): request_id = str(uuid.uuid4())[:8] start_time = time.time() try: info_result = await run_in_threadpool(get_csv_basic_info, csv_path=payload.csv_url) if "error" in info_result: return CsvInfoResponse(success=False, error=info_result["error"], request_id=request_id, duration=time.time() - start_time) return CsvInfoResponse(success=True, data=info_result, request_id=request_id, duration=time.time() - start_time) except Exception as e: raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id}) @app.post("/api/csv_data") async def get_csv_data_endpoint(payload: CsvDataRequest, token: str = Depends(validate_token)): try: result = await run_in_threadpool(get_robust_csv_rows, csv_url=payload.csv_url) if isinstance(result, dict) and "error" in result: raise HTTPException(status_code=400, detail=result["error"]) return result except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.post("/api/execute_python", response_model=PythonExecutionResponse) async def execute_python_endpoint(payload: PythonExecutionRequest, token: str = Depends(validate_token)): request_id = str(uuid.uuid4())[:8] try: execution_result = await run_in_threadpool(execute_python_logic, code=payload.code, custom_context=payload.context) return PythonExecutionResponse(success=execution_result['error'] is None, output=execution_result['output'], result=jsonable_encoder(execution_result['result']), isStructured=execution_result['isStructured'], error=execution_result['error'], request_id=request_id) except Exception as e: raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id}) # --- Batch Handlers --- async def batch_parallel_handler(func, requests: List[Any], token: str): tasks = [func(req, token) for req in requests] results = await asyncio.gather(*tasks, return_exceptions=True) return [res if not isinstance(res, Exception) else {"success": False, "error": str(res)} for res in results] @app.post("/api/batch/execute_sql_query", response_model=BatchResponse[SqlQueryResponse]) async def batch_execute_sql(payload: BatchRequest[SqlQueryRequest], token: str = Depends(validate_token)): responses = await batch_parallel_handler(execute_mysql_endpoint, payload.requests, token) return BatchResponse(responses=responses) @app.post("/api/batch/execute_postgres_query", response_model=BatchResponse[PgQueryResponse]) async def batch_execute_pg(payload: BatchRequest[PgQueryRequest], token: str = Depends(validate_token)): responses = await batch_parallel_handler(execute_postgres_endpoint, payload.requests, token) return BatchResponse(responses=responses) @app.post("/api/batch/execute_mongo", response_model=BatchResponse[ExecutorResponse]) async def batch_execute_mongo(payload: BatchRequest[ExecutorPayload], token: str = Depends(validate_token)): responses = await batch_parallel_handler(execute_mongo_endpoint, payload.requests, token) 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)) @app.get("/") async def root(): return {"message": "Python Code Execution Server is running"} @app.get("/ping") async def ping(): return {"message": "I am alive!"} if __name__ == "__main__": host = os.getenv("HOST", "0.0.0.0") port = int(os.getenv("PORT", 7860)) num_workers = max(1, multiprocessing.cpu_count() - 2) print(f"Starting production server on {host}:{port} with {num_workers} workers...") uvicorn.run("controller:app", host=host, port=port, workers=num_workers, loop="asyncio")