feat(mcp): treemap chart type plugin

superset/mcp_service/chart/chart_utils.py

@@ -51,6 +51,7 @@
     PivotTableChartConfig,
     SortByConfig,
     TableChartConfig,
+    TreemapChartConfig,
     WaterfallChartConfig,
     XYChartConfig,
 )
@@ -1260,6 +1261,25 @@ def map_gauge_config(config: GaugeChartConfig) -> Dict[str, Any]:
     return form_data
 
 
+def map_treemap_config(config: TreemapChartConfig) -> Dict[str, Any]:
+    """Map treemap config to Superset form_data (viz_type ``treemap_v2``).
+
+    Matches the frontend Treemap buildQuery contract: one ``metric`` plus an
+    ordered ``groupby`` hierarchy (first column outermost). When
+    ``sort_by_metric`` is set the query orders by the metric descending.
+    """
+    form_data: Dict[str, Any] = {
+        "viz_type": "treemap_v2",
+        "groupby": [g.name for g in config.groupby],
+        "metric": create_metric_object(config.metric),
+        "sort_by_metric": config.sort_by_metric,
+        "row_limit": config.row_limit,
+        "color_scheme": config.color_scheme or "supersetColors",
+    }
+    _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).
 
@@ -1782,6 +1802,17 @@ def _gauge_chart_what(config: GaugeChartConfig) -> str:
     return f"{metric_label}"
 
 
+def _treemap_chart_what(config: TreemapChartConfig) -> str:
+    """Build the 'what' portion for a treemap chart name."""
+    metric_label = (
+        config.metric.label or config.metric.name or config.metric.sql_expression
+    )
+    dims = ", ".join(g.name for g in config.groupby if g.name)
+    if dims:
+        return f"{dims} by {metric_label}"
+    return f"{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

@@ -41,6 +41,7 @@
 from superset.mcp_service.chart.plugins.pie import PieChartPlugin
 from superset.mcp_service.chart.plugins.pivot_table import PivotTableChartPlugin
 from superset.mcp_service.chart.plugins.table import TableChartPlugin
+from superset.mcp_service.chart.plugins.treemap import TreemapChartPlugin
 from superset.mcp_service.chart.plugins.waterfall import WaterfallChartPlugin
 from superset.mcp_service.chart.plugins.xy import XYChartPlugin
 from superset.mcp_service.chart.registry import register
@@ -50,6 +51,7 @@
 register(TableChartPlugin())
 register(PieChartPlugin())
 register(GaugeChartPlugin())
+register(TreemapChartPlugin())
 register(PivotTableChartPlugin())
 register(InteractivePivotChartPlugin())
 register(MixedTimeseriesChartPlugin())
@@ -70,6 +72,7 @@
     "PieChartPlugin",
     "PivotTableChartPlugin",
     "TableChartPlugin",
+    "TreemapChartPlugin",
     "WaterfallChartPlugin",
     "XYChartPlugin",
 ]

superset/mcp_service/chart/plugins/treemap.py

