diff --git a/superset/mcp_service/explore/schemas.py b/superset/mcp_service/explore/schemas.py new file mode 100644 index 000000000000..3e673388b550 --- /dev/null +++ b/superset/mcp_service/explore/schemas.py @@ -0,0 +1,93 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Pydantic schemas for explore-related MCP tool inputs and outputs. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class GenerateExploreLinkResponse(BaseModel): + """ + Output schema for the generate_explore_link tool. + + On success, url is a fully-qualified Superset Explore URL the user can + open immediately. form_data_key can be used to reconstruct or share the + same configuration. On failure, url is empty and error contains a + human-readable description of what went wrong. + """ + + url: str = Field( + ..., + description=( + "Explore URL — open in a browser to view the interactive chart. " + "Empty string on failure." + ), + ) + form_data: dict[str, Any] = Field( + default_factory=dict, + description="Raw Superset form_data dict that was encoded into the URL.", + ) + form_data_key: str | None = Field( + None, + description=( + "Short cache key that represents this form_data configuration. " + "Can be passed to the Explore UI as ?form_data_key=." + ), + ) + error: str | None = Field( + None, + description="Human-readable error message when generation fails, else null.", + ) + success: bool = Field( + True, + description="True when a valid URL was produced, False on any error.", + ) + + +class ExploreLinkError(BaseModel): + """ + Structured error returned when explore link generation fails. + + error_type distinguishes actionable failure modes so callers can + respond appropriately without parsing the free-text error message. + """ + + error: str = Field(..., description="Human-readable description of the failure.") + error_type: Literal[ + "dataset_not_found", + "generation_failed", + "permission_denied", + ] = Field( + ..., + description=( + "Machine-readable failure category. " + "'dataset_not_found' — dataset_id does not exist; use list_datasets. " + "'generation_failed' — URL could not be constructed from the config. " + "'permission_denied' — current user lacks access to the dataset." + ), + ) + dataset_id: int | str | None = Field( + None, + description="The dataset_id from the original request, for correlation.", + ) + success: bool = Field(False) diff --git a/superset/mcp_service/explore/tool/generate_explore_link.py b/superset/mcp_service/explore/tool/generate_explore_link.py index 1ff0ea0d6c19..30e9f3c07ef5 100644 --- a/superset/mcp_service/explore/tool/generate_explore_link.py +++ b/superset/mcp_service/explore/tool/generate_explore_link.py @@ -22,8 +22,6 @@ chart configuration. """ -from typing import Any, Dict - from fastmcp import Context from superset_core.mcp.decorators import tool, ToolAnnotations @@ -33,9 +31,8 @@ generate_explore_link as generate_url, map_config_to_form_data, ) -from superset.mcp_service.chart.schemas import ( - GenerateExploreLinkRequest, -) +from superset.mcp_service.chart.schemas import GenerateExploreLinkRequest +from superset.mcp_service.explore.schemas import GenerateExploreLinkResponse @tool( @@ -49,7 +46,7 @@ ) async def generate_explore_link( request: GenerateExploreLinkRequest, ctx: Context -) -> Dict[str, Any]: +) -> GenerateExploreLinkResponse: """Generate explore URL for interactive visualization. PREFERRED TOOL for most visualization requests. @@ -121,15 +118,16 @@ async def generate_explore_link( await ctx.warning( "Dataset not found: dataset_id=%s" % (request.dataset_id,) ) - return { - "url": "", - "form_data": {}, - "form_data_key": None, - "error": ( + return GenerateExploreLinkResponse( + url="", + form_data={}, + form_data_key=None, + error=( f"Dataset not found: {request.dataset_id}. " "Use list_datasets to find valid dataset IDs." ), - } + success=False, + ) await ctx.report_progress(2, 4, "Converting configuration to form data") with event_logger.log_context(action="mcp.generate_explore_link.form_data"): @@ -184,12 +182,13 @@ async def generate_explore_link( % (len(explore_url or ""), request.dataset_id, form_data_key) ) - return { - "url": explore_url, - "form_data": form_data, - "form_data_key": form_data_key, - "error": None, - } + return GenerateExploreLinkResponse( + url=explore_url, + form_data=form_data, + form_data_key=form_data_key, + error=None, + success=True, + ) except Exception as e: await ctx.error( @@ -201,9 +200,10 @@ async def generate_explore_link( str(e), ) ) - return { - "url": "", - "form_data": {}, - "form_data_key": None, - "error": f"Failed to generate explore link: {str(e)}", - } + return GenerateExploreLinkResponse( + url="", + form_data={}, + form_data_key=None, + error=f"Failed to generate explore link: {str(e)}", + success=False, + ) diff --git a/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py b/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py index fb8aee539b04..9f94988fe7b2 100644 --- a/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py +++ b/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py @@ -121,9 +121,9 @@ async def test_generate_table_explore_link_minimal( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=test_form_data_key_123" ) mock_create_form_data.assert_called_once() @@ -160,9 +160,9 @@ async def test_generate_table_explore_link_with_features( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=comprehensive_key_456" ) mock_create_form_data.assert_called_once() @@ -199,9 +199,9 @@ async def test_generate_line_chart_explore_link( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=line_chart_key_789" ) mock_create_form_data.assert_called_once() @@ -233,9 +233,9 @@ async def test_generate_bar_chart_explore_link( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=bar_chart_key_abc" ) mock_create_form_data.assert_called_once() @@ -270,9 +270,9 @@ async def test_generate_area_chart_explore_link( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=area_chart_key_def" ) mock_create_form_data.assert_called_once() @@ -305,9 +305,9 @@ async def test_generate_scatter_chart_explore_link( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=scatter_chart_key_ghi" ) mock_create_form_data.assert_called_once() @@ -337,9 +337,9 @@ async def test_generate_explore_link_cache_failure_fallback( ) # Should fallback to basic URL format - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?datasource_type=table&datasource_id=1" ) @@ -373,9 +373,9 @@ async def test_generate_explore_link_database_lock_fallback( ) # Should fallback to basic dataset URL - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?datasource_type=table&datasource_id=5" ) @@ -409,9 +409,9 @@ async def test_generate_explore_link_with_many_columns( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=many_columns_key" ) mock_create_form_data.assert_called_once() @@ -452,9 +452,9 @@ async def test_generate_explore_link_with_many_filters( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=many_filters_key" ) mock_create_form_data.assert_called_once() @@ -509,10 +509,10 @@ async def test_explore_link_url_format_consistency( # All URLs should follow the same format assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=consistency_test_key" ) - assert result.data["error"] is None + assert result.structured_content["error"] is None @patch("superset.daos.dataset.DatasetDAO.find_by_id") @patch( @@ -539,9 +539,9 @@ async def test_generate_explore_link_dataset_id_types( result = await client.call_tool( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=dataset_test_key" ) @@ -583,9 +583,9 @@ async def test_generate_explore_link_complex_configuration( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None assert ( - result.data["url"] + result.structured_content["url"] == "http://localhost:9001/explore/?form_data_key=complex_config_key" ) mock_create_form_data.assert_called_once() @@ -619,8 +619,8 @@ async def test_fallback_url_different_datasets( # Should fallback to basic URL with correct dataset_id expected_url = f"http://localhost:9001/explore/?datasource_type=table&datasource_id={dataset_id}" - assert result.data["error"] is None - assert result.data["url"] == expected_url + assert result.structured_content["error"] is None + assert result.structured_content["url"] == expected_url @patch("superset.daos.dataset.DatasetDAO.find_by_id") @pytest.mark.asyncio @@ -655,10 +655,10 @@ def raise_error(*args, **kwargs): ) # Should return error response with empty URL - assert result.data["url"] == "" - assert result.data["form_data"] == {} - assert result.data["form_data_key"] is None - assert "Invalid config structure" in result.data["error"] + assert result.structured_content["url"] == "" + assert result.structured_content["form_data"] == {} + assert result.structured_content["form_data_key"] is None + assert "Invalid config structure" in result.structured_content["error"] finally: # Restore original function explore_module.map_config_to_form_data = original_func @@ -685,9 +685,14 @@ async def test_generate_explore_link_returns_form_data_key( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None - assert result.data["form_data_key"] == "extracted_form_key_xyz" - assert "form_data_key=extracted_form_key_xyz" in result.data["url"] + assert result.structured_content["error"] is None + assert ( + result.structured_content["form_data_key"] == "extracted_form_key_xyz" + ) + assert ( + "form_data_key=extracted_form_key_xyz" + in result.structured_content["url"] + ) @patch("superset.daos.dataset.DatasetDAO.find_by_id") @patch( @@ -714,13 +719,18 @@ async def test_generate_explore_link_returns_form_data( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None - assert "form_data" in result.data - assert isinstance(result.data["form_data"], dict) - assert result.data["form_data"].get("viz_type") == "echarts_timeseries_line" - assert result.data["form_data"].get("x_axis") == "date" + assert result.structured_content["error"] is None + assert "form_data" in result.structured_content + assert isinstance(result.structured_content["form_data"], dict) + assert ( + result.structured_content["form_data"].get("viz_type") + == "echarts_timeseries_line" + ) + assert result.structured_content["form_data"].get("x_axis") == "date" # Verify datasource field format: "{dataset_id}__table" - assert result.data["form_data"].get("datasource") == "1__table" + assert ( + result.structured_content["form_data"].get("datasource") == "1__table" + ) @patch("superset.daos.dataset.DatasetDAO.find_by_id") @pytest.mark.asyncio @@ -740,11 +750,11 @@ async def test_generate_explore_link_nonexistent_dataset( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["url"] == "" - assert result.data["form_data"] == {} - assert result.data["form_data_key"] is None - assert "Dataset not found: 99999" in result.data["error"] - assert "list_datasets" in result.data["error"] + assert result.structured_content["url"] == "" + assert result.structured_content["form_data"] == {} + assert result.structured_content["form_data_key"] is None + assert "Dataset not found: 99999" in result.structured_content["error"] + assert "list_datasets" in result.structured_content["error"] @patch("superset.daos.dataset.DatasetDAO.find_by_id") @pytest.mark.asyncio @@ -766,10 +776,10 @@ async def test_generate_explore_link_nonexistent_uuid_dataset( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["url"] == "" - assert result.data["form_data"] == {} - assert result.data["form_data_key"] is None - assert "Dataset not found" in result.data["error"] + assert result.structured_content["url"] == "" + assert result.structured_content["form_data"] == {} + assert result.structured_content["form_data_key"] is None + assert "Dataset not found" in result.structured_content["error"] class TestGenerateExploreLinkColumnNormalization: @@ -823,9 +833,9 @@ async def test_xy_chart_x_axis_normalized_in_form_data( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None # x-axis should be normalized from 'orderdate' to 'OrderDate' - assert result.data["form_data"]["x_axis"] == "OrderDate" + assert result.structured_content["form_data"]["x_axis"] == "OrderDate" @patch( "superset.mcp_service.chart.validation.dataset_validator.DatasetValidator._get_dataset_context" @@ -873,8 +883,8 @@ async def test_filter_column_normalized_in_form_data( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None - form_data = result.data["form_data"] + assert result.structured_content["error"] is None + form_data = result.structured_content["form_data"] # x-axis normalized assert form_data["x_axis"] == "OrderDate" # filter subject normalized to match x-axis @@ -919,6 +929,6 @@ async def test_normalization_fallback_when_dataset_not_found( "generate_explore_link", {"request": request.model_dump()} ) - assert result.data["error"] is None + assert result.structured_content["error"] is None # original names should pass through unchanged - assert result.data["form_data"]["x_axis"] == "orderdate" + assert result.structured_content["form_data"]["x_axis"] == "orderdate"