diff --git a/superset/mcp_service/chart/preview_utils.py b/superset/mcp_service/chart/preview_utils.py index d585e3cda99..a5ae10b3de3 100644 --- a/superset/mcp_service/chart/preview_utils.py +++ b/superset/mcp_service/chart/preview_utils.py @@ -36,6 +36,8 @@ from superset.mcp_service.chart.schemas import ( logger = logging.getLogger(__name__) +SUPPORTED_FORM_DATA_PREVIEW_FORMATS = frozenset({"ascii", "table", "vega_lite"}) + def _build_query_columns(form_data: Dict[str, Any]) -> list[str]: """Build query columns list from form_data, including both x_axis and groupby.""" diff --git a/superset/mcp_service/chart/tool/generate_chart.py b/superset/mcp_service/chart/tool/generate_chart.py index 8ad907c1abe..aa31afecc4b 100644 --- a/superset/mcp_service/chart/tool/generate_chart.py +++ b/superset/mcp_service/chart/tool/generate_chart.py @@ -43,6 +43,7 @@ from superset.mcp_service.chart.compile import ( CompileResult, validate_and_compile, ) +from superset.mcp_service.chart.preview_utils import SUPPORTED_FORM_DATA_PREVIEW_FORMATS from superset.mcp_service.chart.schemas import ( AccessibilityMetadata, CHART_FORM_DATA_EXCLUDED_FIELD_NAMES, @@ -630,11 +631,7 @@ async def generate_chart( # noqa: C901 # For preview-only mode (save_chart=false) # Note: Screenshot-based URL previews are not # supported. Use explore_url to view interactively. - if format_type in [ - "ascii", - "table", - "vega_lite", - ]: + if format_type in SUPPORTED_FORM_DATA_PREVIEW_FORMATS: # Generate preview from form data from superset.mcp_service.chart.preview_utils import ( generate_preview_from_form_data, diff --git a/superset/mcp_service/chart/tool/update_chart_preview.py b/superset/mcp_service/chart/tool/update_chart_preview.py index 82ace0df0ad..d130411f677 100644 --- a/superset/mcp_service/chart/tool/update_chart_preview.py +++ b/superset/mcp_service/chart/tool/update_chart_preview.py @@ -40,8 +40,13 @@ from superset.mcp_service.chart.chart_utils import ( map_config_to_form_data, ) from superset.mcp_service.chart.compile import validate_and_compile +from superset.mcp_service.chart.preview_utils import ( + generate_preview_from_form_data, + SUPPORTED_FORM_DATA_PREVIEW_FORMATS, +) from superset.mcp_service.chart.schemas import ( AccessibilityMetadata, + ChartError, PerformanceMetadata, UpdateChartPreviewRequest, ) @@ -229,9 +234,39 @@ def update_chart_preview( # noqa: C901 high_contrast_available=False, ) - # Note: Screenshot-based previews are not supported. - # Use the explore_url to view the chart interactively. previews: Dict[str, Any] = {} + if request.generate_preview: + try: + with event_logger.log_context( + action="mcp.update_chart_preview.preview" + ): + for format_type in request.preview_formats: + # URL previews are represented by explore_url/chart.url. + # Screenshot-based previews are not supported. + if format_type not in SUPPORTED_FORM_DATA_PREVIEW_FORMATS: + continue + + preview_result = generate_preview_from_form_data( + form_data=new_form_data, + dataset_id=dataset.id, + preview_format=format_type, + ) + + if isinstance(preview_result, ChartError): + logger.warning( + "Preview '%s' failed: %s", + format_type, + preview_result.error, + ) + else: + previews[format_type] = ( + preview_result.model_dump(mode="json") + if hasattr(preview_result, "model_dump") + else preview_result + ) + + except (CommandException, ValueError, KeyError) as e: + logger.warning("Preview generation failed: %s", e) # Return enhanced data result = { diff --git a/tests/unit_tests/mcp_service/chart/tool/test_update_chart_preview.py b/tests/unit_tests/mcp_service/chart/tool/test_update_chart_preview.py index 2aa63efdf1c..caeb1e7a747 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_update_chart_preview.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_update_chart_preview.py @@ -32,6 +32,7 @@ from superset.mcp_service.chart.schemas import ( FilterConfig, LegendConfig, TableChartConfig, + TablePreview, UpdateChartPreviewRequest, XYChartConfig, ) @@ -698,6 +699,79 @@ class TestUpdateChartPreview: assert result["warnings"] == [] mock_get_previous_form_data.assert_called_once_with("valid_key_12345") + @patch.object(update_chart_preview_module, "validate_and_compile") + @patch.object(update_chart_preview_module, "has_dataset_access", return_value=True) + @patch("superset.daos.dataset.DatasetDAO.find_by_id") + @patch.object(update_chart_preview_module, "generate_preview_from_form_data") + @patch.object(update_chart_preview_module, "analyze_chart_semantics") + @patch.object(update_chart_preview_module, "analyze_chart_capabilities") + @patch.object(update_chart_preview_module, "generate_explore_link") + @patch.object(update_chart_preview_module, "_get_previous_form_data") + @patch("superset.mcp_service.auth.get_user_from_request") + @pytest.mark.asyncio + async def test_returns_requested_table_preview( + self, + mock_get_user_from_request, + mock_get_previous_form_data, + mock_generate_explore_link, + mock_analyze_chart_capabilities, + mock_analyze_chart_semantics, + mock_generate_preview_from_form_data, + mock_find_by_id, + unused_access_mock, + mock_validate_and_compile, + ) -> None: + """Preview updates honor supported preview_formats.""" + mock_user = Mock() + mock_user.id = 1 + mock_get_user_from_request.return_value = mock_user + mock_find_by_id.return_value = _mock_dataset(id=3) + mock_validate_and_compile.return_value = Mock(success=True) + mock_get_previous_form_data.return_value = {} + mock_generate_explore_link.return_value = ( + "http://localhost:8088/explore/?form_data_key=new_preview_key" + ) + mock_analyze_chart_capabilities.return_value = None + mock_analyze_chart_semantics.return_value = None + table_preview = TablePreview( + table_data="Table Preview", + row_count=1, + supports_sorting=True, + ) + expected_table_preview = { + "type": "table", + "table_data": "Table Preview", + "row_count": 1, + "supports_sorting": True, + } + mock_generate_preview_from_form_data.return_value = table_preview + + request = UpdateChartPreviewRequest( + form_data_key="valid_key_12345", + dataset_id=3, + config=TableChartConfig( + chart_type="table", + columns=[ + ColumnRef(name="country", label="Country"), + ColumnRef(name="sales", label="Sales", aggregate="SUM"), + ], + ), + generate_preview=True, + preview_formats=["url", "table"], + ) + + result = update_chart_preview_module.update_chart_preview( + request=request, ctx=Mock() + ) + + assert result["success"] is True + assert result["previews"] == {"table": expected_table_preview} + mock_generate_preview_from_form_data.assert_called_once() + preview_kwargs = mock_generate_preview_from_form_data.call_args.kwargs + assert preview_kwargs["dataset_id"] == 3 + assert preview_kwargs["preview_format"] == "table" + assert preview_kwargs["form_data"]["viz_type"] == "table" + class TestUpdateChartPreviewValidation: """Tier-1 validation gate and dataset access checks."""