feat: Add query timeout and thread pool for SELECT queries (ClickHouse#20)

Author: iskakaushik

mcp_timeplus/mcp_server.py

@@ -1,5 +1,7 @@
 import logging
 from typing import Sequence
+import concurrent.futures
+import atexit
 
 import timeplus_connect
 from timeplus_connect.driver.binding import quote_identifier, format_query_value
@@ -20,6 +22,10 @@
 )
 logger = logging.getLogger(MCP_SERVER_NAME)
 
+QUERY_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=10)
+atexit.register(lambda: QUERY_EXECUTOR.shutdown(wait=True))
+SELECT_QUERY_TIMEOUT_SECS = 30
+
 load_dotenv()
 
 deps = [
@@ -113,6 +119,15 @@ def get_table_info(table):
 @mcp.tool()
 def run_sql(query: str):
     logger.info(f"Executing query: {query}")
+    future = QUERY_EXECUTOR.submit(execute_query, query)
+    try:
+        result = future.result(timeout=SELECT_QUERY_TIMEOUT_SECS)
+        return result
+    except concurrent.futures.TimeoutError:
+        logger.warning(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds: {query}")
+        future.cancel()
+        return f"Queries taking longer than {SELECT_QUERY_TIMEOUT_SECS} seconds are currently not supported."
+
     client = create_timeplus_client()
     try:
         readonly = 1 if config.readonly else 0