go

controller.py

@@ -5,10 +5,11 @@ 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
+import asyncio
+import multiprocessing
+import hashlib
+import json
+from typing import List, Optional, Dict, Any, TypeVar, Generic
 
 # --- FastAPI & Core ---
 from fastapi import FastAPI, HTTPException, Depends
@@ -16,15 +17,20 @@ 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 ---
 from bson import ObjectId
-import mysql.connector
-import psycopg2
 from psycopg2.extras import RealDictCursor
+from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
+
+# --- Database Pool ---
+from mysql.connector import pooling
+from psycopg2 import pool as pg_pool
+from threading import Lock
+
 import uvicorn
 
 # --- Existing Services ---
@@ -53,7 +59,6 @@ async def lifespan(app: FastAPI):
     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)", 
@@ -104,7 +109,6 @@ async def validate_token(credentials: HTTPAuthorizationCredentials = Depends(sec
 # PYDANTIC MODELS
 # ==============================================================================
 
-# --- MySQL Models ---
 class SqlQueryRequest(BaseModel):
     database_url: str
     sql_query: str
@@ -123,7 +127,6 @@ class SqlQueryResponse(BaseModel):
     limited: bool = False
     message: Optional[str] = None
 
-# --- PostgreSQL Models ---
 class PgQueryRequest(BaseModel):
     database_url: str
     sql_query: str
@@ -141,14 +144,107 @@ class PgQueryResponse(BaseModel):
     is_aggregate: bool = False
     limited: bool = False
     message: Optional[str] = None
+    
+# ==============================================================================
+# CONNECTION POOL MANAGER
+# ==============================================================================
+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
+                }
+                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
+                }
+                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):
+        if db_url in self._pg_pools and conn:
+            self._pg_pools[db_url].putconn(conn)
+
+pool_manager = ConnectionPoolManager()
+
+
+# ==============================================================================
+# REQUEST COALESCER (SINGLEFLIGHT)
+# ==============================================================================
+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:
+            # Execute the function (run_in_threadpool) with args/kwargs
+            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()
 
 
 # ==============================================================================
-# SHARED HELPER FUNCTIONS
+# DB EXECUTION LOGIC
 # ==============================================================================
 
 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*\(', 
@@ -160,10 +256,6 @@ def is_aggregate_query(query: str) -> bool:
             return True
     return False
 
-# ==============================================================================
-# MYSQL LOGIC
-# ==============================================================================
-
 def normalize_mysql_uri(uri: str) -> str:
     try:
         parsed_uri = urlparse(uri)
@@ -175,6 +267,15 @@ def normalize_mysql_uri(uri: str) -> str:
     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
+
 def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
     connection = None
@@ -182,13 +283,7 @@ def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, li
     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)
+        connection = pool_manager.get_mysql_connection(db_url)
         cursor = connection.cursor(dictionary=True)
         
         clean_query = sql_query.strip()
@@ -219,7 +314,6 @@ def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, li
                 
                 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
@@ -234,20 +328,7 @@ def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, li
         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
+        if connection: connection.close() # Returns to pool
 
 def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
     start_time = time.time()
@@ -255,19 +336,17 @@ def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20,
     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)
+        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
+            
         if not limited:
             cursor.execute(clean_query)
             results = cursor.fetchall()
@@ -291,6 +370,7 @@ def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20,
                 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})
@@ -301,18 +381,7 @@ def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20,
         return response
     finally:
         if cursor: cursor.close()
-        if connection: connection.close()
-        
-# ==============================================================================
-# KEEP-ALIVE ROUTE
-# ==============================================================================
-@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 connection: pool_manager.return_postgres_connection(db_url, connection)
 
 
 # ==============================================================================
@@ -324,9 +393,43 @@ async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(
     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)
-        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)
+        
+        # --- 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
+        }
+
+        # 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))
 
@@ -334,33 +437,18 @@ async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(
 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
-            )
+            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
-            )
-
+            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))
 
