fix(dataset): check if dataset exists

superset/mcp_service/chart/schemas.py

@@ -1574,7 +1574,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: Dict[str, Any] = Field(..., description=_CHART_CONFIG_DESCRIPTION)
     generate_preview: bool = True

superset/mcp_service/chart/tool/update_chart_preview.py

@@ -53,6 +53,23 @@ from superset.utils import json as utils_json
 logger = logging.getLogger(__name__)
 
 
+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_old_adhoc_filters(form_data_key: str) -> list[Dict[str, Any]] | None:
     """Retrieve adhoc_filters from the previously cached form_data."""
     from superset.commands.exceptions import CommandException
@@ -91,10 +108,11 @@ def update_chart_preview(
 
     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
@@ -131,6 +149,33 @@ def update_chart_preview(
                 "api_version": "v1",
             }
 
+        # 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

@@ -19,8 +19,13 @@
 Unit tests for update_chart_preview MCP tool
 """
 
-import pytest
+import importlib
+from unittest.mock import Mock, patch
 
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
 from superset.mcp_service.chart.schemas import (
     AxisConfig,
     ColumnRef,
@@ -374,6 +379,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."""
@@ -472,3 +495,54 @@ class TestUpdateChartPreview:
         )
         assert request.config["sort_by"] == ["sales", "profit"]
         assert len(request.config["columns"]) == 3
+
+
+@pytest.fixture
+def mcp_server():
+    return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+    """Mock authentication for all tool tests."""
+    with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+        mock_user = Mock()
+        mock_user.id = 1
+        mock_user.username = "admin"
+        mock_get_user.return_value = mock_user
+        yield mock_get_user
+
+
+class TestUpdateChartPreviewTool:
+    """Tests for update_chart_preview tool execution."""
+
+    @patch.object(
+        importlib.import_module("superset.mcp_service.chart.tool.update_chart_preview"),
+        "_find_dataset",
+        return_value=None,
+    )
+    @pytest.mark.asyncio
+    async def test_update_chart_preview_dataset_not_found(
+        self, mock_find_dataset, mcp_server
+    ):
+        """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"]