Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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: bool = False


class McpConfig(BaseSettings):
Expand Down
8 changes: 6 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,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

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 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")