Port a1bdc6 to1dc0c4 (#11)

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

* add action to publish to pypi (ClickHouse#19)

* add client_name with mcp_clickhouse (ClickHouse#21)

* Add descriptions for tools required by Bedrock (ClickHouse#23)

Addresses ClickHouse#22

* 0.1.5

* fix: prevent BrokenResourceError by returning structured responses for query errors (ClickHouse#26)

fixes ClickHouse#25

* Update README.md - Added link to the Youtube overview (ClickHouse#27)

* fix cherry pick issue and add desc

* Update publish.yml

* change mcp client name

---------

Co-authored-by: Kaushik Iska <[email protected]>
Co-authored-by: Philip Dubé <[email protected]>
Co-authored-by: Mikhail Shustov <[email protected]>
Co-authored-by: Dmitry Pavlov <[email protected]>
Co-authored-by: Philip Dubé <[email protected]>
Co-authored-by: Ryadh DAHIMENE <[email protected]>

Author: 6 people

.github/workflows/publish.yml

@@ -0,0 +1,18 @@
+on:
+  workflow_dispatch:
+
+jobs:
+  publish:
+    name: Upload release to PyPI
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: "https://pypi.org/p/mcp-timeplus"
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv python install
+      - run: uv build
+      - uses: pypa/gh-action-pypi-publish@release/v1

mcp_timeplus/mcp_env.py

@@ -121,6 +121,7 @@ def get_client_config(self) -> dict:
             "verify": self.verify,
             "connect_timeout": self.connect_timeout,
             "send_receive_timeout": self.send_receive_timeout,
+            "client_name": "mcp_timeplus",
         }
 
         # Add optional database if set

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 = [
@@ -35,6 +41,7 @@
 
 @mcp.tool()
 def list_databases():
+    """List available Timeplus databases"""
     logger.info("Listing all databases")
     client = create_timeplus_client()
     result = client.command("SHOW DATABASES")
@@ -44,6 +51,7 @@ def list_databases():
 
 @mcp.tool()
 def list_tables(database: str = 'default', like: str = None):
+    """List available tables/streams in the given database"""
     logger.info(f"Listing tables in database '{database}'")
     client = create_timeplus_client()
     query = f"SHOW STREAMS FROM {quote_identifier(database)}"
@@ -109,10 +117,7 @@ def get_table_info(table):
     logger.info(f"Found {len(tables)} tables")
     return tables
 
-
-@mcp.tool()
-def run_sql(query: str):
-    logger.info(f"Executing query: {query}")
+def execute_query(query: str):
     client = create_timeplus_client()
     try:
         readonly = 1 if config.readonly else 0
@@ -128,7 +133,36 @@ def run_sql(query: str):
         return rows
     except Exception as err:
         logger.error(f"Error executing query: {err}")
-        return f"error running query: {err}"
+        # Return a structured dictionary rather than a string to ensure proper serialization
+        # by the MCP protocol. String responses for errors can cause BrokenResourceError.
+        return {"error": str(err)}
+
+@mcp.tool()
+def run_sql(query: str):
+    """Run a query in a Timeplus database"""
+    logger.info(f"Executing query: {query}")
+    try:
+        future = QUERY_EXECUTOR.submit(execute_query, query)
+        try:
+            result = future.result(timeout=SELECT_QUERY_TIMEOUT_SECS)
+            # Check if we received an error structure from execute_query
+            if isinstance(result, dict) and "error" in result:
+                logger.warning(f"Query failed: {result['error']}")
+                # MCP requires structured responses; string error messages can cause
+                # serialization issues leading to BrokenResourceError
+                return {"status": "error", "message": f"Query failed: {result['error']}"}
+            return result
+        except concurrent.futures.TimeoutError:
+            logger.warning(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds: {query}")
+            future.cancel()
+            # Return a properly structured response for timeout errors
+            return {"status": "error", "message": f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds"}
+    except Exception as e:
+        logger.error(f"Unexpected error in run_select_query: {str(e)}")
+        # Catch all other exceptions and return them in a structured format
+        # to prevent MCP serialization failures
+        return {"status": "error", "message": f"Unexpected error: {str(e)}"}
+
 
 @mcp.prompt()
 def generate_sql(requirements: str) -> str:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp-timeplus"
-version = "0.1.4"
+version = "0.1.5"
 description = "An MCP server for Timeplus."
 readme = "README.md"
 license = "Apache-2.0"

tests/test_tool.py

@@ -70,8 +70,10 @@ def test_run_select_query_failure(self):
         """Test running a SELECT query with an error."""
         query = f"SELECT * FROM {self.test_db}.non_existent_table"
         result = run_sql(query)
-        self.assertIsInstance(result, str)
-        self.assertIn("error running query", result)
+        self.assertIsInstance(result, dict)
+        self.assertEqual(result["status"], "error")
+        self.assertIn("Query failed", result["message"])
+
 
     def test_table_and_column_comments(self):
         """Test that table and column comments are correctly retrieved."""