diff --git a/README.md b/README.md index 4b11591..7d976c3 100644 --- a/README.md +++ b/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 diff --git a/src/spark_history_mcp/config/config.py b/src/spark_history_mcp/config/config.py index 639f9e6..230472e 100644 --- a/src/spark_history_mcp/config/config.py +++ b/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: Optional[bool] = None class McpConfig(BaseSettings): diff --git a/src/spark_history_mcp/tools/tools.py b/src/spark_history_mcp/tools/tools.py index 94d3dc7..70aae67 100644 --- a/src/spark_history_mcp/tools/tools.py +++ b/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,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 diff --git a/tests/unit/test_tools.py b/tests/unit/test_tools.py index 580a3f9..c9f25d5 100644 --- a/tests/unit/test_tools.py +++ b/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,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, "")