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

fixes #25

Author: iskakaushik

mcp_clickhouse/mcp_server.py

@@ -128,21 +128,36 @@ def execute_query(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_select_query(query: str):
     """Run a SELECT query in a ClickHouse database"""
     logger.info(f"Executing SELECT 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."
+        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)}"}
 
 
 def create_clickhouse_client():

tests/test_tool.py

@@ -71,8 +71,9 @@ 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_select_query(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."""