@@ -0,0 +1,144 @@
+# 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.
+
+"""Treemap 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 (
+    _summarize_filters,
+    _treemap_chart_what,
+    map_treemap_config,
+)
+from superset.mcp_service.chart.plugin import BaseChartPlugin
+from superset.mcp_service.chart.schemas import ColumnRef, TreemapChartConfig
+from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator
+from superset.mcp_service.common.error_schemas import ChartGenerationError
+
+
+class TreemapChartPlugin(BaseChartPlugin):
+    """Plugin for treemap chart type."""
+
+    chart_type = "treemap_v2"
+    display_name = "Treemap"
+    native_viz_types: ClassVar[Mapping[str, str]] = {
+        "treemap_v2": "Treemap",
+    }
+
+    def pre_validate(
+        self,
+        config: dict[str, Any],
+    ) -> ChartGenerationError | None:
+        missing_fields = []
+
+        if not config.get("groupby"):
+            missing_fields.append("'groupby' (ordered hierarchy columns)")
+        if "metric" not in config:
+            missing_fields.append("'metric' (value metric sizing the tiles)")
+
+        if missing_fields:
+            return ChartGenerationError(
+                error_type="missing_treemap_fields",
+                message=(
+                    f"Treemap chart missing required fields: "
+                    f"{', '.join(missing_fields)}"
+                ),
+                details=(
+                    "Treemaps size tiles by a metric and nest them by an "
+                    "ordered groupby hierarchy (first column is the outermost "
+                    "level)"
+                ),
+                suggestions=[
+                    "Add 'groupby': [{'name': 'region'}, {'name': 'product'}]",
+                    "Add 'metric': {'name': 'revenue', 'aggregate': 'SUM'}",
+                    "Example: {'chart_type': 'treemap_v2', "
+                    "'groupby': [{'name': 'region'}], "
+                    "'metric': {'name': 'revenue', 'aggregate': 'SUM'}}",
+                ],
+                error_code="MISSING_TREEMAP_FIELDS",
+            )
+
+        return None
+
+    def extract_column_refs(self, config: Any) -> list[ColumnRef]:
+        if not isinstance(config, TreemapChartConfig):
+            return []
+        refs: list[ColumnRef] = [*config.groupby, 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_treemap_config(config)
+
+    def generate_name(self, config: Any, dataset_name: str | None = None) -> str:
+        what = _treemap_chart_what(config)
+        context = _summarize_filters(config.filters)
+        return self._with_context(what, context)
+
+    def resolve_viz_type(self, config: Any) -> str:
+        return "treemap_v2"
+
+    def normalize_column_refs(self, config: Any, dataset_context: Any) -> Any:
+        config_dict = config.model_dump()
+
+        for col in config_dict.get("groupby") or []:
+            if 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 TreemapChartConfig.model_validate(config_dict)
+
+    def schema_error_hint(self) -> ChartGenerationError | None:
+        return ChartGenerationError(
+            error_type="treemap_validation_error",
+            message="Treemap chart configuration validation failed",
+            details=(
+                "The treemap chart configuration is missing required "
+                "fields or has invalid structure"
+            ),
+            suggestions=[
+                "Ensure 'groupby' has at least one column for the hierarchy",
+                "Ensure 'metric' field has 'name' and 'aggregate'",
+                "Example: {'chart_type': 'treemap_v2', "
+                "'groupby': [{'name': 'region'}], "
+                "'metric': {'name': 'revenue', 'aggregate': 'SUM'}}",
+            ],
+            error_code="TREEMAP_VALIDATION_ERROR",
+        )

superset/mcp_service/chart/schemas.py

@@ -1358,6 +1358,62 @@ def parse_numbers(value: str, field_name: str) -> list[float]:
         return self
 
 
+class TreemapChartConfig(BaseChartConfig):
+    """Config for treemap charts (viz_type ``treemap_v2``).
+
+    Matches the frontend Treemap buildQuery contract: one ``metric`` sizing
+    the tiles plus an ordered ``groupby`` hierarchy — the first column is the
+    outermost level and each subsequent column nests inside it. When
+    ``sort_by_metric`` is set, tiles are ordered by the metric descending.
+    """
+
+    model_config = ConfigDict(extra="ignore", populate_by_name=True)
+
+    chart_type: Literal["treemap_v2"] = "treemap_v2"
+    groupby: List[ColumnRef] = Field(
+        ...,
+        min_length=1,
+        description="Ordered category columns forming the treemap hierarchy "
+        "(first = outermost level; order defines nesting)",
+    )
+    metric: ColumnRef = Field(
+        ...,
+        description="Value metric sizing the tiles (use aggregate e.g. SUM, "
+        "COUNT for ad-hoc, or set saved_metric=True for a saved dataset metric)",
+    )
+    sort_by_metric: bool = Field(
+        True,
+        description="Order tiles by the metric descending (frontend default)",
+    )
+    row_limit: int = Field(100, description="Max rows queried", ge=1, le=10000)
+    filters: List[FilterConfig] | None = Field(
+        None,
+        description="Structured filters (column/op/value). "
+        "Do NOT use adhoc_filters or raw SQL expressions.",
+    )
+    color_scheme: str | None = Field(
+        None,
+        description=(
+            "Superset color scheme ID (e.g. 'supersetColors', 'lyftColors', "
+            "'googleCategory10c', 'd3Category10'). Defaults to 'supersetColors'."
+        ),
+        max_length=100,
+    )
+
+    @model_validator(mode="after")
+    def reject_metric_style_groupby(self) -> "TreemapChartConfig":
+        """groupby entries are hierarchy dimensions, not metrics."""
+        for i, col in enumerate(self.groupby or []):
+            _reject_sql_expression_on_dimension(col, f"groupby[{i}]")
+            if col.is_metric:
+                raise ValueError(
+                    f"groupby[{i}] 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)
 
@@ -2601,6 +2657,7 @@ def reject_metric_style_dimensions(self) -> "WaterfallChartConfig":
     | TableChartConfig
     | PieChartConfig
     | GaugeChartConfig
+    | TreemapChartConfig
     | PivotTableChartConfig
     | InteractivePivotChartConfig
     | MixedTimeseriesChartConfig
@@ -2613,7 +2670,7 @@ def reject_metric_style_dimensions(self) -> "WaterfallChartConfig":
         discriminator="chart_type",
         description=(
             "Chart configuration - specify chart_type as 'xy', 'table', "
-            "'pie', 'gauge', 'pivot_table', 'interactive_pivot', "
+            "'pie', 'gauge', 'treemap_v2', 'pivot_table', 'interactive_pivot', "
             "'mixed_timeseries', 'handlebars', "
             "'big_number', 'histogram', 'box_plot', or 'waterfall'"
         ),

superset/mcp_service/chart/tool/generate_chart.py

@@ -86,7 +86,8 @@ async def generate_chart(  # noqa: C901
     - LLM clients MUST display returned chart URL to users
     - Use numeric dataset ID or UUID (NOT schema.table_name format)
     - MUST include chart_type in config (one of: 'xy', 'table', 'pie',
-      'gauge', 'pivot_table', 'mixed_timeseries', 'handlebars',
+      'gauge', 'treemap_v2', 'pivot_table', 'mixed_timeseries',
+      'handlebars',
       'big_number', 'histogram', 'box_plot', 'waterfall', plus host-gated
       types returned by get_chart_type_schema such as 'interactive_pivot')
 
@@ -128,6 +129,9 @@ async def generate_chart(  # noqa: C901
       Required fields: metric; optional: groupby (one dial per value),
       min_val, max_val
 
+    - chart_type='treemap_v2' for hierarchical part-to-whole.
+      Required fields: groupby (ordered hierarchy), metric
+
     - chart_type='histogram' for value-distribution charts.
       Required fields: column (numeric); optional: bins, groupby, normalize,
       cumulative
@@ -151,6 +155,7 @@ async def generate_chart(  # noqa: C901
     - "compare two metrics over time" -> chart_type='mixed_timeseries'
     - "single number" / "KPI" / "scorecard" -> chart_type='big_number'
     - "gauge" / "dial" / "speedometer" -> chart_type='gauge'
+    - "treemap" / "hierarchy" -> chart_type='treemap_v2'
     - "custom HTML template" -> chart_type='handlebars'
     - "histogram" / "distribution" -> chart_type='histogram'
     - "box plot" / "box and whisker" -> chart_type='box_plot'

superset/mcp_service/chart/tool/get_chart_type_schema.py

@@ -39,6 +39,7 @@
     PieChartConfig,
     PivotTableChartConfig,
     TableChartConfig,
+    TreemapChartConfig,
     WaterfallChartConfig,
     XYChartConfig,
 )
@@ -51,6 +52,7 @@
     "table": TypeAdapter(TableChartConfig),
     "pie": TypeAdapter(PieChartConfig),
     "gauge": TypeAdapter(GaugeChartConfig),
+    "treemap_v2": TypeAdapter(TreemapChartConfig),
     "pivot_table": TypeAdapter(PivotTableChartConfig),
     "interactive_pivot": TypeAdapter(InteractivePivotChartConfig),
     "mixed_timeseries": TypeAdapter(MixedTimeseriesChartConfig),
@@ -212,6 +214,13 @@
             "metric": {"name": "progress", "aggregate": "AVG"},
         },
     ],
+    "treemap_v2": [
+        {
+            "chart_type": "treemap_v2",
+            "groupby": [{"name": "region"}, {"name": "product"}],
+            "metric": {"name": "revenue", "aggregate": "SUM"},
+        },
+    ],
 }
 
 
@@ -301,8 +310,9 @@ def get_chart_type_schema(
     for a chart configuration before calling generate_chart or update_chart.
 
     Valid chart_type values depend on the host deployment. Core types are xy,
-    table, pie, gauge, pivot_table, mixed_timeseries, handlebars, big_number,
-    histogram, box_plot, and waterfall. Deployments that enable an AG Grid
+    table, pie, gauge, treemap_v2, pivot_table, mixed_timeseries, handlebars,
+    big_number, histogram, box_plot, and waterfall. Deployments that enable an
+    AG Grid
     pivot extension also expose interactive_pivot.
 
     Returns the JSON Schema for the requested chart type, optionally