feat(mcp): make form_data_key optional in update_chart_preview (#39680)

superset/mcp_service/chart/schemas.py

@@ -1771,7 +1771,13 @@ class UpdateChartRequest(QueryCacheControl):
 
 
 class UpdateChartPreviewRequest(FormDataCacheControl):
-    form_data_key: str = Field(..., description="Existing form_data_key to update")
+    form_data_key: str | None = Field(
+        None,
+        description=(
+            "Existing form_data_key to update"
+            " (omit for fresh preview from config + dataset_id)"
+        ),
+    )
     dataset_id: int | str = Field(..., description="Dataset ID or UUID")
     config: ChartConfig = Field(..., description="Chart configuration")
     generate_preview: bool = True

superset/mcp_service/chart/tool/update_chart_preview.py

@@ -66,6 +66,23 @@ INVALID_FORM_DATA_KEY_WARNING = (
 )
 
 
+def _find_dataset(dataset_id: int | str) -> Any | None:
+    """Look up a dataset by numeric ID or UUID and check access."""
+    from superset.daos.dataset import DatasetDAO
+    from superset.mcp_service.auth import has_dataset_access
+
+    if isinstance(dataset_id, int) or (
+        isinstance(dataset_id, str) and dataset_id.isdigit()
+    ):
+        dataset = DatasetDAO.find_by_id(int(dataset_id))
+    else:
+        dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
+
+    if dataset and not has_dataset_access(dataset):
+        return None
+    return dataset
+
+
 def _get_previous_form_data(form_data_key: str) -> dict[str, Any] | None:
     """Retrieve the previously cached form_data."""
     from superset.commands.exceptions import CommandException
@@ -102,10 +119,11 @@ def update_chart_preview(  # noqa: C901
 
     IMPORTANT:
     - Modifies cached form_data from generate_chart (save_chart=False)
-    - Original form_data_key is invalidated, new one returned
+    - Original form_data_key (when provided) is invalidated, new one returned
     - LLM clients MUST display explore_url to users
 
     Use when:
+    - Creating a fresh preview from config + dataset_id (omit form_data_key)
     - Modifying preview before deciding to save
     - Iterating on chart design without creating permanent charts
     - Testing different configurations
@@ -118,6 +136,33 @@ def update_chart_preview(  # noqa: C901
         # config is already a typed ChartConfig (validated by Pydantic)
         config = request.config
 
+        # Validate dataset exists and user has access
+        with event_logger.log_context(action="mcp.update_chart_preview.dataset_lookup"):
+            dataset = _find_dataset(request.dataset_id)
+
+            if not dataset:
+                return {
+                    "chart": None,
+                    "error": {
+                        "error_type": "dataset_not_found",
+                        "message": (f"Dataset not found: {request.dataset_id}"),
+                        "details": (
+                            f"No dataset found with identifier "
+                            f"'{request.dataset_id}'. This could "
+                            f"be an invalid ID/UUID or a "
+                            f"permissions issue."
+                        ),
+                        "suggestions": [
+                            "Verify the dataset ID or UUID",
+                            "Check dataset access permissions",
+                            "Use list_datasets to find available datasets",
+                        ],
+                    },
+                    "success": False,
+                    "schema_version": "2.0",
+                    "api_version": "v1",
+                }
+
         with event_logger.log_context(action="mcp.update_chart_preview.form_data"):
             # Map the new config to form_data format
             # Pass dataset_id to enable column type checking

tests/unit_tests/mcp_service/chart/tool/test_update_chart_preview.py

@@ -422,6 +422,24 @@ class TestUpdateChartPreview:
             )
             assert request.form_data_key == key
 
+    @pytest.mark.asyncio
+    async def test_update_chart_preview_form_data_key_optional(self):
+        """Test that form_data_key can be omitted for fresh previews."""
+        config = TableChartConfig(
+            chart_type="table",
+            columns=[ColumnRef(name="col1")],
+        )
+
+        # Omit form_data_key entirely
+        request = UpdateChartPreviewRequest(dataset_id=1, config=config)
+        assert request.form_data_key is None
+
+        # Explicitly pass None
+        request2 = UpdateChartPreviewRequest(
+            form_data_key=None, dataset_id=1, config=config
+        )
+        assert request2.form_data_key is None
+
     @pytest.mark.asyncio
     async def test_update_chart_preview_cache_control(self):
         """Test cache control parameters in update preview request."""
@@ -575,11 +593,13 @@ class TestUpdateChartPreview:
     @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.object(update_chart_preview_module, "_find_dataset")
     @patch("superset.mcp_service.auth.get_user_from_request")
     @pytest.mark.asyncio
     async def test_warns_when_previous_form_data_key_is_missing(
         self,
         mock_get_user_from_request,
+        mock_find_dataset,
         mock_get_previous_form_data,
         mock_generate_explore_link,
         mock_analyze_chart_capabilities,
@@ -592,6 +612,7 @@ class TestUpdateChartPreview:
         mock_user = Mock()
         mock_user.id = 1
         mock_get_user_from_request.return_value = mock_user
+        mock_find_dataset.return_value = _mock_dataset(id=3)
         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 = None
@@ -636,11 +657,13 @@ class TestUpdateChartPreview:
     @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.object(update_chart_preview_module, "_find_dataset")
     @patch("superset.mcp_service.auth.get_user_from_request")
     @pytest.mark.asyncio
     async def test_preserves_previous_adhoc_filters_without_warning(
         self,
         mock_get_user_from_request,
+        mock_find_dataset,
         mock_get_previous_form_data,
         mock_generate_explore_link,
         mock_analyze_chart_capabilities,
@@ -653,6 +676,7 @@ class TestUpdateChartPreview:
         mock_user = Mock()
         mock_user.id = 1
         mock_get_user_from_request.return_value = mock_user
+        mock_find_dataset.return_value = _mock_dataset(id=3)
         mock_find_by_id.return_value = _mock_dataset(id=3)
         mock_validate_and_compile.return_value = Mock(success=True)
         cached_adhoc_filters = [
@@ -707,11 +731,13 @@ class TestUpdateChartPreview:
     @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.object(update_chart_preview_module, "_find_dataset")
     @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_find_dataset,
         mock_get_previous_form_data,
         mock_generate_explore_link,
         mock_analyze_chart_capabilities,
@@ -725,6 +751,7 @@ class TestUpdateChartPreview:
         mock_user = Mock()
         mock_user.id = 1
         mock_get_user_from_request.return_value = mock_user
+        mock_find_dataset.return_value = _mock_dataset(id=3)
         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 = {}
@@ -773,9 +800,45 @@ class TestUpdateChartPreview:
         assert preview_kwargs["form_data"]["viz_type"] == "table"
 
 
+class TestUpdateChartPreviewTool:
+    """Tests for update_chart_preview tool execution."""
+
+    @patch.object(
+        update_chart_preview_module,
+        "_find_dataset",
+        return_value=None,
+    )
+    @pytest.mark.asyncio
+    async def test_update_chart_preview_dataset_not_found(
+        self, mock_find_dataset, mcp_server, mock_auth
+    ):
+        """Test that a non-existent dataset returns a clear error."""
+
+        request = {
+            "dataset_id": 99999,
+            "config": {
+                "chart_type": "table",
+                "columns": [{"name": "col1"}],
+            },
+        }
+
+        async with Client(mcp_server) as client:
+            result = await client.call_tool(
+                "update_chart_preview", {"request": request}
+            )
+
+            data = result.structured_content
+            assert data["success"] is False
+            assert data["chart"] is None
+            error = data["error"]
+            assert error["error_type"] == "dataset_not_found"
+            assert "99999" in error["message"]
+
+
 class TestUpdateChartPreviewValidation:
     """Tier-1 validation gate and dataset access checks."""
 
+    @patch.object(update_chart_preview_module, "_find_dataset")
     @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, "validate_and_compile")
@@ -787,8 +850,9 @@ class TestUpdateChartPreviewValidation:
         self,
         mock_create_form_data,
         mock_validate,
-        mock_find_dataset,
+        mock_find_by_id,
         unused_access_mock,
+        mock_find_dataset,
         mcp_server,
         mock_auth,
     ):
@@ -797,6 +861,7 @@ class TestUpdateChartPreviewValidation:
         from superset.mcp_service.common.error_schemas import ChartGenerationError
 
         mock_find_dataset.return_value = _mock_dataset(id=3)
+        mock_find_by_id.return_value = _mock_dataset(id=3)
         mock_validate.return_value = CompileResult(
             success=False,
             error="Column 'num_boys' does not exist in dataset",
@@ -834,6 +899,7 @@ class TestUpdateChartPreviewValidation:
             assert "sum_boys" in error["suggestions"]
             mock_create_form_data.assert_not_called()
 
+    @patch.object(update_chart_preview_module, "_find_dataset")
     @patch.object(update_chart_preview_module, "has_dataset_access", return_value=False)
     @patch("superset.daos.dataset.DatasetDAO.find_by_id")
     @patch(
@@ -843,13 +909,15 @@ class TestUpdateChartPreviewValidation:
     async def test_dataset_access_denied_short_circuits(
         self,
         mock_create_form_data,
-        mock_find_dataset,
+        mock_find_by_id,
         unused_access_mock,
+        mock_find_dataset,
         mcp_server,
         mock_auth,
     ):
         """has_dataset_access=False → DatasetNotAccessible, no cache write."""
         mock_find_dataset.return_value = _mock_dataset(id=3)
+        mock_find_by_id.return_value = _mock_dataset(id=3)
 
         config = TableChartConfig(
             chart_type="table", columns=[ColumnRef(name="region")]