feat(mcp): heatmap chart type plugin

superset/common/form_data_query_context.py

@@ -130,6 +130,23 @@ def freeform_where_having(form_data: dict[str, Any]) -> dict[str, str]:
     return extras
 
 
+def _as_column_list(value: Any) -> list[Any]:
+    """
+    Normalize a ``groupby``/``columns`` value into a list.
+
+    Single-select controls (e.g. the heatmap ``groupby`` Y axis, which is
+    ``multi: false``, and heatmap charts migrated via ``MigrateHeatmapChart``)
+    store the dimension as a bare string. Wrap a scalar in a one-element list,
+    mirroring ``chart_helpers.resolve_groupby``, so downstream list operations
+    (``.copy()``, ``.insert()``) do not blow up on a ``str``.
+    """
+    if value is None:
+        return []
+    if isinstance(value, str):
+        return [value]
+    return list(value)
+
+
 def columns_from_form_data(form_data: dict[str, Any]) -> list[Any]:
     """
     Derive the query's grouping/raw columns from form data.
@@ -141,10 +158,10 @@ def columns_from_form_data(form_data: dict[str, Any]) -> list[Any]:
     if form_data.get("query_mode") == "raw" and (
         form_data.get("all_columns") or form_data.get("columns")
     ):
-        return list(form_data.get("all_columns") or form_data.get("columns") or [])
+        return _as_column_list(form_data.get("all_columns") or form_data.get("columns"))
 
-    groupby_columns: list[Any] = form_data.get("groupby") or []
-    raw_columns: list[Any] = form_data.get("columns") or []
+    groupby_columns: list[Any] = _as_column_list(form_data.get("groupby"))
+    raw_columns: list[Any] = _as_column_list(form_data.get("columns"))
     # Prefer explicit raw columns only when they are actually present; a stale
     # empty ``columns: []`` key must not shadow the group-by dimensions (which
     # would silently drop the grouping and change the aggregation).

superset/mcp_service/app.py

@@ -448,11 +448,11 @@ def get_default_instructions(
 chart_type_display_name field with a human-readable name when available.
 This field is populated for chart types known to the MCP registry
 (xy, pie, table, pivot_table, big_number, mixed_timeseries, handlebars,
-histogram, box_plot, waterfall, gantt, bubble_v2, and interactive_pivot).
-Availability gates creation and schema discovery, not display names for
-existing charts.
-For all other viz_types (Funnel, Gauge, Heatmap, etc.) it will be null —
-use the raw viz_type field instead when referring to those chart types.
+histogram, box_plot, waterfall, gantt, bubble_v2, heatmap_v2, and
+interactive_pivot). Availability gates creation and schema discovery, not
+display names for existing charts. For all other viz_types it will be
+null — use the raw viz_type field instead when referring to those chart
+types.
 
 Query Examples:
 - List all tables:

superset/mcp_service/chart/chart_helpers.py

@@ -882,7 +882,10 @@ def build_query_dicts_from_form_data(
                 qd["filters"] = [*(qd.get("filters") or []), *null_filters]
         return [qd]
 
-    if viz_type.startswith("echarts_timeseries"):
+    # Heatmap puts its x_axis in the query columns too: the frontend folds it
+    # into groupby in buildQuery, and the MCP path builds the query dict
+    # directly, so without this the x axis never reaches GROUP BY.
+    if viz_type.startswith("echarts_timeseries") or viz_type == "heatmap_v2":
         groupby = with_x_axis_column(form_data, groupby)
 
     return [

superset/mcp_service/chart/chart_utils.py

@@ -48,6 +48,7 @@
     GanttChartConfig,
     GaugeChartConfig,
     HandlebarsChartConfig,
+    HeatmapChartConfig,
     HistogramChartConfig,
     MixedTimeseriesChartConfig,
     PieChartConfig,
@@ -1694,6 +1695,27 @@ def map_bubble_config(config: BubbleChartConfig) -> Dict[str, Any]:
     return form_data
 
 
+def map_heatmap_config(config: HeatmapChartConfig) -> Dict[str, Any]:
+    """Map heatmap config to Superset form_data (viz_type ``heatmap_v2``).
+
+    Matches the frontend Heatmap buildQuery contract: an ``x_axis`` column and
+    a single ``groupby`` Y column form the two axes, one ``metric`` colours
+    the cells, and ``normalize_across`` selects the rank-normalization range.
+    The Y axis is a single-select ``groupby`` (not a list).
+    """
+    form_data: Dict[str, Any] = {
+        "viz_type": "heatmap_v2",
+        "x_axis": config.x_axis.name,
+        "groupby": config.y_axis.name,
+        "metric": create_metric_object(config.metric),
+        "normalize_across": config.normalize_across,
+        "normalized": config.normalized,
+        "row_limit": config.row_limit,
+    }
+    _add_adhoc_filters(form_data, config.filters)
+    return form_data
+
+
 def map_histogram_config(config: "HistogramChartConfig") -> Dict[str, Any]:
     """Map histogram config to Superset form_data (viz_type histogram_v2).
 
