diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index 29dddb33fd9..f8e9d818607 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -1155,6 +1155,16 @@ class GenerateChartResponse(BaseModel): default_factory=dict, description="Related API endpoints for data/updates" ) + # Form data for rendering charts in external clients (chatbot rendering) + form_data: Dict[str, Any] = Field( + default_factory=dict, + description="Complete form_data configuration for rendering the chart", + ) + form_data_key: str | None = Field( + None, + description="Cache key for the form_data, used in explore URLs", + ) + # Performance and accessibility performance: PerformanceMetadata | None = Field( None, description="Performance metrics" diff --git a/superset/mcp_service/chart/tool/generate_chart.py b/superset/mcp_service/chart/tool/generate_chart.py index e218f4cd202..dfbc466a57f 100644 --- a/superset/mcp_service/chart/tool/generate_chart.py +++ b/superset/mcp_service/chart/tool/generate_chart.py @@ -20,6 +20,7 @@ MCP tool: generate_chart (simplified schema) import logging import time +from urllib.parse import parse_qs, urlparse from fastmcp import Context from superset_core.mcp import tool @@ -284,6 +285,43 @@ async def generate_chart( # noqa: C901 raise # Update explore URL to use saved chart explore_url = f"{get_superset_base_url()}/explore/?slice_id={chart.id}" + + # Generate form_data_key for saved charts (needed for chatbot rendering) + try: + from superset.commands.explore.form_data.parameters import ( + CommandParameters, + ) + from superset.mcp_service.commands.create_form_data import ( + MCPCreateFormDataCommand, + ) + from superset.utils.core import DatasourceType + + # Add datasource to form_data for the cache + form_data_with_datasource = { + **form_data, + "datasource": f"{dataset.id}__table", + } + + cmd_params = CommandParameters( + datasource_type=DatasourceType.TABLE, + datasource_id=dataset.id, + chart_id=chart.id, + tab_id=None, + form_data=json.dumps(form_data_with_datasource), + ) + form_data_key = MCPCreateFormDataCommand(cmd_params).run() + await ctx.debug( + "Generated form_data_key for saved chart: form_data_key=%s" + % (form_data_key,) + ) + except Exception as fdk_error: + logger.warning( + "Failed to generate form_data_key for saved chart: %s", fdk_error + ) + await ctx.warning( + "Failed to generate form_data_key: error=%s" % (str(fdk_error),) + ) + # form_data_key remains None but chart is still valid else: await ctx.report_progress(2, 5, "Generating temporary chart preview") # Generate explore link with cached form_data for preview-only mode @@ -292,9 +330,13 @@ async def generate_chart( # noqa: C901 explore_url = generate_explore_link(request.dataset_id, form_data) await ctx.debug("Generated explore link: explore_url=%s" % (explore_url,)) - # Extract form_data_key from the explore URL - if explore_url and "form_data_key=" in explore_url: - form_data_key = explore_url.split("form_data_key=")[1].split("&")[0] + # Extract form_data_key from the explore URL using proper URL parsing + if explore_url: + parsed = urlparse(explore_url) + query_params = parse_qs(parsed.query) + form_data_key_list = query_params.get("form_data_key", []) + if form_data_key_list: + form_data_key = form_data_key_list[0] # Generate semantic analysis capabilities = analyze_chart_capabilities(chart, request.config) @@ -405,25 +447,37 @@ async def generate_chart( # noqa: C901 # Return enhanced data while maintaining backward compatibility await ctx.report_progress(4, 5, "Building response") + + # Build chart info using serialize_chart_object for saved charts + chart_info = None + if request.save_chart and chart: + from superset.mcp_service.chart.schemas import serialize_chart_object + + chart_info = serialize_chart_object(chart) + if chart_info: + # Override the URL with explore_url + chart_info.url = explore_url + + # Safely serialize chart_info - handle both Pydantic models and dicts + chart_data = None + if chart_info: + if hasattr(chart_info, "model_dump"): + chart_data = chart_info.model_dump() + elif isinstance(chart_info, dict): + chart_data = chart_info + else: + chart_data = chart_info # Pass through as-is + result = { - "chart": { - "id": chart.id if chart else None, - "slice_name": chart.slice_name - if chart - else generate_chart_name(request.config), - "viz_type": chart.viz_type if chart else form_data.get("viz_type"), - "url": explore_url, - "uuid": str(chart.uuid) if chart and chart.uuid else None, - "saved": request.save_chart, - } - if request.save_chart and chart - else None, + "chart": chart_data, "error": None, # Enhanced fields for better LLM integration "previews": previews, "capabilities": capabilities.model_dump() if capabilities else None, "semantics": semantics.model_dump() if semantics else None, "explore_url": explore_url, + # Form data fields - REQUIRED for chatbot/external client rendering + "form_data": form_data, "form_data_key": form_data_key, "api_endpoints": { "data": f"{get_superset_base_url()}/api/v1/chart/{chart.id}/data/" diff --git a/superset/mcp_service/explore/tool/generate_explore_link.py b/superset/mcp_service/explore/tool/generate_explore_link.py index 964140a4489..611659792c1 100644 --- a/superset/mcp_service/explore/tool/generate_explore_link.py +++ b/superset/mcp_service/explore/tool/generate_explore_link.py @@ -23,6 +23,7 @@ chart configuration. """ from typing import Any, Dict +from urllib.parse import parse_qs, urlparse from fastmcp import Context from superset_core.mcp import tool @@ -106,14 +107,26 @@ async def generate_explore_link( # Generate explore link using shared utilities explore_url = generate_url(dataset_id=request.dataset_id, form_data=form_data) + # Extract form_data_key from the explore URL using proper URL parsing + form_data_key = None + if explore_url: + parsed = urlparse(explore_url) + query_params = parse_qs(parsed.query) + form_data_key_list = query_params.get("form_data_key", []) + if form_data_key_list: + form_data_key = form_data_key_list[0] + await ctx.report_progress(3, 3, "URL generation complete") await ctx.info( - "Explore link generated successfully: url_length=%s, dataset_id=%s" - % (len(explore_url), request.dataset_id) + "Explore link generated successfully: url_length=%s, dataset_id=%s, " + "form_data_key=%s" + % (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, } @@ -124,5 +137,7 @@ async def generate_explore_link( ) return { "url": "", + "form_data": {}, + "form_data_key": None, "error": f"Failed to generate explore link: {str(e)}", } 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 9e209050ee1..0769844538e 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 @@ -589,3 +589,97 @@ class TestGenerateExploreLink: 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 + + @pytest.mark.asyncio + async def test_generate_explore_link_tool_exception_handling(self, mcp_server): + """Test that tool-level exceptions are properly handled and return error.""" + import sys + + # Get the actual module object from sys.modules (not via __init__.py which + # returns the function) + explore_module = sys.modules[ + "superset.mcp_service.explore.tool.generate_explore_link" + ] + + original_func = explore_module.map_config_to_form_data + + def raise_error(*args, **kwargs): + raise ValueError("Invalid config structure") + + explore_module.map_config_to_form_data = raise_error + try: + config = TableChartConfig( + chart_type="table", columns=[ColumnRef(name="test_col")] + ) + request = GenerateExploreLinkRequest(dataset_id="1", config=config) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "generate_explore_link", {"request": request.model_dump()} + ) + + # 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"] + finally: + # Restore original function + explore_module.map_config_to_form_data = original_func + + @patch("superset.daos.dataset.DatasetDAO.find_by_id") + @patch( + "superset.mcp_service.commands.create_form_data.MCPCreateFormDataCommand.run" + ) + @pytest.mark.asyncio + async def test_generate_explore_link_returns_form_data_key( + self, mock_create_form_data, mock_find_dataset, mcp_server + ): + """Test that form_data_key is properly extracted from URL.""" + mock_create_form_data.return_value = "extracted_form_key_xyz" + mock_find_dataset.return_value = _mock_dataset(id=1) + + config = TableChartConfig( + chart_type="table", columns=[ColumnRef(name="region")] + ) + request = GenerateExploreLinkRequest(dataset_id="1", config=config) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "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"] + + @patch("superset.daos.dataset.DatasetDAO.find_by_id") + @patch( + "superset.mcp_service.commands.create_form_data.MCPCreateFormDataCommand.run" + ) + @pytest.mark.asyncio + async def test_generate_explore_link_returns_form_data( + self, mock_create_form_data, mock_find_dataset, mcp_server + ): + """Test that form_data dict is returned for external rendering.""" + mock_create_form_data.return_value = "form_data_test_key" + mock_find_dataset.return_value = _mock_dataset(id=1) + + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="sales", aggregate="SUM")], + kind="line", + ) + request = GenerateExploreLinkRequest(dataset_id="1", config=config) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "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"