diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index a711ce535bb..5a4fee99bff 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/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 diff --git a/superset/mcp_service/chart/tool/update_chart_preview.py b/superset/mcp_service/chart/tool/update_chart_preview.py index d130411f677..4fc9a26c955 100644 --- a/superset/mcp_service/chart/tool/update_chart_preview.py +++ b/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 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 caeb1e7a747..fe5c88b9d91 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 @@ -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")]