fix(mcp): add dataset validation for chart tools (#37185)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

superset/mcp_service/CLAUDE.md

@@ -423,6 +423,43 @@ def my_tool(id: int) -> MyResponse:
         raise ValueError(f"Resource {id} not found")
 ```
 
+### 8. Dataset Validation for Chart Tools
+
+**IMPORTANT**: All chart-related tools must validate that the chart's dataset is accessible before performing operations. Use the shared `validate_chart_dataset` utility from `chart_utils.py`.
+
+**Why this matters**: Charts can reference datasets that have been deleted or become inaccessible. Without validation, users may see confusing errors when trying to preview or get data from charts with missing datasets.
+
+**Usage pattern**:
+```python
+from superset.mcp_service.chart.chart_utils import validate_chart_dataset
+
+# After retrieving a chart, validate its dataset
+validation_result = validate_chart_dataset(chart, check_access=True)
+if not validation_result.is_valid:
+    await ctx.warning("Dataset not accessible: %s" % (validation_result.error,))
+    return ChartError(
+        error=validation_result.error or "Chart's dataset is not accessible",
+        error_type="DatasetNotAccessible",
+    )
+
+# Log any warnings (e.g., virtual dataset warnings)
+for warning in validation_result.warnings:
+    await ctx.warning("Dataset warning: %s" % (warning,))
+```
+
+**Tools that use this pattern**:
+- `get_chart_info` - Validates after retrieving chart metadata
+- `get_chart_preview` - Validates before generating preview
+- `get_chart_data` - Validates before querying data
+- `generate_chart` - Validates after chart creation (post-creation validation)
+
+**The `DatasetValidationResult` contains**:
+- `is_valid`: Whether the dataset exists and is accessible
+- `dataset_id`: The dataset ID being validated
+- `dataset_name`: The dataset name (if found)
+- `warnings`: List of warnings (e.g., "virtual dataset may be deleted")
+- `error`: Error message if validation failed
+
 ## Testing Conventions
 
 ### Unit Tests

superset/mcp_service/chart/chart_utils.py

@@ -23,6 +23,7 @@ generation that can be used by both generate_chart and generate_explore_link too
 """
 
 import logging
+from dataclasses import dataclass
 from typing import Any, Dict
 
 from superset.mcp_service.chart.schemas import (
@@ -38,6 +39,113 @@ from superset.utils import json
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class DatasetValidationResult:
+    """Result of dataset accessibility validation."""
+
+    is_valid: bool
+    dataset_id: int | str | None
+    dataset_name: str | None
+    warnings: list[str]
+    error: str | None = None
+
+
+def validate_chart_dataset(
+    chart: Any,
+    check_access: bool = True,
+) -> DatasetValidationResult:
+    """
+    Validate that a chart's dataset exists and is accessible.
+
+    This shared utility should be called by MCP tools after creating or retrieving
+    charts to detect issues like missing or deleted datasets early.
+
+    Args:
+        chart: A chart-like object with datasource_id, datasource_type attributes
+        check_access: Whether to also check user permissions (default True)
+
+    Returns:
+        DatasetValidationResult with validation status and any warnings
+    """
+    from sqlalchemy.exc import SQLAlchemyError
+
+    from superset.daos.dataset import DatasetDAO
+    from superset.mcp_service.auth import has_dataset_access
+
+    warnings: list[str] = []
+    datasource_id = getattr(chart, "datasource_id", None)
+
+    # Check if chart has a datasource reference
+    if datasource_id is None:
+        return DatasetValidationResult(
+            is_valid=False,
+            dataset_id=None,
+            dataset_name=None,
+            warnings=[],
+            error="Chart has no dataset reference (datasource_id is None)",
+        )
+
+    # Try to look up the dataset
+    try:
+        dataset = DatasetDAO.find_by_id(datasource_id)
+
+        if dataset is None:
+            return DatasetValidationResult(
+                is_valid=False,
+                dataset_id=datasource_id,
+                dataset_name=None,
+                warnings=[],
+                error=(
+                    f"Dataset (ID: {datasource_id}) has been deleted or does not "
+                    f"exist. The chart will not render correctly. "
+                    f"Consider updating the chart to use a different dataset."
+                ),
+            )
+
+        dataset_name = getattr(dataset, "table_name", None) or getattr(
+            dataset, "name", None
+        )
+
+        # Check if it's a virtual dataset (SQL Lab query)
+        is_virtual = bool(getattr(dataset, "sql", None))
+        if is_virtual:
+            warnings.append(
+                f"This chart uses a virtual dataset (SQL-based). "
+                f"If the dataset '{dataset_name}' is deleted, this chart will break."
+            )
+
+        # Check access permissions if requested
+        if check_access and not has_dataset_access(dataset):
+            return DatasetValidationResult(
+                is_valid=False,
+                dataset_id=datasource_id,
+                dataset_name=dataset_name,
+                warnings=warnings,
+                error=(
+                    f"Access denied to dataset '{dataset_name}' (ID: {datasource_id}). "
+                    f"You do not have permission to view this dataset."
+                ),
+            )
+
+        return DatasetValidationResult(
+            is_valid=True,
+            dataset_id=datasource_id,
+            dataset_name=dataset_name,
+            warnings=warnings,
+            error=None,
+        )
+
+    except (AttributeError, ValueError, RuntimeError, SQLAlchemyError) as e:
+        logger.exception("Error validating chart dataset %s: %s", datasource_id, e)
+        return DatasetValidationResult(
+            is_valid=False,
+            dataset_id=datasource_id,
+            dataset_name=None,
+            warnings=[],
+            error=f"Error validating dataset (ID: {datasource_id}): {str(e)}",
+        )
+
+
 def generate_explore_link(dataset_id: int | str, form_data: Dict[str, Any]) -> str:
     """Generate an explore link for the given dataset and form data."""
     from sqlalchemy.exc import SQLAlchemyError
@@ -45,6 +153,7 @@ def generate_explore_link(dataset_id: int | str, form_data: Dict[str, Any]) -> s
     from superset.commands.exceptions import CommandException
     from superset.commands.explore.form_data.parameters import CommandParameters
     from superset.daos.dataset import DatasetDAO
+    from superset.exceptions import SupersetException
     from superset.mcp_service.commands.create_form_data import (
         MCPCreateFormDataCommand,
     )
@@ -95,8 +204,17 @@ def generate_explore_link(dataset_id: int | str, form_data: Dict[str, Any]) -> s
         # Return URL with just the form_data_key
         return f"{base_url}/explore/?form_data_key={form_data_key}"
 
-    except (CommandException, SQLAlchemyError, ValueError, AttributeError):
+    except (
+        CommandException,
+        SupersetException,
+        SQLAlchemyError,
+        KeyError,
+        ValueError,
+        AttributeError,
+        TypeError,
+    ) as e:
         # Fallback to basic explore URL with numeric ID if available
+        logger.debug("Explore link generation fallback due to: %s", e)
         if numeric_dataset_id is not None:
             return (
                 f"{base_url}/explore/?datasource_type=table"

superset/mcp_service/chart/tool/generate_chart.py

@@ -25,6 +25,8 @@ from urllib.parse import parse_qs, urlparse
 from fastmcp import Context
 from superset_core.mcp import tool
 
+from superset.commands.exceptions import CommandException
+from superset.exceptions import SupersetException
 from superset.extensions import event_logger
 from superset.mcp_service.auth import has_dataset_access
 from superset.mcp_service.chart.chart_utils import (
@@ -32,6 +34,7 @@ from superset.mcp_service.chart.chart_utils import (
     analyze_chart_semantics,
     generate_chart_name,
     map_config_to_form_data,
+    validate_chart_dataset,
 )
 from superset.mcp_service.chart.schemas import (
     AccessibilityMetadata,
@@ -185,6 +188,7 @@ async def generate_chart(  # noqa: C901
         chart_id = None
         explore_url = None
         form_data_key = None
+        response_warnings: list[str] = []
 
         # Save chart by default (unless save_chart=False)
         if request.save_chart:
@@ -300,7 +304,25 @@ async def generate_chart(  # noqa: C901
                     )
                 )
 
-            except Exception as e:
+                # Post-creation validation: verify the chart's dataset is accessible
+                dataset_check = validate_chart_dataset(chart, check_access=True)
+                if not dataset_check.is_valid:
+                    # Dataset validation failed - warn but don't fail the operation
+                    await ctx.warning(
+                        "Chart created but dataset validation failed: %s"
+                        % (dataset_check.error,)
+                    )
+                    logger.warning(
+                        "Chart %s created but dataset validation failed: %s",
+                        chart.id,
+                        dataset_check.error,
+                    )
+                    if dataset_check.error:
+                        response_warnings.append(dataset_check.error)
+                # Add any validation warnings (e.g., virtual dataset warnings)
+                response_warnings.extend(dataset_check.warnings)
+
+            except CommandException as e:
                 logger.error("Chart creation failed: %s", e)
                 await ctx.error("Chart creation failed: error=%s" % (str(e),))
                 raise
@@ -338,7 +360,7 @@ async def generate_chart(  # noqa: C901
                         "Generated form_data_key for saved chart: "
                         "form_data_key=%s" % (form_data_key,)
                     )
-            except Exception as fdk_error:
+            except CommandException as fdk_error:
                 logger.warning(
                     "Failed to generate form_data_key for saved chart: %s",
                     fdk_error,
@@ -456,7 +478,7 @@ async def generate_chart(  # noqa: C901
                                 if not hasattr(preview_result, "error"):
                                     previews[format_type] = preview_result
 
-            except Exception as e:
+            except (CommandException, ValueError, KeyError) as e:
                 # Log warning but don't fail the entire request
                 await ctx.warning("Preview generation failed: error=%s" % (str(e),))
                 logger.warning("Preview generation failed: %s", e)
@@ -507,8 +529,8 @@ async def generate_chart(  # noqa: C901
             else {},
             "performance": performance.model_dump() if performance else None,
             "accessibility": accessibility.model_dump() if accessibility else None,
-            # Runtime warnings (informational suggestions, not errors)
-            "warnings": runtime_warnings,
+            # Combined runtime and response warnings
+            "warnings": runtime_warnings + response_warnings,
             "success": True,
             "schema_version": "2.0",
             "api_version": "v1",
@@ -523,7 +545,7 @@ async def generate_chart(  # noqa: C901
         )
         return GenerateChartResponse.model_validate(result)
 
-    except Exception as e:
+    except (CommandException, SupersetException, ValueError, KeyError) as e:
         await ctx.error(
             "Chart generation failed: error=%s, execution_time_ms=%s"
             % (
@@ -540,7 +562,7 @@ async def generate_chart(  # noqa: C901
         try:
             if hasattr(request, "config") and hasattr(request.config, "chart_type"):
                 chart_type = request.config.chart_type
-        except Exception as extract_error:
+        except AttributeError as extract_error:
             # Ignore errors when extracting chart type for error context
             logger.debug("Could not extract chart type: %s", extract_error)
 

superset/mcp_service/chart/tool/get_chart_data.py

@@ -32,7 +32,9 @@ if TYPE_CHECKING:
 
 from superset.commands.exceptions import CommandException
 from superset.commands.explore.form_data.parameters import CommandParameters
+from superset.exceptions import SupersetException
 from superset.extensions import event_logger
+from superset.mcp_service.chart.chart_utils import validate_chart_dataset
 from superset.mcp_service.chart.schemas import (
     ChartData,
     ChartError,
@@ -147,6 +149,22 @@ async def get_chart_data(  # noqa: C901
         )
         logger.info("Getting data for chart %s: %s", chart.id, chart.slice_name)
 
+        # Validate the chart's dataset is accessible before retrieving data
+        validation_result = validate_chart_dataset(chart, check_access=True)
+        if not validation_result.is_valid:
+            await ctx.warning(
+                "Chart found but dataset is not accessible: %s"
+                % (validation_result.error,)
+            )
+            return ChartError(
+                error=validation_result.error
+                or "Chart's dataset is not accessible. Dataset may have been deleted.",
+                error_type="DatasetNotAccessible",
+            )
+        # Log any warnings (e.g., virtual dataset warnings)
+        for warning in validation_result.warnings:
+            await ctx.warning("Dataset warning: %s" % (warning,))
+
         start_time = time.time()
 
         # Track whether we're using unsaved state
@@ -658,7 +676,7 @@ async def get_chart_data(  # noqa: C901
                 cache_status=cache_status,
             )
 
-        except Exception as data_error:
+        except (CommandException, SupersetException, ValueError) as data_error:
             await ctx.error(
                 "Data retrieval failed: chart_id=%s, error=%s, error_type=%s"
                 % (
@@ -673,7 +691,7 @@ async def get_chart_data(  # noqa: C901
                 error_type="DataError",
             )
 
-    except Exception as e:
+    except (CommandException, SupersetException, ValueError, KeyError) as e:
         await ctx.error(
             "Chart data retrieval failed: identifier=%s, error=%s, error_type=%s"
             % (

superset/mcp_service/chart/tool/get_chart_info.py

@@ -28,6 +28,7 @@ from superset_core.mcp import tool
 from superset.commands.exceptions import CommandException
 from superset.commands.explore.form_data.parameters import CommandParameters
 from superset.extensions import event_logger
+from superset.mcp_service.chart.chart_utils import validate_chart_dataset
 from superset.mcp_service.chart.schemas import (
     ChartError,
     ChartInfo,
@@ -158,6 +159,25 @@ async def get_chart_info(
             "Chart information retrieved successfully: chart_name=%s, "
             "is_unsaved_state=%s" % (result.slice_name, result.is_unsaved_state)
         )
+
+        # Validate the chart's dataset is accessible
+        if result.id:
+            chart = ChartDAO.find_by_id(result.id)
+            if chart:
+                validation_result = validate_chart_dataset(chart, check_access=True)
+                if not validation_result.is_valid:
+                    await ctx.warning(
+                        "Chart found but dataset is not accessible: %s"
+                        % (validation_result.error,)
+                    )
+                    return ChartError(
+                        error=validation_result.error
+                        or "Chart's dataset is not accessible",
+                        error_type="DatasetNotAccessible",
+                    )
+                # Log any warnings (e.g., virtual dataset warnings)
+                for warning in validation_result.warnings:
+                    await ctx.warning("Dataset warning: %s" % (warning,))
     else:
         await ctx.warning("Chart retrieval failed: error=%s" % (str(result),))
 

superset/mcp_service/chart/tool/get_chart_preview.py

@@ -25,7 +25,10 @@ from typing import Any, Dict, List, Protocol
 from fastmcp import Context
 from superset_core.mcp import tool
 
+from superset.commands.exceptions import CommandException
+from superset.exceptions import SupersetException
 from superset.extensions import event_logger
+from superset.mcp_service.chart.chart_utils import validate_chart_dataset
 from superset.mcp_service.chart.schemas import (
     AccessibilityMetadata,
     ASCIIPreview,
@@ -176,7 +179,14 @@ class ASCIIPreviewStrategy(PreviewFormatStrategy):
                 height=self.request.ascii_height or 20,
             )
 
-        except Exception as e:
+        except (
+            CommandException,
+            SupersetException,
+            ValueError,
+            KeyError,
+            AttributeError,
+            TypeError,
+        ) as e:
             logger.error("ASCII preview generation failed: %s", e)
             return ChartError(
                 error=f"Failed to generate ASCII preview: {str(e)}",
@@ -237,7 +247,14 @@ class TablePreviewStrategy(PreviewFormatStrategy):
                 row_count=len(data),
             )
 
-        except Exception as e:
+        except (
+            CommandException,
+            SupersetException,
+            ValueError,
+            KeyError,
+            AttributeError,
+            TypeError,
+        ) as e:
             logger.error("Table preview generation failed: %s", e)
             return ChartError(
                 error=f"Failed to generate table preview: {str(e)}",
@@ -256,7 +273,7 @@ class VegaLitePreviewStrategy(PreviewFormatStrategy):
 
                 return utils_json.loads(self.chart.params)
             return None
-        except Exception:
+        except (ValueError, TypeError):
             return None
 
     def generate(self) -> VegaLitePreview | ChartError:
@@ -345,7 +362,14 @@ class VegaLitePreviewStrategy(PreviewFormatStrategy):
                 supports_streaming=False,
             )
 
-        except Exception as e:
+        except (
+            CommandException,
+            SupersetException,
+            ValueError,
+            KeyError,
+            AttributeError,
+            TypeError,
+        ) as e:
             logger.exception(
                 "Error generating Vega-Lite preview for chart %s", self.chart.id
             )
@@ -455,11 +479,11 @@ class VegaLitePreviewStrategy(PreviewFormatStrategy):
                     field_type = self._determine_field_type(sample_values)
                     field_types[field] = field_type
 
-                except Exception as e:
+                except (TypeError, ValueError, KeyError, AttributeError) as e:
                     logger.warning("Error analyzing field '%s': %s", field, e)
                     field_types[field] = "nominal"  # Safe default
 
-        except Exception as e:
+        except (TypeError, ValueError, KeyError, AttributeError) as e:
             logger.warning("Error in field type analysis: %s", e)
             # Return nominal types for all fields as fallback
             return dict.fromkeys(fields, "nominal")
@@ -945,7 +969,7 @@ def generate_ascii_chart(
                 "Unsupported chart type '%s', falling back to table", chart_type
             )
             return _generate_ascii_table(data, width)
-    except Exception as e:
+    except (TypeError, ValueError, KeyError, IndexError) as e:
         logger.error("ASCII chart generation failed: %s", e)
         import traceback
 
@@ -1866,7 +1890,13 @@ async def _get_chart_preview_internal(  # noqa: C901
                                     self.uuid = None
 
                             chart = TransientChart(form_data)
-                    except (ValueError, KeyError, AttributeError, TypeError) as e:
+                    except (
+                        CommandException,
+                        ValueError,
+                        KeyError,
+                        AttributeError,
+                        TypeError,
+                    ) as e:
                         # Form data key not found or invalid
                         logger.debug(
                             "Failed to get form data for key %s: %s",
@@ -1900,6 +1930,24 @@ async def _get_chart_preview_internal(  # noqa: C901
         logger.info("Generating preview for chart %s", getattr(chart, "id", "NO_ID"))
         logger.info("Chart datasource_id: %s", getattr(chart, "datasource_id", "NONE"))
 
+        # Validate the chart's dataset is accessible before generating preview
+        # Skip validation for transient charts (no ID) - different data sources
+        if getattr(chart, "id", None) is not None:
+            validation_result = validate_chart_dataset(chart, check_access=True)
+            if not validation_result.is_valid:
+                await ctx.warning(
+                    "Chart found but dataset is not accessible: %s"
+                    % (validation_result.error,)
+                )
+                return ChartError(
+                    error=validation_result.error
+                    or "Chart's dataset is not accessible. Dataset may be deleted.",
+                    error_type="DatasetNotAccessible",
+                )
+            # Log any warnings (e.g., virtual dataset warnings)
+            for warning in validation_result.warnings:
+                await ctx.warning("Dataset warning: %s" % (warning,))
+
         import time
 
         start_time = time.time()
@@ -1990,7 +2038,14 @@ async def _get_chart_preview_internal(  # noqa: C901
 
         return result
 
-    except Exception as e:
+    except (
+        CommandException,
+        SupersetException,
+        ValueError,
+        KeyError,
+        AttributeError,
+        TypeError,
+    ) as e:
         await ctx.error(
             "Chart preview generation failed: identifier=%s, format=%s, error=%s, "
             "error_type=%s"
@@ -2055,7 +2110,14 @@ async def get_chart_preview(
             )
 
         return result
-    except Exception as e:
+    except (
+        CommandException,
+        SupersetException,
+        ValueError,
+        KeyError,
+        AttributeError,
+        TypeError,
+    ) as e:
         await ctx.error(
             "Chart preview generation failed: identifier=%s, error=%s, error_type=%s"
             % (

tests/unit_tests/mcp_service/chart/test_chart_utils.py

@@ -18,7 +18,7 @@
 """Tests for chart utilities module"""
 
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -32,6 +32,7 @@ from superset.mcp_service.chart.chart_utils import (
     map_filter_operator,
     map_table_config,
     map_xy_config,
+    validate_chart_dataset,
 )
 from superset.mcp_service.chart.schemas import (
     AxisConfig,
@@ -1045,3 +1046,101 @@ class TestFilterConfigValidation:
         """Test IN operator with empty list"""
         f = FilterConfig(column="platform", op="IN", value=[])
         assert f.value == []
+
+
+class TestValidateChartDataset:
+    """Test validate_chart_dataset function"""
+
+    @patch("superset.mcp_service.auth.has_dataset_access")
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    def test_validate_chart_dataset_no_datasource_id(
+        self, mock_find: MagicMock, mock_access: MagicMock
+    ) -> None:
+        """Chart with no datasource_id returns invalid result."""
+        chart = MagicMock(spec=[])  # no datasource_id attribute
+        result = validate_chart_dataset(chart)
+        assert not result.is_valid
+        assert result.dataset_id is None
+        assert "no dataset reference" in (result.error or "").lower()
+        mock_find.assert_not_called()
+
+    @patch("superset.mcp_service.auth.has_dataset_access")
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None)
+    def test_validate_chart_dataset_deleted_dataset(
+        self, mock_find: MagicMock, mock_access: MagicMock
+    ) -> None:
+        """Chart whose dataset was deleted returns invalid result."""
+        chart = MagicMock()
+        chart.datasource_id = 42
+        result = validate_chart_dataset(chart)
+        assert not result.is_valid
+        assert result.dataset_id == 42
+        assert "deleted" in (result.error or "").lower()
+
+    @patch("superset.mcp_service.auth.has_dataset_access", return_value=True)
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    def test_validate_chart_dataset_valid(
+        self, mock_find: MagicMock, mock_access: MagicMock
+    ) -> None:
+        """Valid chart with accessible dataset returns valid result."""
+        dataset = MagicMock()
+        dataset.table_name = "my_table"
+        dataset.sql = None
+        mock_find.return_value = dataset
+        chart = MagicMock()
+        chart.datasource_id = 7
+        result = validate_chart_dataset(chart)
+        assert result.is_valid
+        assert result.dataset_id == 7
+        assert result.dataset_name == "my_table"
+        assert result.warnings == []
+
+    @patch("superset.mcp_service.auth.has_dataset_access", return_value=True)
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    def test_validate_chart_dataset_virtual_warns(
+        self, mock_find: MagicMock, mock_access: MagicMock
+    ) -> None:
+        """Virtual dataset emits a warning."""
+        dataset = MagicMock()
+        dataset.table_name = "virt_ds"
+        dataset.sql = "SELECT 1"
+        mock_find.return_value = dataset
+        chart = MagicMock()
+        chart.datasource_id = 10
+        result = validate_chart_dataset(chart)
+        assert result.is_valid
+        assert len(result.warnings) == 1
+        assert "virtual" in result.warnings[0].lower()
+
+    @patch("superset.mcp_service.auth.has_dataset_access")
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    def test_validate_chart_dataset_sqlalchemy_error(
+        self, mock_find: MagicMock, mock_access: MagicMock
+    ) -> None:
+        """SQLAlchemy errors are caught and produce an invalid result."""
+        from sqlalchemy.exc import SQLAlchemyError
+
+        mock_find.side_effect = SQLAlchemyError("connection lost")
+        chart = MagicMock()
+        chart.datasource_id = 99
+        result = validate_chart_dataset(chart)
+        assert not result.is_valid
+        assert result.dataset_id == 99
+        assert "error" in (result.error or "").lower()
+
+    @patch(
+        "superset.mcp_service.chart.chart_utils.get_superset_base_url",
+        return_value="http://localhost:8088",
+    )
+    @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+    def test_generate_explore_link_sqlalchemy_error(
+        self,
+        mock_find: MagicMock,
+        mock_base_url: MagicMock,
+    ) -> None:
+        """SQLAlchemy errors in generate_explore_link fall back to basic URL."""
+        from sqlalchemy.exc import SQLAlchemyError
+
+        mock_find.side_effect = SQLAlchemyError("db gone")
+        url = generate_explore_link(5, {"viz_type": "table"})
+        assert "datasource_id=5" in url