From 458631c26eb4515be0f6582f353523a409c16a90 Mon Sep 17 00:00:00 2001 From: Greg Neighbors Date: Wed, 9 Sep 2026 11:38:43 -0400 Subject: [PATCH] feat(mcp): treemap chart type plugin Adds a generate_chart plugin for treemap (viz_type 'treemap_v2'): one metric sizing the tiles plus an ordered groupby hierarchy; sort_by_metric orders tiles by the metric descending. Rebased onto current master (past the gauge plugin + its follow-up), and adds the generate_chart.py docstring entry (one-of list, per-type bullet, 'treemap'/'hierarchy' quick-lookup) plus the get_chart_type_schema core list. Co-Authored-By: Claude Opus 4.8 --- superset/mcp_service/chart/chart_utils.py | 31 +++ .../mcp_service/chart/plugins/__init__.py | 3 + superset/mcp_service/chart/plugins/treemap.py | 144 ++++++++++++++ superset/mcp_service/chart/schemas.py | 59 +++++- .../mcp_service/chart/tool/generate_chart.py | 7 +- .../chart/tool/get_chart_type_schema.py | 14 +- .../mcp_service/chart/test_treemap_chart.py | 176 ++++++++++++++++++ 7 files changed, 430 insertions(+), 4 deletions(-) create mode 100644 superset/mcp_service/chart/plugins/treemap.py create mode 100644 tests/unit_tests/mcp_service/chart/test_treemap_chart.py diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py index bc3de192924c..a9934d7727b1 100644 --- a/superset/mcp_service/chart/chart_utils.py +++ b/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. diff --git a/superset/mcp_service/chart/plugins/__init__.py b/superset/mcp_service/chart/plugins/__init__.py index d3440c191b54..247d265be8f5 100644 --- a/superset/mcp_service/chart/plugins/__init__.py +++ b/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", ] diff --git a/superset/mcp_service/chart/plugins/treemap.py b/superset/mcp_service/chart/plugins/treemap.py new file mode 100644 index 000000000000..e566377dcb80 --- /dev/null +++ b/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", + ) diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index 46f0054af103..5bee8c890d23 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/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'" ), diff --git a/superset/mcp_service/chart/tool/generate_chart.py b/superset/mcp_service/chart/tool/generate_chart.py index 0bef46e299be..bee231145a1e 100644 --- a/superset/mcp_service/chart/tool/generate_chart.py +++ b/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' diff --git a/superset/mcp_service/chart/tool/get_chart_type_schema.py b/superset/mcp_service/chart/tool/get_chart_type_schema.py index d3ed18a4baf4..18897163780a 100644 --- a/superset/mcp_service/chart/tool/get_chart_type_schema.py +++ b/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 diff --git a/tests/unit_tests/mcp_service/chart/test_treemap_chart.py b/tests/unit_tests/mcp_service/chart/test_treemap_chart.py new file mode 100644 index 000000000000..74b71088beab --- /dev/null +++ b/tests/unit_tests/mcp_service/chart/test_treemap_chart.py @@ -0,0 +1,176 @@ +# 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. + +"""Tests for the treemap chart type plugin. + +Schema validation, form_data mapping (matching the frontend Treemap +buildQuery contract for viz_type ``treemap_v2`` — one ``metric`` plus an +ordered multi ``groupby`` hierarchy), and registry integration. +""" + +import pytest +from pydantic import TypeAdapter, ValidationError + +from superset.mcp_service.chart.chart_utils import map_treemap_config +from superset.mcp_service.chart.schemas import ChartConfig, TreemapChartConfig + + +class TestTreemapChartConfigSchema: + """TreemapChartConfig schema validation.""" + + def test_basic_treemap_config(self) -> None: + config = TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region"}, {"name": "product"}], + metric={"name": "revenue", "aggregate": "SUM"}, + ) + assert [g.name for g in config.groupby] == ["region", "product"] + assert config.sort_by_metric is True # shared control default + + def test_treemap_missing_groupby(self) -> None: + with pytest.raises(ValidationError): + TreemapChartConfig( + chart_type="treemap_v2", + metric={"name": "revenue", "aggregate": "SUM"}, + ) + + def test_treemap_empty_groupby_rejected(self) -> None: + with pytest.raises(ValidationError): + TreemapChartConfig( + chart_type="treemap_v2", + groupby=[], + metric={"name": "revenue", "aggregate": "SUM"}, + ) + + def test_treemap_missing_metric(self) -> None: + with pytest.raises(ValidationError): + TreemapChartConfig(chart_type="treemap_v2", groupby=[{"name": "region"}]) + + def test_treemap_rejects_extra_fields(self) -> None: + with pytest.raises(ValidationError): + TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region"}], + metric={"name": "revenue", "aggregate": "SUM"}, + bogus=1, + ) + + def test_treemap_groupby_rejects_aggregate(self) -> None: + """An aggregate makes a hierarchy column metric-like; reject it.""" + with pytest.raises(ValidationError): + TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region", "aggregate": "SUM"}], + metric={"name": "revenue", "aggregate": "SUM"}, + ) + + def test_treemap_groupby_rejects_saved_metric(self) -> None: + with pytest.raises(ValidationError): + TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "count", "saved_metric": True}], + metric={"name": "revenue", "aggregate": "SUM"}, + ) + + def test_chart_config_union_dispatches_treemap(self) -> None: + config = TypeAdapter(ChartConfig).validate_python( + { + "chart_type": "treemap_v2", + "groupby": [{"name": "region"}], + "metric": {"name": "revenue", "aggregate": "SUM"}, + } + ) + assert isinstance(config, TreemapChartConfig) + + +class TestMapTreemapConfig: + """form_data mapping must match the frontend Treemap buildQuery.""" + + def test_basic_treemap_form_data(self) -> None: + config = TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region"}, {"name": "product"}], + metric={"name": "revenue", "aggregate": "SUM"}, + ) + form_data = map_treemap_config(config) + assert form_data["viz_type"] == "treemap_v2" + # hierarchy order is preserved (nesting levels) + assert form_data["groupby"] == ["region", "product"] + assert form_data["metric"]["label"] == "SUM(revenue)" + assert form_data["sort_by_metric"] is True + + def test_treemap_form_data_with_filters_and_no_sort(self) -> None: + config = TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region"}], + metric={"name": "revenue", "aggregate": "SUM"}, + sort_by_metric=False, + filters=[{"column": "year", "op": "=", "value": 2026}], + ) + form_data = map_treemap_config(config) + assert form_data["sort_by_metric"] is False + assert form_data["adhoc_filters"], "filters must map to adhoc_filters" + + def test_treemap_saved_metric_maps_to_name_string(self) -> None: + config = TreemapChartConfig( + chart_type="treemap_v2", + groupby=[{"name": "region"}], + metric={"name": "gross_margin", "saved_metric": True}, + ) + assert map_treemap_config(config)["metric"] == "gross_margin" + + +class TestTreemapPluginRegistry: + """Plugin registration and viz-type resolution.""" + + def test_treemap_plugin_registered(self) -> None: + from superset.mcp_service.chart import registry + + plugin = registry.get("treemap_v2") + assert plugin is not None + assert plugin.resolve_viz_type(None) == "treemap_v2" + + def test_display_name_resolves(self) -> None: + from superset.mcp_service.chart.registry import display_name_for_viz_type + + assert display_name_for_viz_type("treemap_v2") == "Treemap" + + def test_pre_validate_missing_fields(self) -> None: + from superset.mcp_service.chart import registry + + plugin = registry.get("treemap_v2") + assert plugin is not None + error = plugin.pre_validate({"chart_type": "treemap_v2"}) + assert error is not None + assert "groupby" in error.message + assert "metric" in error.message + + +class TestTreemapRecommendationCategory: + """Treemap is categorized for chart recommendations and schema discovery.""" + + def test_treemap_in_recommendation_category_map(self) -> None: + from superset.mcp_service.chart.tool.get_chart_data import _VIZ_CATEGORY + + assert _VIZ_CATEGORY.get("treemap_v2") == "treemap" + + def test_get_chart_type_schema_includes_treemap(self) -> None: + from superset.mcp_service.chart.tool.get_chart_type_schema import ( + _CHART_TYPE_ADAPTERS, + ) + + assert "treemap_v2" in _CHART_TYPE_ADAPTERS