Skip to content

Commit d0af0f2

Browse files
authored
Make include_plan_description configurable (#118)
Signed-off-by: zemin-piao <[email protected]>
1 parent a7ccd60 commit d0af0f2

File tree

4 files changed

+83
-3
lines changed

4 files changed

+83
-3
lines changed

README.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ servers:
112112
auth: # optional
113113
username: "user"
114114
password: "pass"
115+
include_plan_description: false # optional, whether to include SQL execution plans by default (default: false)
115116
mcp:
116117
transports:
117118
- streamable-http # streamable-http or stdio.
@@ -185,7 +186,7 @@ The MCP server provides **18 specialized tools** organized by analysis patterns.
185186
*SQL performance analysis and execution plan comparison*
186187
| 🔧 Tool | 📝 Description |
187188
|---------|----------------|
188-
| `list_slowest_sql_queries` | 🐌 Get the top N slowest SQL queries for an application with detailed execution metrics |
189+
| `list_slowest_sql_queries` | 🐌 Get the top N slowest SQL queries for an application with detailed execution metrics and optional plan descriptions |
189190
| `compare_sql_execution_plans` | 🔍 Compare SQL execution plans between two Spark jobs, analyzing logical/physical plans and execution metrics |
190191

191192
### 🚨 Performance & Bottleneck Analysis
@@ -302,6 +303,7 @@ SHS_SERVERS_*_AUTH_TOKEN - Token for a specific server
302303
SHS_SERVERS_*_VERIFY_SSL - Whether to verify SSL for a specific server (true/false)
303304
SHS_SERVERS_*_TIMEOUT - HTTP request timeout in seconds for a specific server (default: 30)
304305
SHS_SERVERS_*_EMR_CLUSTER_ARN - EMR cluster ARN for a specific server
306+
SHS_SERVERS_*_INCLUDE_PLAN_DESCRIPTION - Whether to include SQL execution plans by default for a specific server (true/false, default: false)
305307
```
306308
307309
## 🤖 AI Agent Integration

src/spark_history_mcp/config/config.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ class ServerConfig(BaseSettings):
2828
emr_cluster_arn: Optional[str] = None # EMR specific field
2929
use_proxy: bool = False
3030
timeout: int = 30 # HTTP request timeout in seconds
31+
include_plan_description: Optional[bool] = None
3132

3233

3334
class McpConfig(BaseSettings):

src/spark_history_mcp/tools/tools.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -922,7 +922,7 @@ def list_slowest_sql_queries(
922922
top_n: int = 1,
923923
page_size: int = 100,
924924
include_running: bool = False,
925-
include_plan_description: bool = True,
925+
include_plan_description: Optional[bool] = None,
926926
plan_description_max_length: int = 2000,
927927
) -> List[SqlQuerySummary]:
928928
"""
@@ -935,7 +935,7 @@ def list_slowest_sql_queries(
935935
top_n: Number of slowest queries to return (default: 1)
936936
page_size: Number of executions to fetch per page (default: 100)
937937
include_running: Whether to include running queries (default: False)
938-
include_plan_description: Whether to include execution plans (default: True)
938+
include_plan_description: Whether to include execution plans (uses server config if not specified)
939939
plan_description_max_length: Max characters for plan description (default: 1500)
940940
941941
Returns:
@@ -944,6 +944,12 @@ def list_slowest_sql_queries(
944944
ctx = mcp.get_context()
945945
client = get_client_or_default(ctx, server, app_id)
946946

947+
# Config takes priority: if config is set (True/False), use it; otherwise default to True
948+
if client.config.include_plan_description is not None:
949+
include_plan_description = client.config.include_plan_description
950+
else:
951+
include_plan_description = True
952+
947953
all_executions: List[ExecutionData] = []
948954
offset = 0
949955

tests/unit/test_tools.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from unittest.mock import MagicMock, patch
44

55
from spark_history_mcp.api.spark_client import SparkRestClient
6+
from spark_history_mcp.config.config import ServerConfig
67
from spark_history_mcp.models.spark_types import (
78
ApplicationInfo,
89
ExecutionData,
@@ -1087,3 +1088,73 @@ def test_get_slowest_sql_queries_limit(self, mock_get_client):
10871088
self.assertEqual(result[0].duration, 10000)
10881089
self.assertEqual(result[1].duration, 9000)
10891090
self.assertEqual(result[2].duration, 8000)
1091+
1092+
@patch("spark_history_mcp.tools.tools.get_client_or_default")
1093+
def test_list_slowest_sql_queries_uses_server_config_for_plan_description(
1094+
self, mock_get_client
1095+
):
1096+
"""Test that include_plan_description falls back to server config when not provided"""
1097+
# Setup mock client with server config
1098+
mock_client = MagicMock()
1099+
server_config = ServerConfig(
1100+
url="http://test:18080", include_plan_description=False
1101+
)
1102+
mock_client.config = server_config
1103+
1104+
# Create mock SQL execution
1105+
sql = MagicMock(spec=ExecutionData)
1106+
sql.id = 1
1107+
sql.duration = 5000
1108+
sql.status = "COMPLETED"
1109+
sql.success_job_ids = [1]
1110+
sql.failed_job_ids = []
1111+
sql.running_job_ids = []
1112+
sql.description = "Test Query"
1113+
sql.submission_time = datetime.now()
1114+
sql.plan_description = "Sample plan description"
1115+
1116+
mock_client.get_sql_list.return_value = [sql]
1117+
mock_get_client.return_value = mock_client
1118+
1119+
# Call function without include_plan_description parameter (should use server config)
1120+
result = list_slowest_sql_queries("spark-app-123")
1121+
1122+
# Verify plan description is empty due to server config setting False
1123+
self.assertEqual(len(result), 1)
1124+
self.assertEqual(result[0].plan_description, "")
1125+
1126+
@patch("spark_history_mcp.tools.tools.get_client_or_default")
1127+
def test_list_slowest_sql_queries_explicit_override_server_config(
1128+
self, mock_get_client
1129+
):
1130+
"""Test that server config overrides parameter when config is set"""
1131+
# Setup mock client with server config set to False
1132+
mock_client = MagicMock()
1133+
server_config = ServerConfig(
1134+
url="http://test:18080", include_plan_description=False
1135+
)
1136+
mock_client.config = server_config
1137+
1138+
# Create mock SQL execution
1139+
sql = MagicMock(spec=ExecutionData)
1140+
sql.id = 1
1141+
sql.duration = 5000
1142+
sql.status = "COMPLETED"
1143+
sql.success_job_ids = [1]
1144+
sql.failed_job_ids = []
1145+
sql.running_job_ids = []
1146+
sql.description = "Test Query"
1147+
sql.submission_time = datetime.now()
1148+
sql.plan_description = "Sample plan description"
1149+
1150+
mock_client.get_sql_list.return_value = [sql]
1151+
mock_get_client.return_value = mock_client
1152+
1153+
# Call function with explicit include_plan_description=True (config should override to False)
1154+
result = list_slowest_sql_queries(
1155+
"spark-app-123", include_plan_description=True
1156+
)
1157+
1158+
# Verify plan description is NOT included because server config is False
1159+
self.assertEqual(len(result), 1)
1160+
self.assertEqual(result[0].plan_description, "")

0 commit comments

Comments
 (0)