@@ -370,10 +458,34 @@ 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
-        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)
+        
+        # Unique ID for Coalescing
+        unique_params = {
+            "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
+        )
+
+        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)
@@ -382,12 +494,34 @@ 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
-        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)
+        
+        unique_params = {
+            "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
+        )
+
+        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_csv_analysis", response_model=AnalysisResponse)
 async def execute_analysis_endpoint(payload: AnalysisRequest, token: str = Depends(validate_token)):
     request_id = str(uuid.uuid4())[:8]
@@ -435,15 +569,8 @@ async def execute_python_endpoint(payload: PythonExecutionRequest, token: str =
     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)
-# ==============================================================================
-
+# --- Batch Handlers ---
 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]
@@ -463,6 +590,20 @@ async def batch_execute_mongo(payload: BatchRequest[ExecutorPayload], token: str
     responses = await batch_parallel_handler(execute_mongo_endpoint, payload.requests, token)
     return BatchResponse(responses=responses)
 
+# --- Test Endpoint ---
+
+# ==============================================================================
+# KEEP-ALIVE ROUTE
+# ==============================================================================
+@app.get("/")
+async def root():
+    return {"message": "Python Code Execution Server is running"}
+
+@app.get("/ping")
+async def ping():
+    return {"message": "I am alive!"}
+
+
 # ==============================================================================
 # HIGH PERFORMANCE SERVER EXECUTION
 # ==============================================================================
@@ -484,6 +625,4 @@ if __name__ == "__main__":
         port=port, 
         workers=num_workers,
         loop="asyncio",
-        log_level="info",
-        access_log=False
     )

mongo_service.py

@@ -1,12 +1,11 @@
-
 import logging
 from typing import Dict, List, Union
 from bson import ObjectId
 import ast
 import json
 import re
-from pymongo.collation import Collation
 from pymongo import MongoClient
+from threading import Lock # Needed for thread-safe pooling
 
 logging.basicConfig(
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
@@ -14,6 +13,42 @@ logging.basicConfig(
 )
 logger = logging.getLogger("mongo_service")
 
+# ==============================================================================
+# MONGO CONNECTION POOL MANAGER
+# ==============================================================================
+class MongoConnectionManager:
+    def __init__(self):
+        self._clients: Dict[str, MongoClient] = {}
+        self._lock = Lock()
+
+    def get_client(self, uri: str) -> MongoClient:
+        """
+        Returns a cached MongoClient instance.
+        If it doesn't exist, creates one with a connection pool.
+        """
+        # Double-checked locking for performance
+        if uri not in self._clients:
+            with self._lock:
+                if uri not in self._clients:
+                    logger.info(f"Initialize new MongoDB Client for URI: {uri[:20]}...")
+                    # maxPoolSize=50 allows 50 concurrent ops per URI. 
+                    # The rest will wait in queue automatically.
+                    self._clients[uri] = MongoClient(
+                        uri, 
+                        serverSelectionTimeoutMS=5000, 
+                        connectTimeoutMS=5000,
+                        maxPoolSize=50, 
+                        minPoolSize=1
+                    )
+        return self._clients[uri]
+
+# Global Manager Instance
+mongo_manager = MongoConnectionManager()
+
+# ==============================================================================
+# HELPER FUNCTIONS (Unchanged)
+# ==============================================================================
+
 def convert_oid(obj):
     """Recursively convert $oid format or ObjectId objects to strings."""
     if isinstance(obj, ObjectId):
@@ -34,23 +69,19 @@ def parse_query_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
 
     query_str = str(query_input).strip()
     
-    # 1. Try JSON
     try:
         return json.loads(query_str)
     except json.JSONDecodeError:
         pass
 
-    # 2. Extract code block
     match = re.search(r"```(?:python|json|javascript)?\s*(.*?)\s*```", query_str, re.DOTALL)
     if match:
         query_str = match.group(1).strip()
 
-    # 3. Clean up MongoDB shell syntax
     query_str = re.sub(r'^db\.\w+\.\w+\(', '', query_str)
     if query_str.endswith(')'):
         query_str = query_str[:-1]
 
-    # 4. AST Literal Eval
     try:
         parsed = ast.literal_eval(query_str)
         return convert_oid(parsed)
@@ -59,28 +90,30 @@ def parse_query_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
         raise ValueError(f"Could not parse query string: {str(e)}")
 
 def get_value_ignore_case(d: Dict, keys: List[str], default=None):
-    """Helper to get value checking multiple key variations (camelCase/snake_case)."""
     for k in keys:
         if k in d:
             return d[k]
     return default
 
+# ==============================================================================
+# MAIN EXECUTION LOGIC (Modified)
+# ==============================================================================
+
 def execute_mongo_operation(
     mongo_uri: str, 
     db_name: str, 
     collection_name: str, 
     query: Union[Dict, List],
-    limited: bool = False,   # <--- Added
-    limit_rows: int = 20     # <--- Added
+    limited: bool = False,
+    limit_rows: int = 20
 ):
     """
-    Generic MongoDB Executor.
-    Supports: Aggregation Pipelines (List) and Find Operations (Dict).
-    Now supports explicit row limiting.
+    Executes MongoDB operations using the shared connection pool.
     """
-    client = None
+    # 1. GET CLIENT FROM MANAGER (Do NOT create new MongoClient here)
+    client = mongo_manager.get_client(mongo_uri)
+    
     try:
-        client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=5000)
         db = client[db_name]
         collection = db[collection_name]
         
