Make include_plan_description configurable, so optionally sql query plans containing sensitive data can be masked

Author: zemin-piao

README.md

@@ -112,6 +112,7 @@ servers:
     auth:  # optional
       username: "user"
       password: "pass"
+    include_plan_description: false  # optional, whether to include SQL execution plans by default (default: false)
 mcp:
   transports:
     - streamable-http # streamable-http or stdio.
@@ -185,7 +186,7 @@ The MCP server provides **18 specialized tools** organized by analysis patterns.
 *SQL performance analysis and execution plan comparison*
 | 🔧 Tool | 📝 Description |
 |---------|----------------|
-| `list_slowest_sql_queries` | 🐌 Get the top N slowest SQL queries for an application with detailed execution metrics |
+| `list_slowest_sql_queries` | 🐌 Get the top N slowest SQL queries for an application with detailed execution metrics and optional plan descriptions |
 | `compare_sql_execution_plans` | 🔍 Compare SQL execution plans between two Spark jobs, analyzing logical/physical plans and execution metrics |
 
 ### 🚨 Performance & Bottleneck Analysis
@@ -302,6 +303,7 @@ SHS_SERVERS_*_AUTH_TOKEN - Token for a specific server
 SHS_SERVERS_*_VERIFY_SSL - Whether to verify SSL for a specific server (true/false)
 SHS_SERVERS_*_TIMEOUT - HTTP request timeout in seconds for a specific server (default: 30)
 SHS_SERVERS_*_EMR_CLUSTER_ARN - EMR cluster ARN for a specific server
+SHS_SERVERS_*_INCLUDE_PLAN_DESCRIPTION - Whether to include SQL execution plans by default for a specific server (true/false, default: false)
 ```
 
 ## 🤖 AI Agent Integration

src/spark_history_mcp/config/config.py

@@ -28,6 +28,7 @@ class ServerConfig(BaseSettings):
     emr_cluster_arn: Optional[str] = None  # EMR specific field
     use_proxy: bool = False
     timeout: int = 30  # HTTP request timeout in seconds
+    include_plan_description: bool = False
 
 
 class McpConfig(BaseSettings):

src/spark_history_mcp/tools/tools.py

@@ -922,7 +922,7 @@ def list_slowest_sql_queries(
     top_n: int = 1,
     page_size: int = 100,
     include_running: bool = False,
-    include_plan_description: bool = True,
+    include_plan_description: Optional[bool] = None,
     plan_description_max_length: int = 2000,
 ) -> List[SqlQuerySummary]:
     """
@@ -935,7 +935,7 @@ def list_slowest_sql_queries(
         top_n: Number of slowest queries to return (default: 1)
         page_size: Number of executions to fetch per page (default: 100)
         include_running: Whether to include running queries (default: False)
-        include_plan_description: Whether to include execution plans (default: True)
+        include_plan_description: Whether to include execution plans (uses server config if not specified)
         plan_description_max_length: Max characters for plan description (default: 1500)
 
     Returns:
@@ -944,6 +944,10 @@ def list_slowest_sql_queries(
     ctx = mcp.get_context()
     client = get_client_or_default(ctx, server, app_id)
 
+    # Use server config if include_plan_description not explicitly provided
+    if include_plan_description is None:
+        include_plan_description = client.config.include_plan_description
+
     all_executions: List[ExecutionData] = []
     offset = 0
 

tests/unit/test_tools.py

@@ -3,6 +3,7 @@
 from unittest.mock import MagicMock, patch
 
 from spark_history_mcp.api.spark_client import SparkRestClient
+from spark_history_mcp.config.config import ServerConfig
 from spark_history_mcp.models.spark_types import (
     ApplicationInfo,
     ExecutionData,
@@ -1087,3 +1088,63 @@ def test_get_slowest_sql_queries_limit(self, mock_get_client):
         self.assertEqual(result[0].duration, 10000)
         self.assertEqual(result[1].duration, 9000)
         self.assertEqual(result[2].duration, 8000)
+
+    @patch("spark_history_mcp.tools.tools.get_client_or_default")
+    def test_list_slowest_sql_queries_uses_server_config_for_plan_description(self, mock_get_client):
+        """Test that include_plan_description falls back to server config when not provided"""
+        # Setup mock client with server config
+        mock_client = MagicMock()
+        server_config = ServerConfig(url="http://test:18080", include_plan_description=False)
+        mock_client.config = server_config
+
+        # Create mock SQL execution
+        sql = MagicMock(spec=ExecutionData)
+        sql.id = 1
+        sql.duration = 5000
+        sql.status = "COMPLETED"
+        sql.success_job_ids = [1]
+        sql.failed_job_ids = []
+        sql.running_job_ids = []
+        sql.description = "Test Query"
+        sql.submission_time = datetime.now()
+        sql.plan_description = "Sample plan description"
+
+        mock_client.get_sql_list.return_value = [sql]
+        mock_get_client.return_value = mock_client
+
+        # Call function without include_plan_description parameter (should use server config)
+        result = list_slowest_sql_queries("spark-app-123")
+
+        # Verify plan description is empty due to server config setting False
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].plan_description, "")
+
+    @patch("spark_history_mcp.tools.tools.get_client_or_default")
+    def test_list_slowest_sql_queries_explicit_override_server_config(self, mock_get_client):
+        """Test that explicit include_plan_description parameter overrides server config"""
+        # Setup mock client with server config set to False
+        mock_client = MagicMock()
+        server_config = ServerConfig(url="http://test:18080", include_plan_description=False)
+        mock_client.config = server_config
+
+        # Create mock SQL execution
+        sql = MagicMock(spec=ExecutionData)
+        sql.id = 1
+        sql.duration = 5000
+        sql.status = "COMPLETED"
+        sql.success_job_ids = [1]
+        sql.failed_job_ids = []
+        sql.running_job_ids = []
+        sql.description = "Test Query"
+        sql.submission_time = datetime.now()
+        sql.plan_description = "Sample plan description"
+
+        mock_client.get_sql_list.return_value = [sql]
+        mock_get_client.return_value = mock_client
+
+        # Call function with explicit include_plan_description=True (should override server config)
+        result = list_slowest_sql_queries("spark-app-123", include_plan_description=True)
+
+        # Verify plan description is included despite server config being False
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].plan_description, "Sample plan description")