Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/spark_history_mcp/config/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: Optional[bool] = None


class McpConfig(BaseSettings):
Expand Down
10 changes: 8 additions & 2 deletions src/spark_history_mcp/tools/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
"""
Expand All @@ -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:
Expand All @@ -944,6 +944,12 @@ def list_slowest_sql_queries(
ctx = mcp.get_context()
client = get_client_or_default(ctx, server, app_id)

# Config takes priority: if config is set (True/False), use it; otherwise default to True
if client.config.include_plan_description is not None:
include_plan_description = client.config.include_plan_description
else:
include_plan_description = True

all_executions: List[ExecutionData] = []
offset = 0

Expand Down
71 changes: 71 additions & 0 deletions tests/unit/test_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1087,3 +1088,73 @@ 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 server config overrides parameter when config is set"""
# 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 (config should override to False)
result = list_slowest_sql_queries(
"spark-app-123", include_plan_description=True
)

# Verify plan description is NOT included because server config is False
self.assertEqual(len(result), 1)
self.assertEqual(result[0].plan_description, "")