@@ -2315,6 +2337,14 @@ def _bubble_chart_what(config: BubbleChartConfig) -> str:
     return f"{config.entity.name}: {x_label} vs {y_label}"
 
 
+def _heatmap_chart_what(config: HeatmapChartConfig) -> str:
+    """Build the 'what' portion for a heatmap chart name."""
+    metric_label = (
+        config.metric.label or config.metric.name or config.metric.sql_expression
+    )
+    return f"{config.x_axis.name} vs {config.y_axis.name} by {metric_label}"
+
+
 def _pivot_table_what(config: PivotTableChartConfig) -> str:
     """Build the 'what' portion for a pivot table chart name."""
     # Pivot rows reject sql_expression at validation, so name is set.

superset/mcp_service/chart/plugins/__init__.py

@@ -33,6 +33,7 @@
 from superset.mcp_service.chart.plugins.gantt import GanttChartPlugin
 from superset.mcp_service.chart.plugins.gauge import GaugeChartPlugin
 from superset.mcp_service.chart.plugins.handlebars import HandlebarsChartPlugin
+from superset.mcp_service.chart.plugins.heatmap import HeatmapChartPlugin
 from superset.mcp_service.chart.plugins.histogram import HistogramChartPlugin
 from superset.mcp_service.chart.plugins.interactive_pivot import (
     InteractivePivotChartPlugin,
@@ -62,6 +63,7 @@
 register(BigNumberChartPlugin())
 register(HistogramChartPlugin())
 register(BoxPlotChartPlugin())
+register(HeatmapChartPlugin())
 register(WaterfallChartPlugin())
 register(GanttChartPlugin())
 
@@ -72,6 +74,7 @@
     "GanttChartPlugin",
     "GaugeChartPlugin",
     "HandlebarsChartPlugin",
+    "HeatmapChartPlugin",
     "HistogramChartPlugin",
     "InteractivePivotChartPlugin",
     "MixedTimeseriesChartPlugin",

superset/mcp_service/chart/plugins/heatmap.py

@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Heatmap chart type plugin."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, ClassVar
+
+from superset.mcp_service.chart.chart_utils import (
+    _heatmap_chart_what,
+    _summarize_filters,
+    map_heatmap_config,
+)
+from superset.mcp_service.chart.plugin import BaseChartPlugin
+from superset.mcp_service.chart.schemas import ColumnRef, HeatmapChartConfig
+from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator
+from superset.mcp_service.common.error_schemas import ChartGenerationError
+
+
+class HeatmapChartPlugin(BaseChartPlugin):
+    """Plugin for heatmap chart type."""
+
+    chart_type = "heatmap_v2"
+    display_name = "Heatmap"
+    native_viz_types: ClassVar[Mapping[str, str]] = {
+        "heatmap_v2": "Heatmap",
+    }
+
+    def pre_validate(
+        self,
+        config: dict[str, Any],
+    ) -> ChartGenerationError | None:
+        missing_fields = []
+
+        if "x_axis" not in config:
+            missing_fields.append("'x_axis' (column along the X axis)")
+        if "y_axis" not in config and "groupby" not in config:
+            missing_fields.append("'y_axis' (column along the Y axis)")
+        if "metric" not in config:
+            missing_fields.append("'metric' (value colouring each cell)")
+
+        if missing_fields:
+            return ChartGenerationError(
+                error_type="missing_heatmap_fields",
+                message=(
+                    f"Heatmap chart missing required fields: "
+                    f"{', '.join(missing_fields)}"
+                ),
+                details=(
+                    "Heatmaps plot a metric across two dimensions — one on the "
+                    "x_axis and one on the y_axis — colouring each cell by the "
+                    "metric value"
+                ),
+                suggestions=[
+                    "Add 'x_axis': {'name': 'day_of_week'}",
+                    "Add 'y_axis': {'name': 'hour'}",
+                    "Add 'metric': {'name': 'trips', 'aggregate': 'COUNT'}",
+                    "Example: {'chart_type': 'heatmap_v2', "
+                    "'x_axis': {'name': 'day_of_week'}, "
+                    "'y_axis': {'name': 'hour'}, "
+                    "'metric': {'name': 'trips', 'aggregate': 'COUNT'}}",
+                ],
+                error_code="MISSING_HEATMAP_FIELDS",
+            )
+
+        return None
+
+    def extract_column_refs(self, config: Any) -> list[ColumnRef]:
+        if not isinstance(config, HeatmapChartConfig):
+            return []
+        refs: list[ColumnRef] = [config.x_axis, config.y_axis, config.metric]
+        if config.filters:
+            for f in config.filters:
+                refs.append(ColumnRef(name=f.column))
+        return refs
+
+    def to_form_data(
+        self, config: Any, dataset_id: int | str | None = None
+    ) -> dict[str, Any]:
+        return map_heatmap_config(config)
+
+    def generate_name(self, config: Any, dataset_name: str | None = None) -> str:
+        what = _heatmap_chart_what(config)
+        context = _summarize_filters(config.filters)
+        return self._with_context(what, context)
+
+    def resolve_viz_type(self, config: Any) -> str:
+        return "heatmap_v2"
+
+    def normalize_column_refs(self, config: Any, dataset_context: Any) -> Any:
+        config_dict = config.model_dump(exclude_unset=True)
+
+        for key in ("x_axis", "y_axis"):
+            col = config_dict.get(key)
+            if col and not col.get("sql_expression") and not col.get("saved_metric"):
+                col["name"] = DatasetValidator.get_canonical_column_name(
+                    col["name"], dataset_context
+                )
+        if config_dict.get("metric"):
+            if config_dict["metric"].get("sql_expression"):
+                pass
+            elif config_dict["metric"].get("saved_metric"):
+                config_dict["metric"]["name"] = (
+                    DatasetValidator.get_canonical_metric_name(
+                        config_dict["metric"]["name"], dataset_context
+                    )
+                )
+            else:
+                config_dict["metric"]["name"] = (
+                    DatasetValidator.get_canonical_column_name(
+                        config_dict["metric"]["name"], dataset_context
+                    )
+                )
+        DatasetValidator.normalize_filters(config_dict, dataset_context)
+        return HeatmapChartConfig.model_validate(config_dict)
+
+    def schema_error_hint(self) -> ChartGenerationError | None:
+        return ChartGenerationError(
+            error_type="heatmap_validation_error",
+            message="Heatmap chart configuration validation failed",
+            details=(
+                "The heatmap chart configuration is missing required "
+                "fields or has invalid structure"
+            ),
+            suggestions=[
+                "Ensure 'x_axis' and 'y_axis' each have a 'name'",
+                "Ensure 'metric' field has 'name' and 'aggregate'",
+                "Example: {'chart_type': 'heatmap_v2', "
+                "'x_axis': {'name': 'day_of_week'}, "
+                "'y_axis': {'name': 'hour'}, "
+                "'metric': {'name': 'trips', 'aggregate': 'COUNT'}}",
+            ],
+            error_code="HEATMAP_VALIDATION_ERROR",
+        )

superset/mcp_service/chart/schemas.py

@@ -1727,6 +1727,67 @@ def record_implicit_metric_aggregate(self) -> "BubbleChartConfig":
         return self
 
 
+class HeatmapChartConfig(BaseChartConfig):
+    """Config for heatmap charts (viz_type ``heatmap_v2``).
+
+    Matches the frontend Heatmap buildQuery contract: an ``x_axis`` column, a
+    single ``groupby`` column for the Y axis, and one ``metric`` colouring each
+    cell. ``normalize_across`` drives the server-side rank normalization
+    (whole heatmap, per-x, or per-y).
+    """
+
+    model_config = ConfigDict(extra="ignore", populate_by_name=True)
+
+    chart_type: Literal["heatmap_v2"] = "heatmap_v2"
+    x_axis: ColumnRef = Field(
+        ...,
+        description="Column along the X axis",
+    )
+    y_axis: ColumnRef = Field(
+        ...,
+        description="Column along the Y axis (form_data 'groupby'; single-select)",
+        validation_alias=AliasChoices("y_axis", "groupby"),
+    )
+    metric: ColumnRef = Field(
+        ...,
+        description="Value metric colouring each cell (use aggregate e.g. SUM, "
+        "COUNT for ad-hoc, or set saved_metric=True for a saved dataset metric)",
+    )
+    normalize_across: Literal["heatmap", "x", "y"] = Field(
+        "heatmap",
+        description="Range the cell colour is normalized against: the whole "
+        "'heatmap', each 'x' column, or each 'y' row (frontend default: "
+        "'heatmap'). Only takes effect when normalized=true.",
+    )
+    normalized: bool = Field(
+        False,
+        description="Colour cells by rank within 'normalize_across' rather than "
+        "the raw metric value. When false (the default) 'normalize_across' has "
+        "no visual effect.",
+    )
+    row_limit: int = Field(
+        10000, description="Max rows queried (cells = X × Y)", ge=1, le=100000
+    )
+    filters: List[FilterConfig] | None = Field(
+        None,
+        description="Structured filters (column/op/value). "
+        "Do NOT use adhoc_filters or raw SQL expressions.",
+    )
+
+    @model_validator(mode="after")
+    def reject_metric_style_dimensions(self) -> "HeatmapChartConfig":
+        """x_axis and y_axis are dimensions, not metrics."""
+        for col, name in ((self.x_axis, "x_axis"), (self.y_axis, "y_axis")):
+            _reject_sql_expression_on_dimension(col, name)
+            if col and col.is_metric:
+                raise ValueError(
+                    f"{name} must be a plain column, not a metric; drop "
+                    "'aggregate'/'saved_metric' (metrics belong in the 'metric' "
+                    "field)"
+                )
+        return self
+
+
 class PivotTableChartConfig(BaseChartConfig):
     model_config = ConfigDict(extra="ignore", populate_by_name=True)
 
@@ -3745,6 +3806,7 @@ def validate_gantt_roles(self) -> "GanttChartConfig":
     | GaugeChartConfig
     | TreemapChartConfig
     | BubbleChartConfig
+    | HeatmapChartConfig
     | PivotTableChartConfig
     | InteractivePivotChartConfig
     | MixedTimeseriesChartConfig
@@ -3758,7 +3820,8 @@ def validate_gantt_roles(self) -> "GanttChartConfig":
         discriminator=CHART_TYPE_DISCRIMINATOR,
         description=(
             "Chart configuration - specify chart_type as 'xy', 'table', "
-            "'pie', 'gauge', 'treemap_v2', 'bubble_v2', 'pivot_table', "
+            "'pie', 'gauge', 'treemap_v2', 'bubble_v2', 'heatmap_v2', "
+            "'pivot_table', "
             "'interactive_pivot', 'mixed_timeseries', 'handlebars', "
             "'big_number', 'histogram', 'box_plot', 'waterfall', or 'gantt'"
         ),