import base64 from contextlib import asynccontextmanager import logging import time import uuid import os import re import asyncio # Added for parallel batching import multiprocessing # Added for worker calculation from typing import List, Optional, Dict, Any, TypeVar, Generic # Added Generic/TypeVar from urllib.parse import urlparse, parse_qs, urlencode, urlunparse # --- 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 # Added for AnyIO 4.x concurrency tuning from dotenv import load_dotenv from pydantic import BaseModel # --- Database Drivers --- from bson import ObjectId import mysql.connector import psycopg2 from psycopg2.extras import RealDictCursor # --- 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") @asynccontextmanager async def lifespan(app: FastAPI): # Startup logic to_thread.current_default_thread_limiter().total_tokens = 2000 logger.info("Worker Process Started: Thread pool capacity set to 2000.") yield # Shutdown logic (if any) goes here app = FastAPI( title="Unified Data Executor API (Mongo, SQL, CSV)", lifespan=lifespan ) # ============================================================================== # HIGH-CONCURRENCY BATCH MODELS & STARTUP # ============================================================================== T = TypeVar("T") class BatchRequest(BaseModel, Generic[T]): requests: List[T] class BatchResponse(BaseModel, Generic[T]): responses: List[T] # --- Directory Setup --- 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(",")] 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") 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 # ============================================================================== # --- MySQL 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 # --- PostgreSQL Models --- 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 # ============================================================================== # SHARED HELPER FUNCTIONS # ============================================================================== def is_aggregate_query(query: str) -> bool: """Checks for aggregate keywords.""" 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 # ============================================================================== # MYSQL LOGIC # ============================================================================== 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 _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": ""} try: 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 } connection = mysql.connector.connect(**db_config) 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): cursor.execute(clean_query) results = cursor.fetchall() response["message"] = f"Aggregate query completed. Returned {len(results)} row(s)." response["is_aggregate"] = True 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: response["error"] = str(e) return response finally: if cursor: cursor.close() if connection and connection.is_connected(): connection.close() # ============================================================================== # POSTGRES LOGIC # ============================================================================== 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 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": ""} try: 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} connection = psycopg2.connect(**db_config) 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 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): 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 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: if connection: connection.rollback() response["error"] = str(e) return response finally: if cursor: cursor.close() if connection: connection.close() # ============================================================================== # 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: parsed_query = parse_query_input(payload.generated_query) result_data = await 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: # 1. Execute Code 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) # 2. Handle Output Format if payload.return_base64: # OPTION A: Return Base64 (No Supabase Upload) 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: # OPTION B: Upload to Supabase (Standard behavior) 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_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 result_dict = await run_in_threadpool(_run_mysql_synchronously, db_url=normalized_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited) result_dict["request_id"] = request_id return SqlQueryResponse(**result_dict) 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 result_dict = await run_in_threadpool(_run_postgres_synchronously, db_url=clean_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited) result_dict["request_id"] = request_id return PgQueryResponse(**result_dict) except Exception as e: raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id}) @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}) # ============================================================================== # NEW: BATCH HANDLER LOGIC (Scaling for 1000+ Requests) # ============================================================================== async def batch_parallel_handler(func, requests: List[Any], token: str): """ Executes multiple requests in parallel within a single worker process using asyncio.gather, utilizing the high-token thread pool. """ 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) # ============================================================================== # HIGH PERFORMANCE SERVER EXECUTION # ============================================================================== if __name__ == "__main__": import uvicorn host = os.getenv("HOST", "0.0.0.0") port = int(os.getenv("PORT", 8000)) # Scale processes: 16 Cores - 2 = 14 Workers 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="auto", )