@@ -88,13 +121,9 @@ def execute_mongo_operation(
         
         # --- Aggregation ---
         if isinstance(query, list):
-            logger.info("Executing Aggregation Pipeline")
-            
             # Apply Limit if requested and not already present at the end
             if limited:
-                # Check if the last stage is already a $limit
                 if not (query and "$limit" in query[-1]):
-                    logger.info(f"Injecting $limit: {limit_rows} to pipeline")
                     query.append({"$limit": limit_rows})
             
             cursor = collection.aggregate(query, allowDiskUse=True)
@@ -102,20 +131,16 @@ def execute_mongo_operation(
             
         # --- Find / Command ---
         elif isinstance(query, dict):
-            logger.info("Executing Structured Find Query")
-            
             command_keys = {'filter', 'query', '$query', 'projection', 'sort', 'limit', 'skip'}
             has_command_keys = bool(set(query.keys()) & command_keys)
             
             cursor = None
 
             if has_command_keys and ('filter' in query or 'query' in query or '$query' in query):
-                # It is a Structured Command
                 query_filter = get_value_ignore_case(query, ['filter', 'query', '$query'], {})
                 projection = get_value_ignore_case(query, ['projection', 'fields'], None)
                 if projection == {}: projection = None 
                 
-                # Check internal limit
                 internal_limit = int(get_value_ignore_case(query, ['limit'], 0))
                 skip = int(get_value_ignore_case(query, ['skip'], 0))
                 sort_val = get_value_ignore_case(query, ['sort', '$orderby'], None)
@@ -129,18 +154,13 @@ def execute_mongo_operation(
                 
                 if skip > 0: cursor = cursor.skip(skip)
                 
-                # Logic: If 'limited' is forced (Chat Mode), use limit_rows. 
-                # Otherwise use the query's internal limit if it exists.
                 if limited:
                     cursor = cursor.limit(limit_rows)
                 elif internal_limit > 0: 
                     cursor = cursor.limit(internal_limit)
 
             else:
-                # It is a Raw Filter
-                logger.info("Treating input dictionary as Raw Filter")
                 cursor = collection.find(query)
-                
                 if limited:
                     cursor = cursor.limit(limit_rows)
             
@@ -153,6 +173,7 @@ def execute_mongo_operation(
     except Exception as e:
         logger.error(f"DB Execution Error: {e}")
         raise e
-    finally:
-        if client:
-            client.close()
+    
+    # IMPORTANT: DO NOT CLOSE THE CLIENT
+    # finally:
+    #     if client: client.close()  <-- REMOVED