This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new c14ac29f477 fix(mcp): complete Treemap round-trip and previews (#44152)
c14ac29f477 is described below
commit c14ac29f4778ba8e27b3cdb5832c37dfba8d4bdb
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Wed Sep 23 16:25:30 2026 +1000
fix(mcp): complete Treemap round-trip and previews (#44152)
---
docs/admin_docs/configuration/mcp-server.mdx | 61 +-
superset/mcp_service/chart/chart_helpers.py | 35 +
superset/mcp_service/chart/chart_utils.py | 124 ++
superset/mcp_service/chart/compile.py | 21 +-
superset/mcp_service/chart/plugins/treemap.py | 10 +-
superset/mcp_service/chart/preview_utils.py | 14 +-
superset/mcp_service/chart/query_result.py | 113 +-
superset/mcp_service/chart/schemas.py | 357 +++--
superset/mcp_service/chart/tool/get_chart_data.py | 9 +
.../mcp_service/chart/tool/get_chart_preview.py | 30 +-
.../chart/tool/get_chart_type_schema.py | 21 +-
superset/mcp_service/chart/tool/update_chart.py | 30 +-
.../mcp_service/chart/tool/update_chart_preview.py | 51 +-
superset/mcp_service/chart/treemap_preview.py | 297 +++++
.../explore/tool/generate_explore_link.py | 3 +-
.../mcp_service/chart/test_treemap_completeness.py | 1373 ++++++++++++++++++++
.../mcp_service/chart/tool/test_get_chart_data.py | 155 +++
.../chart/tool/test_get_chart_type_schema.py | 26 +
.../explore/tool/test_generate_explore_link.py | 49 +
19 files changed, 2598 insertions(+), 181 deletions(-)
diff --git a/docs/admin_docs/configuration/mcp-server.mdx
b/docs/admin_docs/configuration/mcp-server.mdx
index 319db21ccc4..086dce6b9a4 100644
--- a/docs/admin_docs/configuration/mcp-server.mdx
+++ b/docs/admin_docs/configuration/mcp-server.mdx
@@ -1048,7 +1048,7 @@ Specifying a tool name that does not exist logs a warning
at startup and is othe
## Disabling chart type plugins
-The `generate_chart` tool dispatches per chart type (`xy`, `table`, `pie`,
`gauge`, `pivot_table`, `interactive_pivot`, `mixed_timeseries`, `handlebars`,
`big_number`, `histogram`, `box_plot`, `waterfall`) to a registered chart type
plugin. Gauge requests use the public MCP discriminator `chart_type: "gauge"`;
Superset stores the corresponding native Explore visualization as `viz_type:
"gauge_chart"`. Two settings let operators enable or disable individual chart
type plugins at runtime, [...]
+The `generate_chart` tool dispatches per chart type (`xy`, `table`, `pie`,
`gauge`, `treemap_v2`, `pivot_table`, `interactive_pivot`, `mixed_timeseries`,
`handlebars`, `big_number`, `histogram`, `box_plot`, `waterfall`) to a
registered chart type plugin. Gauge requests use the public MCP discriminator
`chart_type: "gauge"`; Superset stores the corresponding native Explore
visualization as `viz_type: "gauge_chart"`. Two settings let operators enable
or disable individual chart type plugin [...]
Gauge data inspection and JSON/CSV/XLSX exports preserve source groups,
including NULL and non-finite metric values, using the same row-count and
completeness semantics as other chart types. The default XLSX exporter writes
non-finite floats as `nan`, `inf`, or `-inf` text, matching CSV and
distinguishing them from blank NULL cells. Only previews and compile checks
skip NULL, NaN, non-finite, and other nonnumeric dial values while preserving
finite groups. In preview/compile paths, a non [...]
@@ -1212,6 +1212,65 @@ The tool does not rewrite dots into SQL path separators:
quoting and nested-fiel
support depend on the dataset's database dialect, and dots can also be literal
characters in column names.
+
+### Treemap charts
+
+Saved and cached previews honor the configured Treemap row limit across ASCII,
+table, and Vega representations. SQL Decimal metrics are accepted when finite;
+booleans and nonfinite values are rejected. Categories with identical display
+names but distinct raw types retain separate geometry groups. Currency settings
+use the native `symbolPosition` property. Adding filters while omitting the
+temporal column preserves saved predicates and the existing neutral time
binding;
+explicit empty filters or a cleared temporal column remove those controls.
+
+
+Use `chart_type: "treemap_v2"` with an ordered `groupby` hierarchy and one
+`metric`. Hierarchy columns accept typed references (`{"name": "region"}`)
+or native strings. Metrics accept typed column references with an aggregate,
+saved metric names, or native SIMPLE/SQL metric objects. SQL metrics still
+require a label and pass the shared SQL-expression validation. Unknown controls
+and executable JavaScript hooks are not accepted.
+
+```json
+{
+ "dataset_id": 7,
+ "config": {
+ "chart_type": "treemap_v2",
+ "groupby": ["region", "product"],
+ "metric": {"name": "revenue", "aggregate": "SUM"},
+ "row_limit": 100,
+ "sort_by_metric": true,
+ "show_labels": true,
+ "show_upper_labels": true,
+ "label_type": "key_value",
+ "number_format": ",.2f"
+ }
+}
+```
+
+Bounded queries sort by the metric descending when requested, then by hierarchy
+columns ascending to break ties, matching the frontend query contract.
+`update_chart` and `update_chart_preview` preserve omitted same-dataset Treemap
+controls, including hierarchy, metrics, formatting, filters, temporal state and
+`template_params`. Partial configurations are accepted only for updates; a new
+chart still requires hierarchy and metric roles. Explicit empty filters clear
+filters, and explicit null clears nullable controls. Rebinding to another
dataset
+requires complete roles and inherits only supported presentation controls.
+
+ASCII and table previews show hierarchy and values. Vega-Lite previews use
+bounded slice-and-dice rectangles with metric-proportional areas, hierarchy
+paths, categorical colors, labels and tooltips. This is not the native ECharts
+layout: use the Explore URL for native formatting and interactions. The preview
+supports up to 1,000 nonnegative rows, 20 hierarchy levels, and the built-in
+`supersetColors` and `lyftColors` schemes. Unsupported color/currency formats,
+negative values, zero totals and over-limit geometry return
+`UnsupportedTreemapPreview` rather than a substitute scatter or bar chart.
+`SMART_NUMBER` is approximated with SI formatting in this preview.
+
+Queries must return every hierarchy output and the metric's exact output label.
+Malformed/error envelopes and nonnumeric or nonfinite metric outputs return
+structured errors instead of apparently successful previews, charts or exports.
+
## Dataset discovery and routing
`list_datasets` searches table name, **description**, schema, and SQL using
diff --git a/superset/mcp_service/chart/chart_helpers.py
b/superset/mcp_service/chart/chart_helpers.py
index 3233dc12c28..f191697e72c 100644
--- a/superset/mcp_service/chart/chart_helpers.py
+++ b/superset/mcp_service/chart/chart_helpers.py
@@ -534,6 +534,14 @@ def resolve_metrics_and_groupby(
viz_type = (
form_data.get("viz_type", getattr(chart, "viz_type", "") if chart else
"") or ""
)
+ if viz_type == "treemap_v2":
+ # Treemap has exactly these roles; stale controls from another plugin
+ # must not override its singular metric or ordered hierarchy.
+ metric = form_data.get("metric")
+ hierarchy = form_data.get("groupby") or []
+ return ([metric] if metric else []), (
+ [hierarchy] if isinstance(hierarchy, str) else list(hierarchy)
+ )
singular_metric_no_groupby = (
"big_number",
"big_number_total",
@@ -694,6 +702,31 @@ def resolve_sort_metric(form_data: dict[str, Any]) -> Any
| None:
return raw or None
+def _apply_treemap_query_fields(
+ qd: dict[str, Any],
+ form_data: dict[str, Any],
+ columns: list[Any],
+ effective_row_limit: int | None,
+) -> None:
+ """Apply Treemap temporal binding and bounded hierarchy ordering."""
+ # extractExtras maps the selected SQL time column to QueryObject
granularity.
+ # A normalized dashboard override takes precedence, including a clear.
+ granularity = form_data.get("granularity",
form_data.get("granularity_sqla"))
+ if granularity is not None:
+ qd["granularity"] = granularity
+ # Match Treemap buildQuery/applyOrderBy, including hierarchy tie-breakers.
+ ordering = qd.pop("orderby", [])
+ ordering.extend(
+ (column, True) for column in columns if isinstance(column, str) and
column
+ )
+ try:
+ bounded = float(effective_row_limit or 0) != 0
+ except (ValueError, TypeError):
+ bounded = True
+ if bounded and ordering:
+ qd["orderby"] = ordering
+
+
def _build_single_query_dict(
form_data: dict[str, Any],
columns: list[Any],
@@ -728,6 +761,8 @@ def _build_single_query_dict(
order_desc if order_desc is not None else
form_data.get("order_desc", True)
)
qd["orderby"] = [(sort_metric, not descending)]
+ if form_data.get("viz_type") == "treemap_v2":
+ _apply_treemap_query_fields(qd, form_data, columns,
effective_row_limit)
apply_form_data_filters_to_query(qd, form_data)
return qd
diff --git a/superset/mcp_service/chart/chart_utils.py
b/superset/mcp_service/chart/chart_utils.py
index 1551693d958..adf06de1acc 100644
--- a/superset/mcp_service/chart/chart_utils.py
+++ b/superset/mcp_service/chart/chart_utils.py
@@ -55,6 +55,7 @@ from superset.mcp_service.chart.schemas import (
SortByConfig,
TableChartConfig,
TreemapChartConfig,
+ TreemapChartUpdateConfig,
WaterfallChartConfig,
XYChartConfig,
)
@@ -857,6 +858,115 @@ def _without_generated_gauge_time_filter(
]
+def resolve_treemap_update_config(
+ config: ChartConfig | TreemapChartUpdateConfig,
+ existing: dict[str, Any],
+ *,
+ dataset_rebind: bool = False,
+) -> ChartConfig:
+ """Fill omitted required roles only from an authorized same-dataset
Treemap."""
+ if not isinstance(config, TreemapChartUpdateConfig) or isinstance(
+ config, TreemapChartConfig
+ ):
+ return config
+ values = config.model_dump(exclude_unset=True)
+ if existing.get("viz_type") == "treemap_v2" and not dataset_rebind:
+ for field in ("groupby", "metric"):
+ if field not in config.model_fields_set and field in existing:
+ values[field] = existing[field]
+ resolved = TreemapChartConfig.model_validate(values)
+ resolved.__pydantic_fields_set__ = set(config.model_fields_set)
+ return resolved
+
+
+_TREEMAP_PRESENTATION_KEYS = frozenset(
+ {
+ "color_scheme",
+ "show_labels",
+ "show_upper_labels",
+ "label_type",
+ "label_position",
+ "number_format",
+ "date_format",
+ "currency_format",
+ }
+)
+
+
+def _merge_treemap_filters(
+ existing: dict[str, Any],
+ patch: dict[str, Any],
+ config: TreemapChartConfig,
+ dataset_rebind: bool,
+) -> None:
+ """Separate explicit filter/temporal changes from mapper-generated
defaults."""
+ fields = config.model_fields_set
+ if "temporal_column" in fields and config.temporal_column is None:
+ patch["adhoc_filters"] = _without_generated_gauge_time_filter(patch)
+ patch.pop(MCP_DASHBOARD_TIME_FILTER_SUBJECT, None)
+ if dataset_rebind:
+ return
+ if "temporal_column" not in fields:
+ # Discard the mapper's default binding before removing its provenance.
+ patch["adhoc_filters"] = _without_generated_gauge_time_filter(patch)
+ patch.pop(MCP_DASHBOARD_TIME_FILTER_SUBJECT, None)
+ if "filters" in fields and config.filters:
+ if subject := existing.get(MCP_DASHBOARD_TIME_FILTER_SUBJECT):
+ patch[MCP_DASHBOARD_TIME_FILTER_SUBJECT] = subject
+ preserve_previous_adhoc_filters(patch, existing)
+ if "filters" not in fields:
+ if "temporal_column" in fields:
+ inherited = (
+ [] if dataset_rebind else
_without_generated_gauge_time_filter(existing)
+ )
+ patch["adhoc_filters"] = [*inherited, *patch.get("adhoc_filters",
[])]
+ else:
+ patch.pop("adhoc_filters", None)
+
+
+def _merge_treemap_form_data(
+ existing: dict[str, Any],
+ generated: dict[str, Any],
+ config: TreemapChartConfig,
+ dataset_rebind: bool,
+) -> dict[str, Any]:
+ """Apply only explicit controls; never inherit query roles across
datasets."""
+ fields = config.model_fields_set
+ merged = {
+ key: value
+ for key, value in existing.items()
+ if not dataset_rebind or key in _TREEMAP_PRESENTATION_KEYS
+ }
+ patch = dict(generated)
+ for field in type(config).model_fields:
+ if field not in fields and (
+ not dataset_rebind or field in _TREEMAP_PRESENTATION_KEYS
+ ):
+ patch.pop(field, None)
+ _merge_treemap_filters(existing, patch, config, dataset_rebind)
+ merged.update(patch)
+ for field in fields:
+ if getattr(config, field) is None:
+ merged.pop(field, None)
+ if "filters" in fields and not config.filters:
+ merged.pop("adhoc_filters", None)
+ merged.pop(MCP_DASHBOARD_TIME_FILTER_SUBJECT, None)
+ if "temporal_column" in fields and config.temporal_column is None:
+ merged.pop(MCP_DASHBOARD_TIME_FILTER_SUBJECT, None)
+ # These roles belong to other plugins and must not affect Treemap queries.
+ for key in (
+ "metrics",
+ "columns",
+ "all_columns",
+ "x_axis",
+ "groupby_b",
+ "metrics_b",
+ "order_by_cols",
+ ):
+ merged.pop(key, None)
+ return merged
+
+
def merge_chart_form_data( # noqa: C901
existing_form_data: dict[str, Any],
new_form_data: dict[str, Any],
@@ -872,6 +982,10 @@ def merge_chart_form_data( # noqa: C901
"""
if existing_form_data.get("viz_type") != new_form_data.get("viz_type"):
return dict(new_form_data)
+ if isinstance(config, TreemapChartConfig):
+ return _merge_treemap_form_data(
+ existing_form_data, new_form_data, config, dataset_rebind
+ )
if isinstance(config, GanttChartConfig):
merged = dict(new_form_data)
if not dataset_rebind:
@@ -1505,6 +1619,16 @@ def map_treemap_config(config: TreemapChartConfig) ->
Dict[str, Any]:
"row_limit": config.row_limit,
"color_scheme": config.color_scheme or "supersetColors",
}
+ for key in _TREEMAP_PRESENTATION_KEYS | {
+ "time_range",
+ "granularity_sqla",
+ "template_params",
+ }:
+ value = getattr(config, key)
+ if value is not None:
+ form_data[key] = (
+ value.to_form_data() if isinstance(value, CurrencyFormat) else
value
+ )
_add_adhoc_filters(form_data, config.filters)
return form_data
diff --git a/superset/mcp_service/chart/compile.py
b/superset/mcp_service/chart/compile.py
index f4f3e80728e..4befcb8f5be 100644
--- a/superset/mcp_service/chart/compile.py
+++ b/superset/mcp_service/chart/compile.py
@@ -42,7 +42,7 @@ from sqlalchemy.exc import SQLAlchemyError
from superset.commands.exceptions import CommandException
from superset.errors import SupersetErrorType
from superset.mcp_service.chart.query_result import (
- normalize_gauge_query_result,
+ normalize_chart_query_result,
query_result_failure,
)
from superset.mcp_service.chart.schemas import ChartError
@@ -136,7 +136,7 @@ def _compile_chart(
query_context = build_query_context_from_form_data(
query_form_data,
row_limit=min(10, int(form_data.get("row_limit") or 10))
- if form_data.get("viz_type") == "gauge_chart"
+ if form_data.get("viz_type") in ("gauge_chart", "treemap_v2")
else 2,
force=False,
)
@@ -156,22 +156,31 @@ def _compile_chart(
tier="compile",
error_obj=_build_compile_error(error_str),
)
- result = normalize_gauge_query_result(result, form_data)
+ result = normalize_chart_query_result(result, form_data)
if isinstance(result, ChartError):
+ is_treemap = form_data.get("viz_type") == "treemap_v2"
+ error_code = (
+ "INVALID_TREEMAP_RESULT" if is_treemap else
"INVALID_GAUGE_RESULT"
+ )
+ message = (
+ "Treemap metric query returned invalid values"
+ if is_treemap
+ else "Gauge metric query returned invalid values"
+ )
return CompileResult(
success=False,
error=result.error,
- error_code="INVALID_GAUGE_RESULT",
+ error_code=error_code,
tier="compile",
error_obj=ChartGenerationError(
error_type=result.error_type,
- message="Gauge metric query returned invalid values",
+ message=message,
details=result.error,
suggestions=[
"Use a numeric-producing metric",
"Check the metric alias and SQL expression",
],
- error_code="INVALID_GAUGE_RESULT",
+ error_code=error_code,
),
)
for query in result.get("queries", []):
diff --git a/superset/mcp_service/chart/plugins/treemap.py
b/superset/mcp_service/chart/plugins/treemap.py
index 170c7d5d476..e907495b572 100644
--- a/superset/mcp_service/chart/plugins/treemap.py
+++ b/superset/mcp_service/chart/plugins/treemap.py
@@ -81,6 +81,8 @@ class TreemapChartPlugin(BaseChartPlugin):
if not isinstance(config, TreemapChartConfig):
return []
refs: list[ColumnRef] = [*config.groupby, config.metric]
+ if config.granularity_sqla:
+ refs.append(ColumnRef(name=config.granularity_sqla))
if config.filters:
for f in config.filters:
refs.append(ColumnRef(name=f.column))
@@ -122,8 +124,14 @@ class TreemapChartPlugin(BaseChartPlugin):
config_dict["metric"]["name"], dataset_context
)
)
+ if granularity := config_dict.get("granularity_sqla"):
+ config_dict["granularity_sqla"] = (
+ DatasetValidator.get_canonical_column_name(granularity,
dataset_context)
+ )
DatasetValidator.normalize_filters(config_dict, dataset_context)
- return TreemapChartConfig.model_validate(config_dict)
+ normalized = TreemapChartConfig.model_validate(config_dict)
+ normalized.__pydantic_fields_set__ = set(config.model_fields_set)
+ return normalized
def schema_error_hint(self) -> ChartGenerationError | None:
return ChartGenerationError(
diff --git a/superset/mcp_service/chart/preview_utils.py
b/superset/mcp_service/chart/preview_utils.py
index 1c8c79ec5e7..f5217f09b8d 100644
--- a/superset/mcp_service/chart/preview_utils.py
+++ b/superset/mcp_service/chart/preview_utils.py
@@ -30,6 +30,7 @@ from typing import Any, Dict, List
from superset.mcp_service.chart.query_result import (
metric_result_label,
+ normalize_chart_query_result,
normalize_gauge_query_result,
query_result_failure,
)
@@ -39,6 +40,7 @@ from superset.mcp_service.chart.schemas import (
TablePreview,
VegaLitePreview,
)
+from superset.mcp_service.chart.treemap_preview import treemap_ascii,
treemap_vega_lite
logger = logging.getLogger(__name__)
@@ -105,8 +107,7 @@ def generate_preview_from_form_data(
if query_failure := query_result_failure(result):
return query_failure
-
- result = normalize_gauge_query_result(result, form_data)
+ result = normalize_chart_query_result(result, form_data)
if isinstance(result, ChartError):
return result
if not result or not result.get("queries"):
@@ -144,7 +145,12 @@ def _generate_ascii_preview_from_data(
viz_type = form_data.get("viz_type", "table")
# Handle different chart types
- if viz_type == "gauge_chart":
+ if viz_type == "treemap_v2":
+ content_or_error = treemap_ascii(data, form_data)
+ if isinstance(content_or_error, ChartError):
+ return content_or_error
+ content = content_or_error
+ elif viz_type == "gauge_chart":
content_or_error = generate_gauge_ascii_preview(data, form_data)
if isinstance(content_or_error, ChartError):
return content_or_error
@@ -1256,6 +1262,8 @@ def _generate_vega_lite_preview_from_data( # noqa: C901
) -> VegaLitePreview | ChartError:
"""Generate Vega-Lite preview from raw data and form_data."""
viz_type = form_data.get("viz_type", "table")
+ if viz_type == "treemap_v2":
+ return treemap_vega_lite(data, form_data)
if viz_type == "gantt_chart":
return _generate_gantt_vega_lite_preview(data, form_data)
if viz_type == "gauge_chart":
diff --git a/superset/mcp_service/chart/query_result.py
b/superset/mcp_service/chart/query_result.py
index 16c01122d16..329dee7ec4d 100644
--- a/superset/mcp_service/chart/query_result.py
+++ b/superset/mcp_service/chart/query_result.py
@@ -19,7 +19,9 @@
import math
from collections.abc import Mapping
-from typing import Any
+from decimal import Decimal
+from numbers import Real
+from typing import Any, cast
from superset.mcp_service.chart.schemas import ChartError
@@ -229,3 +231,112 @@ def validate_gauge_query_result(
"""Check Gauge results using the same finite-dial contract as rendering."""
normalized = normalize_gauge_query_result(result, form_data)
return normalized if isinstance(normalized, ChartError) else None
+
+
+def column_result_label(column: Any) -> str | None:
+ """Resolve the query-result key using frontend ``getColumnLabel`` rules.
+
+ Explore's ``DndColumnSelect`` stores adhoc Custom SQL entries as objects,
+ so ``groupby`` on a saved chart may hold either a physical column name or
+ an adhoc column whose output key is its label (or raw SQL expression).
+ """
+ if isinstance(column, str):
+ return column or None
+ if not isinstance(column, Mapping):
+ return None
+ for key in ("label", "sqlExpression"):
+ value = column.get(key)
+ if isinstance(value, str) and value:
+ return value
+ return None
+
+
+def treemap_hierarchy_labels(form_data: Mapping[str, Any]) -> list[str] | None:
+ """Map a Treemap ``groupby`` onto the result keys its rows are keyed by.
+
+ Returns ``None`` when the hierarchy is absent, empty, contains an
+ unresolvable entry, or would collapse onto duplicate output labels.
+ """
+ hierarchy = form_data.get("groupby")
+ if not isinstance(hierarchy, list) or not hierarchy:
+ return None
+ labels = [column_result_label(column) for column in hierarchy]
+ if any(label is None for label in labels):
+ return None
+ resolved = cast(list[str], labels)
+ if len(set(resolved)) != len(resolved):
+ return None
+ return resolved
+
+
+def normalize_chart_query_result(result: Any, form_data: Mapping[str, Any]) ->
Any:
+ """Validate chart-specific result contracts before consumers use rows."""
+ if form_data.get("viz_type") != "treemap_v2":
+ return normalize_gauge_query_result(result, form_data)
+ if failure := query_result_failure(result):
+ return failure
+ label = metric_result_label(form_data.get("metric"))
+ hierarchy = treemap_hierarchy_labels(form_data)
+ if not label or hierarchy is None or label in hierarchy:
+ return ChartError(
+ error=(
+ "Treemap requires unique hierarchy columns and a distinct
metric label."
+ ),
+ error_type="InvalidTreemapFormData",
+ )
+ queries = result.get("queries") if isinstance(result, Mapping) else None
+ if not isinstance(queries, list) or len(queries) != 1:
+ return ChartError(
+ error="Treemap requires exactly one query result.",
+ error_type="InvalidTreemapResult",
+ )
+ query = queries[0]
+ rows = query.get("data") if isinstance(query, Mapping) else None
+ if not isinstance(rows, list):
+ return ChartError(
+ error="Treemap query data must be an array of rows.",
+ error_type="InvalidTreemapResult",
+ )
+ if failure := _validate_treemap_rows(rows, hierarchy, label):
+ return failure
+ return result
+
+
+def _validate_treemap_rows(
+ rows: list[Any], hierarchy: list[str], label: str
+) -> ChartError | None:
+ """Require complete hierarchy outputs and finite numeric metric values."""
+ for index, row in enumerate(rows):
+ if not isinstance(row, Mapping) or any(
+ column not in row for column in [*hierarchy, label]
+ ):
+ return ChartError(
+ error=f"Treemap row {index} is missing hierarchy or metric
outputs.",
+ error_type="InvalidTreemapResult",
+ )
+ value = row[label]
+ try:
+ valid = (
+ not isinstance(value, bool)
+ and isinstance(value, (Real, Decimal))
+ and (
+ value.is_finite()
+ if isinstance(value, Decimal)
+ else math.isfinite(value)
+ )
+ )
+ except (OverflowError, ValueError):
+ valid = False
+ if not valid:
+ return ChartError(
+ error=(
+ f"Treemap row {index} metric {label!r} must be finite and
numeric."
+ ),
+ error_type="InvalidTreemapMetric",
+ )
+ if any(isinstance(row[column], (dict, list)) for column in hierarchy):
+ return ChartError(
+ error=f"Treemap row {index} hierarchy values must be scalar.",
+ error_type="InvalidTreemapResult",
+ )
+ return None
diff --git a/superset/mcp_service/chart/schemas.py
b/superset/mcp_service/chart/schemas.py
index 077bb756c9d..5cfd4820898 100644
--- a/superset/mcp_service/chart/schemas.py
+++ b/superset/mcp_service/chart/schemas.py
@@ -1147,6 +1147,137 @@ class PieChartConfig(BaseChartConfig):
return self
+def _adapt_native_single_metric_form_data(data: Any) -> Any: # noqa: C901
+ """Adapt bounded native single-metric controls before strict validation."""
+ if not isinstance(data, dict):
+ return data
+ data = dict(data)
+
+ # ``gauge`` is the public MCP discriminator; ``gauge_chart`` remains
+ # the native frontend viz_type and is accepted only as an input alias.
+ if data.get("chart_type") == "gauge_chart" or (
+ "chart_type" not in data and data.get("viz_type") == "gauge_chart"
+ ):
+ data["chart_type"] = "gauge"
+ data.pop("viz_type", None)
+
+ # These identify the Explore/chart envelope, not visualization controls.
+ for key in (
+ "datasource",
+ "datasource_id",
+ "datasource_name",
+ "datasource_type",
+ "form_data_key",
+ "slice_id",
+ "slice_name",
+ "url",
+ ):
+ data.pop(key, None)
+ data.pop("_mcp_dashboard_time_filter_subject", None)
+
+ metric = data.get("metric")
+ if isinstance(metric, str):
+ data["metric"] = {"name": metric, "saved_metric": True}
+ elif isinstance(metric, dict) and metric.get("expressionType") in {
+ "SIMPLE",
+ "SQL",
+ }:
+ expression_type = metric.get("expressionType")
+ if expression_type == "SQL":
+ data["metric"] = {
+ "sql_expression": metric.get("sqlExpression"),
+ "label": metric.get("label"),
+ }
+ else:
+ if not isinstance(metric.get("aggregate"), str):
+ raise ValueError("Native SIMPLE metrics require an aggregate")
+ column = metric.get("column")
+ column_name = (
+ column.get("column_name") or column.get("columnName")
+ if isinstance(column, dict)
+ else None
+ )
+ data["metric"] = {
+ "name": column_name,
+ "aggregate": metric.get("aggregate"),
+ "label": metric.get("label"),
+ }
+
+ groupby = data.get("groupby")
+ if isinstance(groupby, str):
+ groupby = [groupby]
+ if isinstance(groupby, list):
+ data["groupby"] = [
+ {"name": value} if isinstance(value, str) else value for value in
groupby
+ ]
+
+ if isinstance(data.get("time_range"), str):
+ data["time_range"] = validate_time_range(data["time_range"]) or None
+
+ # Supported native SIMPLE filters are represented by FilterConfig.
+ # SQL adhoc filters remain intentionally unsupported on the typed MCP
+ # surface. TEMPORAL_RANGE is represented by time_range/granularity.
+ if "adhoc_filters" in data:
+ if "filters" in data:
+ raise ValueError("Use either filters or adhoc_filters, not both")
+ native_filters = data.pop("adhoc_filters")
+ if not isinstance(native_filters, list):
+ raise ValueError("adhoc_filters must be a list")
+ filters: list[dict[str, Any]] = []
+ for index, filter_ in enumerate(native_filters):
+ if not isinstance(filter_, dict):
+ raise ValueError(f"adhoc_filters[{index}] must be an object")
+ if filter_.get("expressionType") not in (None, "SIMPLE"):
+ raise ValueError(
+ f"adhoc_filters[{index}] must use expressionType='SIMPLE'"
+ )
+ if str(filter_.get("clause", "WHERE")).upper() != "WHERE":
+ raise ValueError(f"adhoc_filters[{index}] must use
clause='WHERE'")
+ operator = filter_.get("operator") or filter_.get("op")
+ subject = filter_.get("subject") or filter_.get("col")
+ comparator = filter_.get("comparator", filter_.get("val"))
+ if operator == "TEMPORAL_RANGE":
+ if not isinstance(subject, str) or not subject:
+ raise ValueError(f"adhoc_filters[{index}] has no temporal
subject")
+ if not isinstance(comparator, str):
+ raise ValueError(
+ f"adhoc_filters[{index}] requires a temporal
comparator"
+ )
+ comparator = validate_time_range(comparator) or NO_TIME_RANGE
+ if comparator == NO_TIME_RANGE:
+ if data.get("temporal_column") not in (None, subject):
+ raise ValueError(
+ f"adhoc_filters[{index}] conflicts with another "
+ "dashboard temporal binding"
+ )
+ data["temporal_column"] = subject
+ continue
+ if (
+ data.get("granularity_sqla") not in (None, subject)
+ and data.get("time_range") != NO_TIME_RANGE
+ ) or data.get("time_range") not in (
+ None,
+ NO_TIME_RANGE,
+ comparator,
+ ):
+ raise ValueError(
+ f"adhoc_filters[{index}] conflicts with another
temporal "
+ "range; multiple distinct temporal ranges are not
supported"
+ )
+ data["granularity_sqla"] = subject
+ data["time_range"] = comparator
+ continue
+ if operator == "==":
+ operator = "="
+ if operator not in
get_args(FilterConfig.model_fields["op"].annotation):
+ raise ValueError(
+ f"adhoc_filters[{index}] uses unsupported operator
{operator!r}"
+ )
+ filters.append({"column": subject, "op": operator, "value":
comparator})
+ data["filters"] = filters
+ return data
+
+
class GaugeChartConfig(BaseChartConfig):
"""Config for gauge charts (viz_type ``gauge_chart``).
@@ -1245,136 +1376,9 @@ class GaugeChartConfig(BaseChartConfig):
@model_validator(mode="before")
@classmethod
- def adapt_native_form_data(cls, data: Any) -> Any: # noqa: C901
- """Accept the Gauge plugin's native form_data without weakening
typing."""
- if not isinstance(data, dict):
- return data
- data = dict(data)
-
- # ``gauge`` is the public MCP discriminator; ``gauge_chart`` remains
- # the native frontend viz_type and is accepted only as an input alias.
- if data.get("chart_type") == "gauge_chart" or (
- "chart_type" not in data and data.get("viz_type") == "gauge_chart"
- ):
- data["chart_type"] = "gauge"
- data.pop("viz_type", None)
-
- # These identify the Explore/chart envelope, not Gauge controls.
- for key in (
- "datasource",
- "datasource_id",
- "datasource_name",
- "datasource_type",
- "form_data_key",
- "slice_id",
- "slice_name",
- "url",
- ):
- data.pop(key, None)
- data.pop("_mcp_dashboard_time_filter_subject", None)
-
- metric = data.get("metric")
- if isinstance(metric, str):
- data["metric"] = {"name": metric, "saved_metric": True}
- elif isinstance(metric, dict) and metric.get("expressionType") in {
- "SIMPLE",
- "SQL",
- }:
- expression_type = metric.get("expressionType")
- if expression_type == "SQL":
- data["metric"] = {
- "sql_expression": metric.get("sqlExpression"),
- "label": metric.get("label"),
- }
- else:
- column = metric.get("column")
- column_name = (
- column.get("column_name") or column.get("columnName")
- if isinstance(column, dict)
- else None
- )
- data["metric"] = {
- "name": column_name,
- "aggregate": metric.get("aggregate"),
- "label": metric.get("label"),
- }
-
- groupby = data.get("groupby")
- if isinstance(groupby, str):
- groupby = [groupby]
- if isinstance(groupby, list):
- data["groupby"] = [
- {"name": value} if isinstance(value, str) else value
- for value in groupby
- ]
-
- if isinstance(data.get("time_range"), str):
- data["time_range"] = validate_time_range(data["time_range"]) or
None
-
- # Supported native SIMPLE filters are represented by FilterConfig.
- # SQL adhoc filters remain intentionally unsupported on the typed MCP
- # surface. TEMPORAL_RANGE is represented by time_range/granularity.
- if "adhoc_filters" in data:
- if "filters" in data:
- raise ValueError("Use either filters or adhoc_filters, not
both")
- native_filters = data.pop("adhoc_filters")
- if not isinstance(native_filters, list):
- raise ValueError("adhoc_filters must be a list")
- filters: list[dict[str, Any]] = []
- for index, filter_ in enumerate(native_filters):
- if not isinstance(filter_, dict):
- raise ValueError(f"adhoc_filters[{index}] must be an
object")
- if filter_.get("expressionType") not in (None, "SIMPLE"):
- raise ValueError(
- f"adhoc_filters[{index}] must use
expressionType='SIMPLE'"
- )
- if str(filter_.get("clause", "WHERE")).upper() != "WHERE":
- raise ValueError(f"adhoc_filters[{index}] must use
clause='WHERE'")
- operator = filter_.get("operator") or filter_.get("op")
- subject = filter_.get("subject") or filter_.get("col")
- comparator = filter_.get("comparator", filter_.get("val"))
- if operator == "TEMPORAL_RANGE":
- if not isinstance(subject, str) or not subject:
- raise ValueError(
- f"adhoc_filters[{index}] has no temporal subject"
- )
- if not isinstance(comparator, str):
- raise ValueError(
- f"adhoc_filters[{index}] requires a temporal
comparator"
- )
- comparator = validate_time_range(comparator) or
NO_TIME_RANGE
- if comparator == NO_TIME_RANGE:
- if data.get("temporal_column") not in (None, subject):
- raise ValueError(
- f"adhoc_filters[{index}] conflicts with
another "
- "dashboard temporal binding"
- )
- data["temporal_column"] = subject
- continue
- if (
- data.get("granularity_sqla") not in (None, subject)
- and data.get("time_range") != NO_TIME_RANGE
- ) or data.get("time_range") not in (
- None,
- NO_TIME_RANGE,
- comparator,
- ):
- raise ValueError(
- f"adhoc_filters[{index}] conflicts with another
temporal "
- "range; multiple distinct temporal ranges are not
supported"
- )
- data["granularity_sqla"] = subject
- data["time_range"] = comparator
- continue
- if operator == "==":
- operator = "="
- if operator not in
get_args(FilterConfig.model_fields["op"].annotation):
- raise ValueError(
- f"adhoc_filters[{index}] uses unsupported operator
{operator!r}"
- )
- filters.append({"column": subject, "op": operator, "value":
comparator})
- data["filters"] = filters
- return data
+ def adapt_native_form_data(cls, data: Any) -> Any:
+ """Accept bounded native form data without weakening typed
validation."""
+ return _adapt_native_single_metric_form_data(data)
@field_validator("time_range")
@classmethod
@@ -1458,7 +1462,7 @@ class GaugeChartConfig(BaseChartConfig):
return self
-class TreemapChartConfig(BaseChartConfig):
+class TreemapChartUpdateConfig(BaseChartConfig):
"""Config for treemap charts (viz_type ``treemap_v2``).
Matches the frontend Treemap buildQuery contract: one ``metric`` sizing
@@ -1469,15 +1473,25 @@ class TreemapChartConfig(BaseChartConfig):
model_config = ConfigDict(extra="ignore", populate_by_name=True)
- chart_type: Literal["treemap_v2"] = "treemap_v2"
- groupby: List[ColumnRef] = Field(
+ # Required even though every other field is optional: this partial model
+ # sits beside the discriminated ``ChartConfig`` union in the update
+ # requests. With a defaulted discriminator any config that failed the
+ # discriminated branch — including one that simply omitted
+ # ``chart_type`` — would fall through to this model and silently rewrite
+ # an existing chart of another type into a Treemap.
+ chart_type: Literal["treemap_v2"] = Field(
...,
+ description="Chart type discriminator; must be 'treemap_v2'",
+ )
+ groupby: List[ColumnRef] | None = Field(
+ None,
min_length=1,
+ max_length=20,
description="Ordered category columns forming the treemap hierarchy "
"(first = outermost level; order defines nesting)",
)
- metric: ColumnRef = Field(
- ...,
+ metric: ColumnRef | None = Field(
+ None,
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)",
)
@@ -1500,9 +1514,58 @@ class TreemapChartConfig(BaseChartConfig):
max_length=100,
)
+ show_labels: bool = True
+ show_upper_labels: bool = True
+ label_type: Literal["key", "Key", "value", "key_value"] = "key_value"
+ label_position: Literal[
+ "top",
+ "left",
+ "right",
+ "bottom",
+ "inside",
+ "insideLeft",
+ "insideRight",
+ "insideTop",
+ "insideBottom",
+ "insideTopLeft",
+ "insideBottomLeft",
+ "insideTopRight",
+ "insideBottomRight",
+ ] = "insideTopLeft"
+ number_format: str = Field("SMART_NUMBER", max_length=100)
+ date_format: str = Field("smart_date", max_length=100)
+ currency_format: CurrencyFormat | None = None
+ time_range: str | None = Field(None, max_length=1000)
+ granularity_sqla: str | None = Field(None, min_length=1, max_length=255)
+ template_params: str | None = Field(None, max_length=10000)
+
+ @model_validator(mode="before")
+ @classmethod
+ def adapt_native_form_data(cls, data: Any) -> Any:
+ """Accept native hierarchy and saved, SIMPLE, and SQL metric inputs."""
+ return _adapt_native_single_metric_form_data(data)
+
+ @field_validator("time_range")
+ @classmethod
+ def validate_treemap_time_range(cls, value: str | None) -> str | None:
+ """Validate time ranges using the shared parser."""
+ return validate_time_range(value)
+
@model_validator(mode="after")
- def reject_metric_style_groupby(self) -> "TreemapChartConfig":
+ def reject_metric_style_groupby(self) -> "TreemapChartUpdateConfig":
"""groupby entries are hierarchy dimensions, not metrics."""
+ names = [col.name for col in self.groupby or []]
+ if len(set(names)) != len(names):
+ raise ValueError("groupby must contain unique hierarchy columns")
+ metric_label = (self.metric.label or self.metric.name) if self.metric
else None
+ if (
+ self.metric
+ and metric_label in names
+ and (self.metric.label or self.metric.saved_metric)
+ ):
+ raise ValueError(
+ "metric output label must not collide with hierarchy columns"
+ )
for i, col in enumerate(self.groupby or []):
_reject_sql_expression_on_dimension(col, f"groupby[{i}]")
if col.is_metric:
@@ -1514,6 +1577,22 @@ class TreemapChartConfig(BaseChartConfig):
return self
+class TreemapChartConfig(TreemapChartUpdateConfig):
+ """Complete Treemap configuration required for generation and
compilation."""
+
+ # Restored to a default here: this model is only ever reached through the
+ # discriminated union, which already requires the key in client payloads,
+ # and internal call sites construct it directly.
+ chart_type: Literal["treemap_v2"] = "treemap_v2"
+ groupby: List[ColumnRef] = Field(
+ ...,
+ min_length=1,
+ max_length=20,
+ description="Ordered hierarchy columns, outermost first",
+ )
+ metric: ColumnRef = Field(..., description="Metric sizing the hierarchy
tiles")
+
+
class BubbleChartConfig(BaseChartConfig):
"""Config for bubble charts (viz_type ``bubble_v2``).
@@ -3863,7 +3942,7 @@ class UpdateChartRequest(ChartRequestNormalizerMixin,
QueryCacheControl):
description="Chart ID or UUID",
validation_alias=AliasChoices("identifier", "id", "chart_id"),
)
- config: ChartConfig | None = Field(
+ config: ChartConfig | TreemapChartUpdateConfig | None = Field(
None,
description="Chart configuration. Optional; omit to only update
chart_name.",
)
@@ -3938,7 +4017,9 @@ class
UpdateChartPreviewRequest(ChartRequestNormalizerMixin, FormDataCacheContro
),
)
dataset_id: int | str = Field(..., description="Dataset ID or UUID")
- config: ChartConfig = Field(..., description="Chart configuration")
+ config: ChartConfig | TreemapChartUpdateConfig = Field(
+ ..., description="Chart configuration"
+ )
generate_preview: bool = True
preview_formats: List[Literal["url", "ascii", "vega_lite", "table"]] =
Field(
default_factory=lambda: ["url"],
diff --git a/superset/mcp_service/chart/tool/get_chart_data.py
b/superset/mcp_service/chart/tool/get_chart_data.py
index 5369a436c3c..5ae0237ea34 100644
--- a/superset/mcp_service/chart/tool/get_chart_data.py
+++ b/superset/mcp_service/chart/tool/get_chart_data.py
@@ -45,6 +45,7 @@ from superset.mcp_service.chart.chart_helpers import (
)
from superset.mcp_service.chart.chart_utils import validate_chart_dataset
from superset.mcp_service.chart.query_result import (
+ normalize_chart_query_result,
query_result_failure,
)
from superset.mcp_service.chart.schemas import (
@@ -791,6 +792,10 @@ async def execute_chart_data( # noqa: C901
command.validate()
result = command.run()
+ if form_data.get("viz_type") == "treemap_v2":
+ result = normalize_chart_query_result(result, form_data)
+ if isinstance(result, ChartError):
+ return result
if query_failure := query_result_failure(result):
return query_failure
@@ -1164,6 +1169,10 @@ async def _query_from_form_data( # noqa: C901
command.validate()
result = command.run()
+ if form_data.get("viz_type") == "treemap_v2":
+ result = normalize_chart_query_result(result, form_data)
+ if isinstance(result, ChartError):
+ return result
if query_failure := query_result_failure(result):
return query_failure
diff --git a/superset/mcp_service/chart/tool/get_chart_preview.py
b/superset/mcp_service/chart/tool/get_chart_preview.py
index 31c9b577f9b..0a8e2cd13d1 100644
--- a/superset/mcp_service/chart/tool/get_chart_preview.py
+++ b/superset/mcp_service/chart/tool/get_chart_preview.py
@@ -47,7 +47,7 @@ from superset.mcp_service.chart.preview_utils import (
generate_gauge_vega_lite_preview,
)
from superset.mcp_service.chart.query_result import (
- normalize_gauge_query_result,
+ normalize_chart_query_result,
query_result_failure,
)
from superset.mcp_service.chart.schemas import (
@@ -62,6 +62,7 @@ from superset.mcp_service.chart.schemas import (
URLPreview,
VegaLitePreview,
)
+from superset.mcp_service.chart.treemap_preview import treemap_ascii,
treemap_vega_lite
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
@@ -185,7 +186,14 @@ def _no_query_fields_error(chart: ChartLike) -> ChartError:
def _preview_row_limit(form_data: dict[str, Any], fallback: int) -> int:
- """Keep Gauge preview cardinality aligned with its frontend row limit."""
+ """Keep single-metric previews aligned with their frontend row limits."""
+ if form_data.get("viz_type") == "treemap_v2":
+ value = form_data.get("row_limit", 100)
+ try:
+ limit = int(value)
+ except (TypeError, ValueError, OverflowError):
+ limit = 100
+ return limit if 1 <= limit <= 10000 else 100
if form_data.get("viz_type") != "gauge_chart":
return fallback
value = form_data.get("row_limit", 10)
@@ -284,7 +292,7 @@ class ASCIIPreviewStrategy(PreviewFormatStrategy):
if query_failure := query_result_failure(result):
return query_failure
- result = normalize_gauge_query_result(result, form_data)
+ result = normalize_chart_query_result(result, form_data)
if isinstance(result, ChartError):
return result
@@ -292,12 +300,14 @@ class ASCIIPreviewStrategy(PreviewFormatStrategy):
if result and "queries" in result and len(result["queries"]) > 0:
data = result["queries"][0].get("data") or []
- if form_data.get("viz_type") == "gauge_chart":
+ if form_data.get("viz_type") == "treemap_v2":
+ ascii_chart = treemap_ascii(
+ data, form_data, self.request.ascii_width or 80
+ )
+ elif form_data.get("viz_type") == "gauge_chart":
ascii_chart = generate_gauge_ascii_preview(
data, form_data, self.request.ascii_width or 80
)
- if isinstance(ascii_chart, ChartError):
- return ascii_chart
else:
ascii_chart = generate_ascii_chart(
data,
@@ -306,6 +316,8 @@ class ASCIIPreviewStrategy(PreviewFormatStrategy):
self.request.ascii_height or 20,
)
+ if isinstance(ascii_chart, ChartError):
+ return ascii_chart
return ASCIIPreview(
ascii_content=ascii_chart,
width=self.request.ascii_width or 80,
@@ -364,7 +376,7 @@ class TablePreviewStrategy(PreviewFormatStrategy):
if query_failure := query_result_failure(result):
return query_failure
- result = normalize_gauge_query_result(result, form_data)
+ result = normalize_chart_query_result(result, form_data)
if isinstance(result, ChartError):
return result
@@ -483,7 +495,7 @@ class VegaLitePreviewStrategy(PreviewFormatStrategy):
if query_failure := query_result_failure(result):
return query_failure
- result = normalize_gauge_query_result(result, form_data)
+ result = normalize_chart_query_result(result, form_data)
if isinstance(result, ChartError):
return result
@@ -492,6 +504,8 @@ class VegaLitePreviewStrategy(PreviewFormatStrategy):
if result and "queries" in result and len(result["queries"]) > 0:
chart_data = result["queries"][0].get("data", [])
+ if form_data.get("viz_type") == "treemap_v2":
+ return treemap_vega_lite(chart_data, form_data)
if form_data.get("viz_type") == "gauge_chart":
return generate_gauge_vega_lite_preview(chart_data, form_data)
viz_type = getattr(self.chart, "viz_type", None) or form_data.get(
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 617579c7dc1..3e09d888e1b 100644
--- a/superset/mcp_service/chart/tool/get_chart_type_schema.py
+++ b/superset/mcp_service/chart/tool/get_chart_type_schema.py
@@ -22,6 +22,8 @@ MCP tool: get_chart_type_schema
from __future__ import annotations
import logging
+from copy import deepcopy
+from functools import lru_cache
from typing import Any, Dict
from pydantic import TypeAdapter
@@ -256,10 +258,27 @@ _CHART_EXAMPLES: Dict[str, list[Dict[str, Any]]] = {
"groupby": [{"name": "region"}, {"name": "product"}],
"metric": {"name": "revenue", "aggregate": "SUM"},
},
+ {
+ "chart_type": "treemap_v2",
+ "groupby": ["region", "product"],
+ "metric": "total_revenue",
+ "show_labels": True,
+ "show_upper_labels": True,
+ "label_type": "key_value",
+ "number_format": ",.2f",
+ "sort_by_metric": False,
+ "row_limit": 100,
+ },
],
}
+@lru_cache(maxsize=len(_CHART_TYPE_ADAPTERS))
+def _compiled_chart_schema(chart_type: str) -> dict[str, Any]:
+ """Compile static adapter schemas once; callers must copy before exposing
them."""
+ return _CHART_TYPE_ADAPTERS[chart_type].json_schema()
+
+
def _get_chart_type_schema_impl(
chart_type: str,
include_examples: bool = True,
@@ -314,7 +333,7 @@ def _get_chart_type_schema_impl(
"valid_chart_types": enabled_types,
}
- schema = adapter.json_schema()
+ schema = deepcopy(_compiled_chart_schema(chart_type))
result: ChartTypeSchemaResponse = {
"chart_type": chart_type,
"schema": schema,
diff --git a/superset/mcp_service/chart/tool/update_chart.py
b/superset/mcp_service/chart/tool/update_chart.py
index 43d08f99a0d..7c190c7f5b2 100644
--- a/superset/mcp_service/chart/tool/update_chart.py
+++ b/superset/mcp_service/chart/tool/update_chart.py
@@ -43,6 +43,7 @@ from superset.mcp_service.chart.chart_utils import (
merge_chart_form_data,
merge_interactive_pivot_ui_config,
merge_table_column_config,
+ resolve_treemap_update_config,
validate_gantt_form_data,
)
from superset.mcp_service.chart.compile import validate_and_compile
@@ -55,6 +56,7 @@ from superset.mcp_service.chart.schemas import (
GenerateChartResponse,
PerformanceMetadata,
TableChartConfig,
+ TreemapChartConfig,
UpdateChartRequest,
)
from superset.mcp_service.chart.validation.dataset_validator import (
@@ -375,7 +377,7 @@ def _build_replacement_form_data(
new_form_data.pop("_mcp_warnings", None)
dataset_rebind = replacement_dataset_id is not None
if replacement_dataset_id is not None and not isinstance(
- parsed_config, (GanttChartConfig, GaugeChartConfig)
+ parsed_config, (GanttChartConfig, GaugeChartConfig, TreemapChartConfig)
):
# Drop only the inherited state the replacement dataset cannot
# resolve, then merge as a same-dataset update. Gantt and Gauge keep
the
@@ -817,12 +819,16 @@ async def update_chart( # noqa: C901
request.dataset_id is not None
and request.dataset_id != getattr(chart, "datasource_id", None)
and request.config is None
- and getattr(chart, "viz_type", None) == "gauge_chart"
+ and getattr(chart, "viz_type", None) in ("gauge_chart",
"treemap_v2")
):
return _validation_error_response(
- message="Gauge dataset rebind requires a complete Gauge
config.",
+ message=(
+ "Gauge dataset rebind requires a complete Gauge config."
+ if chart.viz_type == "gauge_chart"
+ else "Treemap dataset rebind requires a complete Treemap
config."
+ ),
details=(
- "Provide chart_type='gauge' and a metric valid on the
target "
+ "Provide the chart type and complete roles valid on the
target "
"dataset. This prevents stale metric, groupby, and filter
roles "
"from the previous dataset from being retained."
),
@@ -859,7 +865,21 @@ async def update_chart( # noqa: C901
new_form_data: dict[str, Any] | None = None
# config is already a typed ChartConfig | None (validated by Pydantic)
- parsed_config = request.config
+ try:
+ parsed_config = (
+ resolve_treemap_update_config(
+ request.config,
+ _get_existing_form_data(chart),
+ dataset_rebind=request.dataset_id is not None
+ and request.dataset_id != chart.datasource_id,
+ )
+ if request.config is not None
+ else None
+ )
+ except ValueError as ex:
+ return _validation_error_response(
+ "Invalid Treemap update configuration", str(ex)
+ )
validation_config = parsed_config
if request.add_columns is not None:
validation_config = TableChartConfig(columns=request.add_columns)
diff --git a/superset/mcp_service/chart/tool/update_chart_preview.py
b/superset/mcp_service/chart/tool/update_chart_preview.py
index 3b50eb95e4b..ff278144717 100644
--- a/superset/mcp_service/chart/tool/update_chart_preview.py
+++ b/superset/mcp_service/chart/tool/update_chart_preview.py
@@ -41,6 +41,7 @@ from superset.mcp_service.chart.chart_utils import (
merge_chart_form_data,
merge_interactive_pivot_ui_config,
merge_table_column_config,
+ resolve_treemap_update_config,
validate_gantt_form_data,
)
from superset.mcp_service.chart.compile import validate_and_compile
@@ -182,6 +183,39 @@ def update_chart_preview( # noqa: C901
NORMALIZATION_EXCEPTIONS,
)
+ warnings: list[str] = []
+ previous_form_data: dict[str, Any] | None = None
+
+ if request.form_data_key:
+ previous_form_data =
_get_previous_form_data(request.form_data_key)
+ if previous_form_data is None:
+ warnings.append(INVALID_FORM_DATA_KEY_WARNING)
+ previous_datasource = str(
+ (previous_form_data or {}).get("datasource")
+ or (previous_form_data or {}).get("datasource_id")
+ or ""
+ ).split("__", 1)[0]
+ dataset_rebind = previous_datasource != str(dataset.id) and (
+ bool(previous_datasource) or config.chart_type == "treemap_v2"
+ )
+ try:
+ config = resolve_treemap_update_config(
+ config,
+ previous_form_data or {},
+ dataset_rebind=dataset_rebind,
+ )
+ except ValueError as ex:
+ return {
+ "chart": None,
+ "error": {
+ "error_type": "ValidationError",
+ "message": "Invalid Treemap update configuration",
+ "details": str(ex),
+ },
+ "success": False,
+ "schema_version": "2.0",
+ "api_version": "v1",
+ }
try:
config = DatasetValidator.normalize_column_names(
config,
@@ -200,25 +234,10 @@ def update_chart_preview( # noqa: C901
config, dataset_id=request.dataset_id
)
new_form_data.pop("_mcp_warnings", None)
- warnings: list[str] = []
- previous_form_data: dict[str, Any] | None = None
-
- if request.form_data_key:
- previous_form_data =
_get_previous_form_data(request.form_data_key)
- if previous_form_data is None:
- warnings.append(INVALID_FORM_DATA_KEY_WARNING)
if previous_form_data:
merge_table_column_config(previous_form_data, new_form_data)
merge_interactive_pivot_ui_config(previous_form_data,
new_form_data)
- previous_datasource = str(
- previous_form_data.get("datasource")
- or previous_form_data.get("datasource_id")
- or ""
- ).split("__", 1)[0]
- dataset_rebind = bool(
- previous_datasource
- ) and previous_datasource != str(dataset.id)
new_form_data = merge_chart_form_data(
previous_form_data,
new_form_data,
@@ -273,7 +292,7 @@ def update_chart_preview( # noqa: C901
config,
new_form_data,
dataset,
- run_compile_check=config.chart_type == "gauge",
+ run_compile_check=config.chart_type in ("gauge", "treemap_v2"),
)
if not compile_result.success:
logger.warning(
diff --git a/superset/mcp_service/chart/treemap_preview.py
b/superset/mcp_service/chart/treemap_preview.py
new file mode 100644
index 00000000000..ebfaf9f0491
--- /dev/null
+++ b/superset/mcp_service/chart/treemap_preview.py
@@ -0,0 +1,297 @@
+# 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.
+
+"""Bounded slice-and-dice Treemap previews using explicit rectangle
geometry."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from superset.mcp_service.chart.query_result import (
+ metric_result_label,
+ normalize_chart_query_result,
+ treemap_hierarchy_labels,
+)
+from superset.mcp_service.chart.schemas import ChartError, VegaLitePreview
+
+# Match the built-in frontend categorical schemes. Unknown schemes are not
guessed.
+_PALETTES = {
+ "supersetColors": [
+ "#1FA8C9",
+ "#454E7C",
+ "#5AC189",
+ "#FF7F44",
+ "#666666",
+ "#E04355",
+ "#FCC700",
+ "#A868B7",
+ "#3CCCCB",
+ "#A38F79",
+ "#8FD3E4",
+ "#A1A6BD",
+ "#ACE1C4",
+ "#FEC0A1",
+ "#B2B2B2",
+ "#EFA1AA",
+ "#FDE380",
+ "#D3B3DA",
+ "#9EE5E5",
+ "#D1C6BC",
+ ],
+ "lyftColors": [
+ "#EA0B8C",
+ "#6C838E",
+ "#29ABE2",
+ "#33D9C1",
+ "#9DACB9",
+ "#7560AA",
+ "#2D5584",
+ "#831C4A",
+ "#333D47",
+ "#AC2077",
+ ],
+}
+_MAX_ROWS = 1000
+
+
+def treemap_ascii(
+ data: list[dict[str, Any]], form_data: dict[str, Any], width: int = 80
+) -> str | ChartError:
+ """Show the ordered hierarchy and values rather than unrelated bar
geometry."""
+ checked = normalize_chart_query_result({"queries": [{"data": data}]},
form_data)
+ if isinstance(checked, ChartError):
+ return checked
+ label = metric_result_label(form_data["metric"])
+ assert label is not None
+ hierarchy = treemap_hierarchy_labels(form_data)
+ assert hierarchy is not None
+ lines = [f"Treemap hierarchy | {label}"]
+ for row in data[:_MAX_ROWS]:
+ path = " > ".join(str(row[column]) for column in hierarchy)
+ lines.append(f"{path} | {row[label]}")
+ if len(data) > _MAX_ROWS:
+ lines.append(f"Showing {_MAX_ROWS} of {len(data)} rows")
+ return "\n".join(line[:width] for line in lines)
+
+
+def treemap_vega_lite( # noqa: C901
+ data: list[dict[str, Any]], form_data: dict[str, Any]
+) -> VegaLitePreview | ChartError:
+ """Render nested metric-proportional rectangles; never substitute
scatter/bar marks.
+
+ Vega-Lite has no hierarchy transform. A bounded slice-and-dice layout is
+ computed here and sent as explicit coordinates, without executable hooks.
+ Native ECharts layout and dashboard interactions remain available in
Explore.
+ """
+ checked = normalize_chart_query_result({"queries": [{"data": data}]},
form_data)
+ if isinstance(checked, ChartError):
+ return checked
+ label = metric_result_label(form_data["metric"])
+ assert label is not None
+ if not data:
+ return ChartError(error="No Treemap data available.",
error_type="NoDataError")
+ scheme = form_data.get("color_scheme") or "supersetColors"
+ if (
+ scheme not in _PALETTES
+ or form_data.get("currency_format")
+ or form_data.get("label_position", "insideTopLeft") != "insideTopLeft"
+ ):
+ return ChartError(
+ error=(
+ "This Treemap color/currency/label format requires the native "
+ "Explore renderer; use url or table preview."
+ ),
+ error_type="UnsupportedTreemapPreview",
+ )
+ if (
+ len(data) > _MAX_ROWS
+ or any(row[label] < 0 for row in data)
+ or not any(row[label] > 0 for row in data)
+ ):
+ return ChartError(
+ error=(
+ "Treemap geometry requires at most 1000 nonnegative rows and "
+ "a positive total; use table or url preview."
+ ),
+ error_type="UnsupportedTreemapPreview",
+ )
+ hierarchy = treemap_hierarchy_labels(form_data)
+ assert hierarchy is not None
+ if len(hierarchy) > 20 or not math.isfinite(sum(float(row[label]) for row
in data)):
+ return ChartError(
+ error="Treemap preview requires at most 20 levels and a finite
total.",
+ error_type="UnsupportedTreemapPreview",
+ )
+ nodes: list[dict[str, Any]] = []
+
+ def layout(
+ rows: list[dict[str, Any]],
+ depth: int,
+ path: list[str],
+ x: float,
+ y: float,
+ width: float,
+ height: float,
+ ) -> None:
+ """Recursively partition the parent rectangle in hierarchy order."""
+ groups: dict[tuple[str, str], list[dict[str, Any]]] = {}
+ for row in rows:
+ value = row[hierarchy[depth]]
+ groups.setdefault((type(value).__name__, str(value)),
[]).append(row)
+ total = sum(float(row[label]) for row in rows)
+ offset = 0.0
+ for (_, name), children in groups.items():
+ value = sum(float(row[label]) for row in children)
+ ratio = value / total if total else 0
+ nx, ny = (
+ (x + offset * width, y) if depth % 2 == 0 else (x, y + offset
* height)
+ )
+ nw, nh = (
+ (width * ratio, height) if depth % 2 == 0 else (width, height
* ratio)
+ )
+ node_path = [*path, name]
+ leaf = depth == len(hierarchy) - 1
+ nodes.append(
+ {
+ "x0": nx,
+ "y0": ny,
+ "x1": nx + nw,
+ "y1": ny + nh,
+ "label_y": ny + 14 * depth,
+ "name": name,
+ "path": " > ".join(node_path),
+ "value": value,
+ "leaf": leaf,
+ "depth": depth,
+ "percent": ratio,
+ }
+ )
+ if not leaf:
+ layout(children, depth + 1, node_path, nx, ny, nw, nh)
+ offset += ratio
+
+ layout(data, 0, [], 0, 0, 600, 400)
+ coordinates: dict[str, dict[str, Any]] = {
+ "x": {
+ "field": "x0",
+ "type": "quantitative",
+ "scale": {"domain": [0, 600]},
+ "axis": None,
+ },
+ "x2": {"field": "x1"},
+ "y": {
+ "field": "y0",
+ "type": "quantitative",
+ "scale": {"domain": [400, 0]},
+ "axis": None,
+ },
+ "y2": {"field": "y1"},
+ }
+ layers: list[dict[str, Any]] = [
+ {
+ "mark": {"type": "rect", "stroke": "white", "strokeWidth": 1},
+ "encoding": {
+ **coordinates,
+ "color": {
+ "field": "name",
+ "type": "nominal",
+ "scale": {
+ "range": _PALETTES[scheme],
+ # Native traversal assigns the metric root first.
+ "domain": list(
+ dict.fromkeys([label, *(node["name"] for node in
nodes)])
+ ),
+ },
+ "legend": None,
+ },
+ "tooltip": [
+ {"field": "path", "title": "Hierarchy"},
+ {"field": "value", "type": "quantitative", "title": label},
+ {
+ "field": "percent",
+ "type": "quantitative",
+ "format": ".2%",
+ "title": "Share of parent",
+ },
+ ],
+ },
+ }
+ ]
+ if form_data.get("show_labels", True) or
form_data.get("show_upper_labels", True):
+ label_type = form_data.get("label_type", "key_value")
+ number_format = form_data.get("number_format") or "SMART_NUMBER"
+ from superset.utils import json
+
+ fmt = json.dumps(".3~s" if number_format == "SMART_NUMBER" else
number_format)
+ expression = (
+ "datum.name"
+ if label_type in ("key", "Key")
+ else f"format(datum.value, {fmt})"
+ )
+ if label_type == "key_value":
+ expression = f"datum.name + ': ' + {expression}"
+ leaves = str(bool(form_data.get("show_labels", True))).lower()
+ parents = str(bool(form_data.get("show_upper_labels", True))).lower()
+ layers.append(
+ {
+ "transform": [
+ {
+ "filter": (
+ f"(datum.leaf ? {leaves} : {parents}) "
+ "&& datum.x1-datum.x0 > 50 "
+ "&& datum.y1-datum.label_y > 20"
+ )
+ },
+ {"calculate": expression, "as": "label"},
+ ],
+ "mark": {
+ "type": "text",
+ "align": "left",
+ "baseline": "top",
+ "dx": 3,
+ "dy": 3,
+ "limit": {"expr": "datum.x1 - datum.x0 - 6"},
+ },
+ "encoding": {
+ "x": coordinates["x"],
+ "y": {**coordinates["y"], "field": "label_y"},
+ "text": {"field": "label"},
+ },
+ }
+ )
+ return VegaLitePreview(
+ specification={
+ "$schema": "https://vega.github.io/schema/vega-lite/v5.json";,
+ "width": 600,
+ "height": 400,
+ "data": {"values": nodes},
+ "layer": layers,
+ "usermeta": {
+ "viz_type": "treemap_v2",
+ "layout": "slice-and-dice",
+ "hierarchy": hierarchy,
+ "metric": label,
+ "row_count": len(data),
+ "native_renderer": False,
+ "format_note": (
+ "SMART_NUMBER uses SI approximation; use Explore for
native "
+ "formatting and interactions."
+ ),
+ },
+ }
+ )
diff --git a/superset/mcp_service/explore/tool/generate_explore_link.py
b/superset/mcp_service/explore/tool/generate_explore_link.py
index 9d8656da848..12c4c87f1b3 100644
--- a/superset/mcp_service/explore/tool/generate_explore_link.py
+++ b/superset/mcp_service/explore/tool/generate_explore_link.py
@@ -280,7 +280,8 @@ async def generate_explore_link(
normalized_config,
form_data,
dataset,
- run_compile_check=normalized_config.chart_type == "gauge",
+ run_compile_check=normalized_config.chart_type
+ in ("gauge", "treemap_v2"),
)
if not compile_result.success:
await ctx.warning(
diff --git a/tests/unit_tests/mcp_service/chart/test_treemap_completeness.py
b/tests/unit_tests/mcp_service/chart/test_treemap_completeness.py
new file mode 100644
index 00000000000..6bf55284d67
--- /dev/null
+++ b/tests/unit_tests/mcp_service/chart/test_treemap_completeness.py
@@ -0,0 +1,1373 @@
+# 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 regressions spanning native requests, query output and update
semantics."""
+
+import re
+from contextlib import nullcontext
+from copy import deepcopy
+from decimal import Decimal
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+from pydantic import ValidationError
+
+from superset.mcp_service.chart.chart_helpers import _build_single_query_dict
+from superset.mcp_service.chart.chart_utils import (
+ map_treemap_config,
+ merge_chart_form_data,
+ resolve_treemap_update_config,
+)
+from superset.mcp_service.chart.query_result import
normalize_chart_query_result
+from superset.mcp_service.chart.schemas import (
+ ChartError,
+ GenerateChartRequest,
+ GetChartPreviewRequest,
+ TreemapChartConfig,
+ UpdateChartPreviewRequest,
+ UpdateChartRequest,
+)
+from superset.mcp_service.chart.treemap_preview import treemap_ascii,
treemap_vega_lite
+from superset.utils import json
+
+FORM_DATA: dict[str, Any] = {
+ "viz_type": "treemap_v2",
+ "groupby": ["region", "product"],
+ "metric": "revenue",
+ "datasource": "7__table",
+ "color_scheme": "lyftColors",
+ "row_limit": 7,
+ "sort_by_metric": False,
+ "show_labels": False,
+ "show_upper_labels": False,
+ "number_format": ",.1f",
+ "date_format": "%Y",
+ "label_type": "value",
+ "currency_format": {"symbol": "USD", "symbolPosition": "prefix"},
+ "adhoc_filters": [
+ {
+ "expressionType": "SIMPLE",
+ "clause": "WHERE",
+ "subject": "region",
+ "operator": "==",
+ "comparator": "West",
+ }
+ ],
+ "time_range": "2025-01-01 : 2026-01-01",
+ "granularity_sqla": "ds",
+ "template_params": '{"scale": 2}',
+}
+ROWS = [
+ {"region": "West", "product": "A", "revenue": 30},
+ {"region": "West", "product": "B", "revenue": 10},
+ {"region": "East", "product": "A", "revenue": 60},
+]
+
+
[email protected]("sort", [False, True])
[email protected]("limit", [None, 0, 1, 7])
+def test_hierarchy_query_order_matches_frontend(sort: bool, limit: int | None)
-> None:
+ """Metric order has precedence, with hierarchy tie-breakers only when
bounded."""
+ form = {**FORM_DATA, "sort_by_metric": sort, "row_limit": limit}
+ query = _build_single_query_dict(form, form["groupby"], [form["metric"]])
+ expected = ([("revenue", False)] if sort else []) + [
+ ("region", True),
+ ("product", True),
+ ]
+ assert query.get("orderby", []) == (expected if limit else [])
+
+
[email protected](
+ "metric",
+ [
+ "revenue",
+ {"name": "revenue", "saved_metric": True},
+ {
+ "expressionType": "SIMPLE",
+ "column": {"column_name": "amount"},
+ "aggregate": "SUM",
+ "label": "revenue",
+ "hasCustomLabel": True,
+ },
+ {"expressionType": "SQL", "sqlExpression": "SUM(amount)", "label":
"revenue"},
+ ],
+)
[email protected](
+ "request_class", [GenerateChartRequest, UpdateChartPreviewRequest]
+)
+def test_native_request_roundtrip(metric: Any, request_class: Any) -> None:
+ """Native columns and all valid metric shapes survive actual request
validation."""
+ native = {**FORM_DATA, "metric": metric, "slice_id": 1}
+ request = request_class(dataset_id=7, config=native)
+ result = map_treemap_config(request.config)
+ assert result["groupby"] == FORM_DATA["groupby"]
+ assert result["show_labels"] is False
+ assert result["number_format"] == ",.1f"
+ assert result["metric"] == "revenue" or result["metric"]["label"] ==
"revenue"
+
+
[email protected](
+ "field,value",
+ [
+ ("javascript", "alert(1)"),
+ ("groupby", []),
+ ("metric", {"expressionType": "SQL", "sqlExpression": "SUM(amount)"}),
+ ("row_limit", 10001),
+ ],
+)
+def test_native_requests_do_not_accept_unbounded_or_hostile_controls(
+ field: str, value: Any
+) -> None:
+ """Unknown executable controls and invalid native roles remain validation
errors."""
+ with pytest.raises(ValidationError):
+ GenerateChartRequest(dataset_id=7, config={**FORM_DATA, field: value})
+
+
[email protected](
+ "patch_data",
+ [{"show_labels": True}, {"groupby": ["product", "region"], "metric":
"revenue"}],
+)
+def test_omitted_same_viz_controls_survive(patch_data: dict[str, Any]) -> None:
+ """Partial and full-role updates preserve all omitted native controls."""
+ request = UpdateChartRequest(
+ identifier=1, config={"chart_type": "treemap_v2", **patch_data}
+ )
+ assert request.config is not None
+ config = resolve_treemap_update_config(request.config, FORM_DATA)
+ merged = merge_chart_form_data(FORM_DATA, map_treemap_config(config),
config)
+ for key, value in FORM_DATA.items():
+ if key not in patch_data:
+ assert merged[key] == value
+ for key, value in patch_data.items():
+ assert merged[key] == value
+
+
[email protected](
+ "field",
+ [
+ "color_scheme",
+ "currency_format",
+ "time_range",
+ "granularity_sqla",
+ "template_params",
+ "filters",
+ ],
+)
+def test_nullable_clears_are_not_replaced_by_defaults(field: str) -> None:
+ """Explicit null clears native nullable controls rather than restoring
defaults."""
+ request = UpdateChartRequest(
+ identifier=1, config={"chart_type": "treemap_v2", field: None}
+ )
+ assert request.config is not None
+ config = resolve_treemap_update_config(request.config, FORM_DATA)
+ merged = merge_chart_form_data(FORM_DATA, map_treemap_config(config),
config)
+ assert ("adhoc_filters" if field == "filters" else field) not in merged
+
+
+def test_empty_filters_and_dataset_rebind() -> None:
+ """Rebinds retain presentation only and scrub stale query/template
roles."""
+ config = TreemapChartConfig(
+ groupby=[{"name": "other"}], metric="other_metric", filters=[]
+ )
+ existing = {**FORM_DATA, "metrics": ["stale"], "x_axis": "stale"}
+ merged = merge_chart_form_data(
+ existing, map_treemap_config(config), config, dataset_rebind=True
+ )
+ assert merged["color_scheme"] == "lyftColors"
+ assert merged["show_labels"] is False
+ assert merged["groupby"] == ["other"]
+ assert merged["metric"] == "other_metric"
+ for key in (
+ "adhoc_filters",
+ "time_range",
+ "granularity_sqla",
+ "template_params",
+ "metrics",
+ "x_axis",
+ ):
+ assert key not in merged
+ incomplete = UpdateChartRequest(
+ identifier=1, config={"chart_type": "treemap_v2", "show_labels": True}
+ )
+ assert incomplete.config is not None
+ with pytest.raises(ValidationError):
+ resolve_treemap_update_config(incomplete.config, existing,
dataset_rebind=True)
+
+
[email protected](
+ "result",
+ [
+ None,
+ {},
+ {"queries": []},
+ {"queries": [{}]},
+ {"queries": [{"data": {}}]},
+ {"queries": [{"data": [None]}]},
+ {"queries": [{"data": [{"region": "West", "revenue": 2}]}]},
+ {"status": "failed", "message": "Database rejected SQL"},
+ {"queries": [{"error": "SQL failed", "data": ROWS}]},
+ ],
+)
+def test_malformed_results_never_produce_success(result: Any) -> None:
+ """Missing hierarchy and malformed/error envelopes are actionable
failures."""
+ assert isinstance(normalize_chart_query_result(result, FORM_DATA),
ChartError)
+
+
[email protected](
+ "value", [None, True, "3", float("nan"), float("inf"), 10**400]
+)
+def test_metric_outputs_must_be_finite_numeric(value: Any) -> None:
+ """A single invalid metric fails instead of silently dropping a hierarchy
node."""
+ rows = [*ROWS, {**ROWS[0], "revenue": value}]
+ result = normalize_chart_query_result({"queries": [{"data": rows}]},
FORM_DATA)
+ assert isinstance(result, ChartError)
+ assert result.error_type == "InvalidTreemapMetric"
+
+
+def test_result_validation_is_non_mutating_and_chart_specific() -> None:
+ """Other charts retain their existing interpretation of data."""
+ result = {"queries": [{"data": ROWS}]}
+ original = deepcopy(result)
+ assert normalize_chart_query_result(result, FORM_DATA) == original
+ assert result == original
+ unrelated = {"queries": [{"data": [{"x": None}]}]}
+ assert normalize_chart_query_result(unrelated, {"viz_type": "table"}) is
unrelated
+
+
+def test_treemap_geometry_area_hierarchy_labels_and_color() -> None:
+ """Coordinates partition parents and area follows the metric, not row
count."""
+ preview = treemap_vega_lite(
+ ROWS,
+ {
+ **FORM_DATA,
+ "currency_format": None,
+ "show_labels": True,
+ "number_format": None,
+ },
+ )
+ assert not isinstance(preview, ChartError)
+ spec = preview.specification
+ nodes = spec["data"]["values"]
+ leaves = [node for node in nodes if node["leaf"]]
+ for leaf in leaves:
+ area = (leaf["x1"] - leaf["x0"]) * (leaf["y1"] - leaf["y0"])
+ assert area / (600 * 400) == pytest.approx(leaf["value"] / 100)
+ parent = next(
+ node for node in nodes if node["path"] == leaf["path"].split(" >
")[0]
+ )
+ assert parent["x0"] <= leaf["x0"] <= leaf["x1"] <= parent["x1"]
+ assert parent["y0"] <= leaf["y0"] <= leaf["y1"] <= parent["y1"]
+ assert spec["layer"][0]["mark"]["type"] == "rect"
+ assert spec["layer"][0]["encoding"]["color"]["scale"]["range"][0] ==
"#EA0B8C"
+ assert spec["layer"][0]["encoding"]["tooltip"][1]["title"] == "revenue"
+ assert spec["layer"][1]["mark"]["type"] == "text"
+ assert "West > A | 30" in treemap_ascii(ROWS, FORM_DATA)
+
+
[email protected](
+ "rows,override",
+ [
+ (ROWS, {"color_scheme": "unavailable"}),
+ (ROWS, {"currency_format": {"symbol": "USD"}}),
+ ([{**ROWS[0], "revenue": -1}], {}),
+ (ROWS * 334, {}),
+ ],
+)
+def test_unsupported_geometry_is_explicit(rows: Any, override: Any) -> None:
+ """Unsupported representation must not fall back to a plausible bar or
scatter."""
+ preview = treemap_vega_lite(
+ rows, {**FORM_DATA, "currency_format": None, **override}
+ )
+ assert isinstance(preview, ChartError)
+ assert preview.error_type == "UnsupportedTreemapPreview"
+
+
[email protected]("format_name", ["ascii", "table", "vega_lite"])
+def test_saved_and_unsaved_preview_dispatch_match(format_name: str) -> None:
+ """Both preview entry paths use the Treemap representation and result
contract."""
+ from superset.mcp_service.chart.preview_utils import
generate_preview_from_form_data
+ from superset.mcp_service.chart.tool.get_chart_preview import (
+ ASCIIPreviewStrategy,
+ TablePreviewStrategy,
+ VegaLitePreviewStrategy,
+ )
+
+ form = {**FORM_DATA, "currency_format": None}
+ chart = SimpleNamespace(
+ id=1,
+ viz_type="treemap_v2",
+ slice_name="Treemap",
+ datasource_id=7,
+ datasource_type="table",
+ params=json.dumps(form),
+ )
+ context = SimpleNamespace(
+ queries=[SimpleNamespace(metrics=["revenue"], columns=form["groupby"])]
+ )
+ with (
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand"
+ ) as command,
+ patch(
+
"superset.mcp_service.chart.tool.get_chart_preview.build_query_context_from_form_data",
+ return_value=context,
+ ),
+ patch(
+
"superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=context,
+ ),
+ patch("superset.extensions.db.session") as session,
+ ):
+ session.get.return_value = Mock(id=7)
+ command.return_value.run.return_value = {"queries": [{"data": ROWS}]}
+ strategy = {
+ "ascii": ASCIIPreviewStrategy,
+ "table": TablePreviewStrategy,
+ "vega_lite": VegaLitePreviewStrategy,
+ }[format_name]
+ saved = strategy(
+ chart, GetChartPreviewRequest(identifier=1, format=format_name)
+ ).generate()
+ unsaved = generate_preview_from_form_data(form, 7, format_name)
+ assert not isinstance(saved, ChartError)
+ assert not isinstance(unsaved, ChartError)
+ if format_name == "vega_lite":
+ assert saved.specification == unsaved.specification
+ elif format_name == "ascii":
+ assert "West > A | 30" in saved.ascii_content
+ assert "West > A | 30" in unsaved.ascii_content
+ else:
+ assert "region" in saved.table_data
+ assert "revenue" in unsaved.table_data
+
+
[email protected]("rebind", [False, True])
+def test_saved_update_and_update_preview_payloads_agree(rebind: bool) -> None:
+ """Both save and preview-first update paths apply the same omission
contract."""
+ from superset.mcp_service.chart.tool.update_chart import (
+ _build_preview_form_data,
+ _build_update_payload,
+ )
+
+ config = TreemapChartConfig(
+ groupby=[{"name": "region"}], metric="revenue", show_labels=True
+ )
+ request = UpdateChartRequest(
+ identifier=1, config=config, dataset_id=8 if rebind else None
+ )
+ chart = Mock(
+ id=1, datasource_id=7, slice_name="Treemap",
params=json.dumps(FORM_DATA)
+ )
+ with patch(
+
"superset.mcp_service.chart.chart_utils._bind_dashboard_time_range_filter"
+ ):
+ preview = _build_preview_form_data(request, chart,
parsed_config=config)
+ payload = _build_update_payload(request, chart, parsed_config=config)
+ saved = json.loads(payload["params"])
+ for key in (
+ "color_scheme",
+ "row_limit",
+ "sort_by_metric",
+ "groupby",
+ "metric",
+ "number_format",
+ "show_labels",
+ ):
+ assert saved[key] == preview[key]
+ assert saved["color_scheme"] == "lyftColors"
+ assert saved["show_labels"] is True
+ assert ("template_params" not in saved) == rebind
+
+
+def test_normalization_preserves_explicit_field_set_and_rejects_ambiguity() ->
None:
+ """Canonicalization must not turn default values into explicit update
intent."""
+ from superset.mcp_service.chart.plugins.treemap import TreemapChartPlugin
+ from superset.mcp_service.common.error_schemas import DatasetContext
+
+ context = DatasetContext(
+ id=7,
+ table_name="sales",
+ database_name="database",
+ available_columns=[{"name": "Region", "type": "STRING"}],
+ available_metrics=[{"name": "Revenue"}],
+ )
+ config = TreemapChartConfig(groupby=[{"name": "region"}], metric="revenue")
+ normalized = TreemapChartPlugin().normalize_column_refs(config, context)
+ assert normalized.groupby[0].name == "Region"
+ assert normalized.metric.name == "Revenue"
+ assert normalized.model_fields_set == config.model_fields_set
+ context.available_columns.append({"name": "REGION", "type": "STRING"})
+ with pytest.raises(ValueError, match="[Aa]mbiguous"):
+ TreemapChartPlugin().normalize_column_refs(config, context)
+
+
[email protected]
[email protected](
+ "metric",
+ [
+ "revenue",
+ {
+ "expressionType": "SIMPLE",
+ "column": {"column_name": "amount"},
+ "aggregate": "SUM",
+ "label": "revenue",
+ },
+ {"expressionType": "SQL", "sqlExpression": "SUM(amount)", "label":
"revenue"},
+ ],
+)
[email protected]("valid_result", [False, True])
+async def test_registered_generate_chart_native_roundtrip(
+ metric: Any, valid_result: bool
+) -> None:
+ """Call the registered FastMCP tool, retaining real request and compile
checks."""
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+
+ request = GenerateChartRequest(
+ dataset_id=7,
+ config={
+ "chart_type": "treemap_v2",
+ "groupby": ["region", "product"],
+ "metric": metric,
+ "currency_format": {"symbol": "USD", "symbolPosition": "suffix"},
+ },
+ preview_formats=["url"],
+ )
+ dataset = Mock(id=7, datasource_name="sales", table_name="sales")
+ user = Mock(id=1, username="admin", roles=[], groups=[])
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request",
return_value=user),
+ patch(
+
"superset.mcp_service.chart.validation.ValidationPipeline.validate_request_with_warnings",
+ return_value=Mock(is_valid=True, request=request, warnings={},
error=None),
+ ),
+ patch(
+ "superset.mcp_service.chart.chart_utils.generate_explore_link",
+ return_value="http://localhost/explore/?form_data_key=treemap";,
+ ),
+ patch("superset.daos.dataset.DatasetDAO.find_by_id",
return_value=dataset),
+ patch(
+
"superset.mcp_service.chart.tool.generate_chart.has_dataset_access",
+ return_value=True,
+ ),
+ patch(
+
"superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=Mock(),
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand"
+ ) as command,
+ ):
+ command.return_value.run.return_value = {
+ "queries": [
+ {"data": ROWS if valid_result else [{"region": "West",
"revenue": 1}]}
+ ]
+ }
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "generate_chart",
+ {
+ "request": {
+ "dataset_id": 7,
+ "config": {
+ "viz_type": "treemap_v2",
+ "groupby": ["region", "product"],
+ "metric": metric,
+ "currency_format": {
+ "symbol": "USD",
+ "symbolPosition": "suffix",
+ },
+ },
+ "preview_formats": ["url"],
+ }
+ },
+ )
+ data = result.structured_content
+ assert data["success"] is valid_result
+ if valid_result:
+ assert data["form_data"]["currency_format"] == {
+ "symbol": "USD",
+ "symbolPosition": "suffix",
+ }
+ assert data["form_data"]["groupby"] == ["region", "product"]
+ assert data["form_data"]["viz_type"] == "treemap_v2"
+ else:
+ assert data["error"]["error_type"] == "InvalidTreemapResult"
+ command.return_value.validate.assert_called_once()
+
+
[email protected]("number_format", [None, ",.1f"])
+def test_vega_scenegraph_renders_nested_metric_geometry(
+ number_format: str | None,
+) -> None:
+ """Compile the specification and inspect rendered rectangles, not a
snapshot."""
+ import os
+ import shutil
+ import subprocess
+
+ if not os.environ.get("NODE_PATH"):
+ pytest.skip("Requires Node vega@5 and vega-lite@5 via NODE_PATH")
+ spec = treemap_vega_lite(
+ ROWS,
+ {
+ **FORM_DATA,
+ "currency_format": None,
+ "show_labels": True,
+ "number_format": number_format,
+ },
+ )
+ assert not isinstance(spec, ChartError)
+ script = r"""
+const assert = require('assert/strict');
+const vl = require('vega-lite');
+const vega = require('vega');
+(async () => {
+ const input = JSON.parse(require('fs').readFileSync(0, 'utf8'));
+ const warnings = [];
+ const compiled = vl.compile(input, {logger: {
+ warn: message => warnings.push(message), info() {}, debug() {},
+ error: message => { throw Error(message); }
+ }}).spec;
+ assert.deepEqual(warnings, []);
+ const view = new vega.View(vega.parse(compiled), {renderer: 'none'});
+ await view.runAsync();
+ const rectangles = [];
+ const labels = [];
+ function walk(item) {
+ if (item.mark?.marktype === 'rect') rectangles.push(item);
+ if (item.mark?.marktype === 'text') labels.push(item);
+ for (const child of item.items || []) walk(child);
+ }
+ walk(view.scenegraph().root);
+ assert.equal(rectangles.length, 5);
+ assert(labels.length > 0);
+ for (const tile of rectangles.filter(tile => tile.datum.leaf)) {
+ assert(tile.width >= 0 && tile.height >= 0);
+ assert(Math.abs(tile.width * tile.height / 240000 - tile.datum.value /
100) < 1e-8);
+ assert(tile.tooltip.Hierarchy.includes(' > '));
+ assert(tile.fill.startsWith('#'));
+ }
+ assert((await view.toSVG()).includes('<svg'));
+ view.finalize();
+})().catch(error => { console.error(error); process.exit(1); });
+"""
+ node = shutil.which("node")
+ assert node is not None
+ # Run a fixed test script; generated data is supplied only on stdin.
+ result = subprocess.run( # noqa: S603
+ [node, "-e", script],
+ input=json.dumps(spec.specification),
+ capture_output=True,
+ text=True,
+ check=False,
+ )
+ assert result.returncode == 0, result.stderr
+
+
[email protected]
[email protected]("source", ["saved", "cached_saved", "cached_unsaved"])
[email protected]("row_limit", [1, 7])
[email protected]("format_name", ["ascii", "table", "vega_lite"])
+async def test_registered_cached_preview_is_treemap(
+ source: str, row_limit: int, format_name: str
+) -> None:
+ """Cached state overrides saved viz and uses identical Treemap dispatch."""
+ import importlib
+
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_preview"
+ )
+ form = {**FORM_DATA, "currency_format": None, "row_limit": row_limit}
+ chart = SimpleNamespace(
+ id=1,
+ viz_type="treemap_v2" if source == "saved" else "table",
+ slice_name="Chart",
+ datasource_id=7,
+ datasource_type="table",
+ params=json.dumps(form if source == "saved" else {"viz_type":
"table"}),
+ )
+ rows = [
+ {**row, "revenue": Decimal(str(row["revenue"]))} for row in
ROWS[:row_limit]
+ ]
+ if row_limit == 7:
+ rows[0]["region"], rows[1]["region"] = 1, "1"
+ context = SimpleNamespace(
+ queries=[SimpleNamespace(metrics=["revenue"], columns=form["groupby"])]
+ )
+ with (
+ patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=Mock(id=1, username="admin", roles=[], groups=[]),
+ ),
+ patch.object(module, "find_chart_by_identifier", return_value=chart),
+ patch.object(module.db.session, "refresh"),
+ patch.object(
+ module.event_logger,
+ "log_context",
+ side_effect=lambda **kwargs: nullcontext(),
+ ),
+ patch.object(
+ module,
+ "validate_chart_dataset",
+ return_value=Mock(is_valid=True, warnings=[]),
+ ),
+ patch.object(
+ module, "build_query_context_from_form_data", return_value=context
+ ) as build_context,
+ patch(
+ "superset.commands.explore.form_data.get.GetFormDataCommand.run",
+ return_value=json.dumps(form),
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand"
+ ) as command,
+ ):
+ command.return_value.run.return_value = {"queries": [{"data": rows}]}
+ request = {"format": format_name}
+ if source != "saved":
+ request["form_data_key"] = "treemap-key"
+ if source != "cached_unsaved":
+ request["identifier"] = "1"
+ async with Client(mcp) as client:
+ result = await client.call_tool("get_chart_preview", {"request":
request})
+ assert not result.is_error
+ assert build_context.call_args.kwargs["row_limit"] == row_limit
+ data = json.loads(result.content[0].text)
+ assert "error_type" not in data, data
+ content = data["content"]
+ assert content["type"] == format_name
+ if format_name == "vega_lite":
+ spec = content["specification"]
+ assert spec["layer"][0]["mark"]["type"] == "rect"
+ if row_limit == 7:
+ parents = [
+ n for n in spec["data"]["values"] if not n["leaf"] and
n["name"] == "1"
+ ]
+ assert len(parents) == 2
+ assert [n["value"] for n in parents] == [30, 10]
+ assert [
+ (n["x1"] - n["x0"]) * (n["y1"] - n["y0"]) / (600 * 400) for n
in parents
+ ] == pytest.approx([0.3, 0.1])
+ elif format_name == "ascii":
+ assert f"{rows[0]['region']} > A | 30" in content["ascii_content"]
+ else:
+ assert "revenue" in content["table_data"]
+
+
[email protected]
[email protected](
+ "patch_data", [{"show_labels": True}, {"filters": []}, {"color_scheme":
None}]
+)
[email protected]("known_dataset", [True, False, None])
+async def test_registered_update_preview_preserves_cached_controls(
+ patch_data: dict[str, Any],
+ known_dataset: bool | None,
+) -> None:
+ """Partial native updates reach real FastMCP hydration, merge, and cache
writes."""
+ import importlib
+
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+ from superset.mcp_service.chart.compile import CompileResult
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.update_chart_preview"
+ )
+ dataset = Mock(id=7, table_name="sales", schema=None, columns=[],
metrics=[])
+ with (
+ patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=Mock(id=1, username="admin", roles=[], groups=[]),
+ ),
+ patch.object(module, "_find_dataset", return_value=dataset),
+ patch.object(
+ module,
+ "_get_previous_form_data",
+ return_value=(
+ FORM_DATA
+ if known_dataset
+ else {
+ key: value
+ for key, value in FORM_DATA.items()
+ if key != "datasource"
+ }
+ if known_dataset is False
+ else None
+ ),
+ ),
+ patch("superset.daos.dataset.DatasetDAO.find_by_id",
return_value=dataset),
+ patch.object(module, "has_dataset_access", return_value=True),
+ patch(
+
"superset.mcp_service.chart.validation.dataset_validator.DatasetValidator.normalize_column_names",
+ side_effect=lambda config, *args, **kwargs: config,
+ ),
+ patch(
+
"superset.mcp_service.chart.chart_utils._bind_dashboard_time_range_filter"
+ ),
+ patch.object(
+ module, "validate_and_compile",
return_value=CompileResult(success=True)
+ ) as compile_check,
+ patch.object(
+ module,
+ "generate_explore_link",
+ return_value="http://localhost/explore/?form_data_key=updated";,
+ ) as cache_write,
+ ):
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "update_chart_preview",
+ {
+ "request": {
+ "dataset_id": 7,
+ "form_data_key": "previous",
+ "config": {"chart_type": "treemap_v2", **patch_data},
+ "generate_preview": False,
+ }
+ },
+ )
+ assert not result.is_error
+ data = result.structured_content
+ if not known_dataset:
+ assert data["error"]["error_type"] == "ValidationError"
+ assert data["success"] is False
+ cache_write.assert_not_called()
+ return
+ assert data["success"] is True, data
+ merged = data["form_data"]
+ assert merged["row_limit"] == 7
+ assert merged["sort_by_metric"] is False
+ assert merged["metric"] == "revenue"
+ assert merged["groupby"] == ["region", "product"]
+ for key, value in FORM_DATA.items():
+ field = "filters" if key == "adhoc_filters" else key
+ if field not in patch_data:
+ assert merged[key] == value
+ assert compile_check.call_args.kwargs["run_compile_check"] is True
+ assert cache_write.call_args.args[1] == merged
+
+
[email protected]
[email protected]("malformed", [False, True])
+async def test_registered_saved_update_preserves_omissions(malformed: bool) ->
None:
+ """The registered save path persists the same partial Treemap merge as
preview."""
+ import importlib
+
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+
+ module =
importlib.import_module("superset.mcp_service.chart.tool.update_chart")
+ chart = Mock(
+ id=1,
+ datasource_id=7,
+ slice_name="Treemap",
+ viz_type="treemap_v2",
+ uuid="11111111-1111-1111-1111-111111111111",
+ params=json.dumps(
+ {**FORM_DATA, "metric": {"expressionType": "SIMPLE"}}
+ if malformed
+ else FORM_DATA
+ ),
+ )
+ with (
+ patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=Mock(id=1, username="admin", roles=[], groups=[]),
+ ),
+ patch.object(module, "find_chart_by_identifier", return_value=chart),
+ patch(
+ "superset.mcp_service.auth.check_chart_data_access",
+ return_value=Mock(is_valid=True),
+ ),
+ patch(
+
"superset.mcp_service.chart.validation.dataset_validator.DatasetValidator.normalize_column_names",
+ side_effect=lambda config, *args, **kwargs: config,
+ ),
+ patch.object(
+ module, "_validate_update_against_dataset", return_value=None
+ ) as validate,
+ patch(
+
"superset.mcp_service.chart.chart_utils._bind_dashboard_time_range_filter"
+ ),
+ patch("superset.commands.chart.update.UpdateChartCommand") as update,
+ patch("superset.db.session"),
+ ):
+ update.return_value.run.return_value = chart
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "update_chart",
+ {
+ "request": {
+ "identifier": 1,
+ "generate_preview": False,
+ "preview_formats": [],
+ "config": {"chart_type": "treemap_v2", "show_labels":
True},
+ }
+ },
+ )
+ data = result.structured_content
+ if malformed:
+ from superset.mcp_service.chart.schemas import GenerateChartResponse
+
+ assert not result.is_error
+ assert GenerateChartResponse.model_validate(data).success is False
+ assert data["success"] is False
+ assert data["error"]["error_type"] == "ValidationError"
+ assert data["error"]["message"] == "Invalid Treemap update
configuration"
+ update.assert_not_called()
+ return
+ assert data["success"] is True, data
+ assert data["chart"]["is_unsaved_state"] is False
+ persisted = json.loads(update.call_args.args[1]["params"])
+ for key, value in FORM_DATA.items():
+ assert persisted[key] == (True if key == "show_labels" else value)
+ validate.assert_called_once()
+
+
+def test_treemap_query_ignores_stale_cross_chart_roles() -> None:
+ """Saved/native Treemaps must not query a previous plugin's raw columns."""
+ from superset.mcp_service.chart.chart_helpers import
resolve_metrics_and_groupby
+
+ form = {
+ **FORM_DATA,
+ "query_mode": "raw",
+ "all_columns": ["stale"],
+ "metrics": ["stale_metric"],
+ }
+ metrics, hierarchy = resolve_metrics_and_groupby(form)
+ assert metrics == ["revenue"]
+ assert hierarchy == ["region", "product"]
+
+
[email protected]("limit", ["0", "0.0", "", "1"])
+def test_native_string_row_limits_match_frontend(limit: str) -> None:
+ """Frontend applyOrderBy numerically parses string row limits."""
+ query = _build_single_query_dict(
+ {**FORM_DATA, "row_limit": limit}, ["region"], ["revenue"]
+ )
+ assert query.get("orderby", []) == ([("region", True)] if limit == "1"
else [])
+
+
+def test_explicit_temporal_clear_keeps_user_filters_and_template_state() ->
None:
+ """Removing a generated time binding must not erase user-owned
predicates."""
+ from superset.mcp_service.chart.chart_utils import
MCP_DASHBOARD_TIME_FILTER_SUBJECT
+ from superset.mcp_service.chart.tool.update_chart import
_build_preview_form_data
+
+ existing = {
+ **FORM_DATA,
+ MCP_DASHBOARD_TIME_FILTER_SUBJECT: "ds",
+ "adhoc_filters": [
+ *FORM_DATA["adhoc_filters"],
+ {
+ "expressionType": "SIMPLE",
+ "clause": "WHERE",
+ "subject": "ds",
+ "operator": "TEMPORAL_RANGE",
+ "comparator": "No filter",
+ },
+ ],
+ }
+ request = UpdateChartRequest(
+ identifier=1, config={"chart_type": "treemap_v2", "temporal_column":
None}
+ )
+ assert request.config is not None
+ config = resolve_treemap_update_config(request.config, existing)
+ chart = Mock(
+ id=1, datasource_id=7, slice_name="Treemap",
params=json.dumps(existing)
+ )
+ result = _build_preview_form_data(request, chart, parsed_config=config)
+ assert isinstance(result, dict)
+ assert result["adhoc_filters"] == FORM_DATA["adhoc_filters"]
+ assert result["template_params"] == FORM_DATA["template_params"]
+ assert MCP_DASHBOARD_TIME_FILTER_SUBJECT not in result
+
+
+def test_canonical_metric_reference_keeps_merged_default_sum() -> None:
+ """Preserve the community plugin's default SUM for a typed physical
column."""
+ config = TreemapChartConfig(groupby=[{"name": "region"}], metric={"name":
"amount"})
+ assert map_treemap_config(config)["metric"]["label"] == "SUM(amount)"
+ with pytest.raises(ValidationError):
+ TreemapChartConfig(
+ groupby=["region"],
+ metric={"expressionType": "SIMPLE", "column": {"column_name":
"amount"}},
+ )
+
+
+def test_compile_respects_treemap_row_limit() -> None:
+ """Compile must not validate rows excluded by a smaller native row
limit."""
+ from superset.mcp_service.chart.compile import _compile_chart
+
+ with (
+ patch(
+
"superset.mcp_service.chart.chart_helpers.build_query_context_from_form_data",
+ return_value=Mock(),
+ ) as build,
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand"
+ ) as command,
+ ):
+ command.return_value.run.return_value = {"queries": [{"data":
ROWS[:1]}]}
+ result = _compile_chart({**FORM_DATA, "row_limit": 1}, 7)
+ assert result.success is True
+ assert build.call_args.kwargs["row_limit"] == 1
+
+
[email protected]("selected", ["event_time", None])
[email protected]("override", ["dashboard_time", None, "omitted"])
+def test_treemap_selected_time_column_reaches_query(
+ selected: str | None,
+ override: str | None,
+) -> None:
+ """The selected time column and explicit dashboard overrides reach
execution."""
+ from superset.mcp_service.chart.chart_helpers import (
+ build_query_dicts_from_form_data,
+ )
+
+ form = {
+ "viz_type": "treemap_v2",
+ "groupby": ["region"],
+ "metric": "revenue",
+ "granularity_sqla": selected,
+ "time_range": "2026-01-01 : 2026-02-01",
+ }
+ extra = {} if override == "omitted" else {"granularity_sqla": override}
+ with patch(
+ "superset.mcp_service.chart.chart_helpers.resolve_datasource_engine",
+ return_value="postgresql",
+ ):
+ query = build_query_dicts_from_form_data(
+ form,
+ 7,
+ "table",
+ extra_form_data=extra,
+ )[0]
+ # Shared dashboard normalization treats a null extra as no override.
+ # Clearing the chart's selected column itself is covered by selected=None.
+ expected = selected if override in (None, "omitted") else override
+ assert query.get("granularity") == expected
+ assert query["time_range"] == form["time_range"]
+ assert query["columns"] == ["region"]
+ assert query["metrics"] == ["revenue"]
+
+
+def test_treemap_mixed_type_categories_remain_separate() -> None:
+ """Match frontend raw-key grouping and display-name categorical colors."""
+ preview = treemap_vega_lite(
+ [{"region": 1, "revenue": 1}, {"region": "1", "revenue": 3}],
+ {
+ "viz_type": "treemap_v2",
+ "groupby": ["region"],
+ "metric": "revenue",
+ "show_labels": True,
+ "color_scheme": "supersetColors",
+ },
+ )
+ assert not isinstance(preview, ChartError)
+ nodes = preview.specification["data"]["values"]
+ assert len(nodes) == 2
+ assert [node["name"] for node in nodes] == ["1", "1"]
+ assert [node["value"] for node in nodes] == [1, 3]
+
+
[email protected](
+ "value,valid",
+ [
+ (Decimal("2.75"), True),
+ (Decimal("1E1000"), True),
+ (Decimal("NaN"), False),
+ (Decimal("sNaN"), False),
+ (Decimal("Infinity"), False),
+ (Decimal("-Infinity"), False),
+ (True, False),
+ ("2.75", False),
+ ],
+)
+def test_decimal_result_contract(value: Any, valid: bool) -> None:
+ """SQL numeric values remain numeric; nonfinite and coercible strings do
not."""
+ rows = [{**ROWS[0], "revenue": value}]
+ result = {"queries": [{"data": rows}]}
+ checked = normalize_chart_query_result(result, FORM_DATA)
+ assert isinstance(checked, ChartError) is not valid
+ if valid:
+ assert checked is result
+ assert rows[0]["revenue"] is value
+
+
[email protected]
[email protected]("save", [False, True])
[email protected]("control", ["filters", "empty", "change_time",
"clear_time"])
+async def test_registered_filter_update_keeps_saved_temporal_binding(
+ save: bool,
+ control: str,
+) -> None:
+ """Exercise real mapping/binding/merge through both registered update
entries."""
+ import importlib
+
+ from fastmcp import Client
+
+ from superset.mcp_service.app import mcp
+ from superset.mcp_service.chart.chart_utils import
MCP_DASHBOARD_TIME_FILTER_SUBJECT
+ from superset.mcp_service.chart.compile import CompileResult
+
+ saved_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.update_chart"
+ )
+ preview_module = importlib.import_module(
+ "superset.mcp_service.chart.tool.update_chart_preview"
+ )
+ dataset = Mock(
+ id=7,
+ table_name="sales",
+ schema=None,
+ columns=[],
+ metrics=[],
+ main_dttm_col="default_time",
+ )
+ existing = {
+ **FORM_DATA,
+ "granularity_sqla": None,
+ MCP_DASHBOARD_TIME_FILTER_SUBJECT: "saved_time",
+ "adhoc_filters": [
+ *FORM_DATA["adhoc_filters"],
+ {
+ "expressionType": "SIMPLE",
+ "clause": "WHERE",
+ "subject": "saved_time",
+ "operator": "TEMPORAL_RANGE",
+ "comparator": "No filter",
+ },
+ ],
+ }
+ config: dict[str, Any] = {
+ "chart_type": "treemap_v2",
+ "currency_format": {"symbol": "USD", "symbolPosition": "suffix"},
+ }
+ if control in ("filters", "empty"):
+ config["filters"] = (
+ [{"column": "product", "op": "=", "value": "A"}]
+ if control == "filters"
+ else []
+ )
+ else:
+ config["temporal_column"] = "new_time" if control == "change_time"
else None
+ chart = Mock(
+ id=1,
+ datasource_id=7,
+ slice_name="Treemap",
+ viz_type="treemap_v2",
+ uuid="11111111-1111-1111-1111-111111111111",
+ params=json.dumps(existing),
+ )
+ with (
+ patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=Mock(id=1, username="admin", roles=[], groups=[]),
+ ),
+ patch.object(saved_module, "find_chart_by_identifier",
return_value=chart),
+ patch(
+ "superset.mcp_service.auth.check_chart_data_access",
+ return_value=Mock(is_valid=True),
+ ),
+ patch.object(preview_module, "_find_dataset", return_value=dataset),
+ patch.object(preview_module, "_get_previous_form_data",
return_value=existing),
+ patch.object(preview_module, "has_dataset_access", return_value=True),
+ patch("superset.daos.dataset.DatasetDAO.find_by_id",
return_value=dataset),
+ patch(
+
"superset.mcp_service.chart.chart_utils._find_dataset_by_id_or_uuid",
+ return_value=dataset,
+ ),
+ patch(
+
"superset.mcp_service.chart.chart_utils._is_temporal_for_dashboard_binding",
+ return_value=True,
+ ),
+ patch(
+
"superset.mcp_service.chart.validation.dataset_validator.DatasetValidator.normalize_column_names",
+ side_effect=lambda value, *args, **kwargs: value,
+ ),
+ patch.object(
+ saved_module, "_validate_update_against_dataset", return_value=None
+ ),
+ patch.object(
+ preview_module,
+ "validate_and_compile",
+ return_value=CompileResult(success=True),
+ ),
+ patch.object(
+ preview_module,
+ "generate_explore_link",
+ return_value="http://localhost/explore/?form_data_key=updated";,
+ ) as cache,
+ patch("superset.commands.chart.update.UpdateChartCommand") as update,
+ patch("superset.db.session"),
+ ):
+ update.return_value.run.return_value = chart
+ request: dict[str, Any] = {"config": config, "generate_preview": False}
+ if save:
+ request.update(identifier=1, preview_formats=[])
+ else:
+ request.update(dataset_id=7, form_data_key="previous")
+ async with Client(mcp) as client:
+ result = await client.call_tool(
+ "update_chart" if save else "update_chart_preview",
+ {"request": request},
+ )
+ assert not result.is_error
+ data = result.structured_content
+ assert data["success"] is True, data
+ merged = (
+ json.loads(update.call_args.args[1]["params"])
+ if save
+ else cache.call_args.args[1]
+ )
+ filters = merged.get("adhoc_filters", [])
+ if control == "empty":
+ assert not filters
+ assert MCP_DASHBOARD_TIME_FILTER_SUBJECT not in merged
+ else:
+ for predicate in FORM_DATA["adhoc_filters"]:
+ assert predicate in filters
+ if control == "clear_time":
+ assert MCP_DASHBOARD_TIME_FILTER_SUBJECT not in merged
+ else:
+ subject = "saved_time" if control == "filters" else "new_time"
+ assert merged[MCP_DASHBOARD_TIME_FILTER_SUBJECT] == subject
+ temporal = [f for f in filters if f["operator"] ==
"TEMPORAL_RANGE"]
+ assert len(temporal) == 1
+ assert temporal[0]["subject"] == subject
+ assert temporal[0]["comparator"] == "No filter"
+ assert not any(f["subject"] == "default_time" for f in filters)
+ if control == "filters":
+ assert any(
+ f["subject"] == "product" and f["comparator"] == "A" for f in
filters
+ )
+ assert merged["currency_format"] == {"symbol": "USD", "symbolPosition":
"suffix"}
+
+
+def test_resolution_validation_error_is_caught_value_error() -> None:
+ """Use the installed Pydantic type, rather than assuming its
inheritance."""
+ request = UpdateChartPreviewRequest(
+ dataset_id=7, config={"chart_type": "treemap_v2", "show_labels": True}
+ )
+ with pytest.raises(ValidationError) as caught:
+ resolve_treemap_update_config(request.config, {})
+ assert isinstance(caught.value, ValueError)
+
+
+# --- Node-free equivalents of the scenegraph render assertion
-----------------
+#
+# ``test_vega_scenegraph_renders_nested_metric_geometry`` only runs where Node
+# with vega/vega-lite is installed, so the checks below reproduce the parts a
+# compile would catch — unresolvable ``datum`` references and unknown Vega
+# expression functions — using nothing but the standard library.
+
+_DATUM_REFERENCE = re.compile(r"datum\.([A-Za-z_$][A-Za-z0-9_$]*)")
+_FUNCTION_CALL = re.compile(r"\b([A-Za-z_$][A-Za-z0-9_$]*)\s*\(")
+# Vega expression functions the Treemap specification is allowed to emit.
+_ALLOWED_VEGA_FUNCTIONS = frozenset({"format", "if", "isValid", "join"})
+
+
+def _expression_strings(node: Any) -> list[str]:
+ """Collect every Vega expression embedded in a Vega-Lite specification."""
+ found: list[str] = []
+ if isinstance(node, dict):
+ for key, value in node.items():
+ if key in ("filter", "calculate", "expr") and isinstance(value,
str):
+ found.append(value)
+ else:
+ found.extend(_expression_strings(value))
+ elif isinstance(node, list):
+ for item in node:
+ found.extend(_expression_strings(item))
+ return found
+
+
+def _resolvable_fields(layer: dict[str, Any], data_keys: set[str]) -> set[str]:
+ """Return data keys plus every field the layer's transforms derive."""
+ derived = {
+ transform["as"]
+ for transform in layer.get("transform", [])
+ if isinstance(transform.get("as"), str)
+ }
+ return data_keys | derived
+
+
[email protected]("number_format", [None, ",.1f"])
[email protected]("show_labels", [False, True])
+def test_vega_spec_field_references_all_resolve(
+ number_format: str | None, show_labels: bool
+) -> None:
+ """Every encoding and expression field must exist in the data or a
transform."""
+ spec = treemap_vega_lite(
+ ROWS,
+ {
+ **FORM_DATA,
+ "currency_format": None,
+ "show_labels": show_labels,
+ "number_format": number_format,
+ },
+ )
+ assert not isinstance(spec, ChartError)
+ specification = spec.specification
+ data_keys = set(specification["data"]["values"][0])
+ assert specification["$schema"].startswith(
+ "https://vega.github.io/schema/vega-lite";
+ )
+ layers = specification["layer"]
+ assert layers
+
+ for layer in layers:
+ resolvable = _resolvable_fields(layer, data_keys)
+ for channel in layer["encoding"].values():
+ channels = channel if isinstance(channel, list) else [channel]
+ for definition in channels:
+ field = definition.get("field")
+ if field is not None:
+ assert field in resolvable, f"unknown encoding field
{field!r}"
+ for expression in _expression_strings(layer):
+ for reference in _DATUM_REFERENCE.findall(expression):
+ assert reference in resolvable, (
+ f"expression {expression!r} references unknown field
{reference!r}"
+ )
+ unknown = set(_FUNCTION_CALL.findall(expression)) -
_ALLOWED_VEGA_FUNCTIONS
+ assert not unknown, f"expression {expression!r} calls
{sorted(unknown)}"
+
+
+def test_vega_spec_field_reference_check_catches_a_bad_calculate() -> None:
+ """The Node-free check must fail on the typo class it exists to catch."""
+ spec = treemap_vega_lite(ROWS, {**FORM_DATA, "currency_format": None})
+ assert not isinstance(spec, ChartError)
+ layer = deepcopy(spec.specification["layer"][0])
+ layer.setdefault("transform", []).append(
+ {"calculate": "format(datum.valeu, ',.1f')", "as": "typo"}
+ )
+ data_keys = set(spec.specification["data"]["values"][0])
+ resolvable = _resolvable_fields(layer, data_keys)
+ references = {
+ reference
+ for expression in _expression_strings(layer)
+ for reference in _DATUM_REFERENCE.findall(expression)
+ }
+ assert references - resolvable == {"valeu"}
+
+
+# --- Update-union discriminator ----------------------------------------------
+
+
[email protected](
+ "config",
+ [
+ {"row_limit": 50},
+ {"groupby": ["region"], "metric": "revenue"},
+ {"color_scheme": "supersetColors"},
+ ],
+)
+def test_update_request_rejects_config_without_discriminator(
+ config: dict[str, Any],
+) -> None:
+ """A config omitting chart_type must not fall through to the Treemap
model."""
+ with pytest.raises(ValidationError):
+ UpdateChartRequest.model_validate({"identifier": 1, "config": config})
+ with pytest.raises(ValidationError):
+ UpdateChartPreviewRequest.model_validate({"dataset_id": 7, "config":
config})
+
+
+def test_update_request_still_accepts_partial_treemap_with_discriminator() ->
None:
+ """Partial Treemap updates remain valid when the discriminator is
present."""
+ request = UpdateChartRequest.model_validate(
+ {"identifier": 1, "config": {"chart_type": "treemap_v2", "row_limit":
50}}
+ )
+ assert request.config is not None
+ assert request.config.chart_type == "treemap_v2"
+ assert request.config.groupby is None
+ assert request.config.metric is None
+ assert "row_limit" in request.config.model_fields_set
+
+
+# --- Adhoc Custom SQL hierarchy columns --------------------------------------
+
+ADHOC_COLUMN: dict[str, Any] = {
+ "label": "Region Bucket",
+ "sqlExpression": "CASE WHEN region = 'West' THEN 'W' ELSE 'O' END",
+ "expressionType": "SQL",
+ "hasCustomLabel": True,
+}
+ADHOC_FORM_DATA: dict[str, Any] = {
+ **FORM_DATA,
+ "currency_format": None,
+ "groupby": ["region", ADHOC_COLUMN],
+}
+ADHOC_ROWS = [
+ {"region": "West", "Region Bucket": "W", "revenue": 30},
+ {"region": "East", "Region Bucket": "O", "revenue": 60},
+]
+
+
+def test_adhoc_hierarchy_column_is_normalized_to_its_output_label() -> None:
+ """Saved Treemaps using Explore's Custom SQL columns must not be
rejected."""
+ result = {"queries": [{"data": ADHOC_ROWS}]}
+ assert normalize_chart_query_result(result, ADHOC_FORM_DATA) is result
+
+
+def test_adhoc_hierarchy_column_without_label_falls_back_to_sql() -> None:
+ """An unlabelled adhoc column is keyed by its raw SQL, matching
getColumnLabel."""
+ form_data = {
+ **ADHOC_FORM_DATA,
+ "groupby": [
+ "region",
+ {"sqlExpression": "lower(city)", "expressionType": "SQL"},
+ ],
+ }
+ rows = [{"region": "West", "lower(city)": "sf", "revenue": 30}]
+ result = {"queries": [{"data": rows}]}
+ assert normalize_chart_query_result(result, form_data) is result
+
+
+def test_adhoc_hierarchy_previews_render_with_resolved_labels() -> None:
+ """Preview formats index rows by output label rather than the adhoc
object."""
+ ascii_preview = treemap_ascii(ADHOC_ROWS, ADHOC_FORM_DATA)
+ assert not isinstance(ascii_preview, ChartError)
+ assert "West > W" in ascii_preview
+ vega = treemap_vega_lite(ADHOC_ROWS, ADHOC_FORM_DATA)
+ assert not isinstance(vega, ChartError)
+ assert vega.specification["usermeta"]["hierarchy"] == ["region", "Region
Bucket"]
+
+
[email protected](
+ "groupby",
+ [
+ ["region", {"expressionType": "SQL"}],
+ ["region", {"label": "region"}],
+ ["region", 7],
+ ],
+)
+def test_unresolvable_or_duplicate_hierarchy_columns_stay_rejected(
+ groupby: list[Any],
+) -> None:
+ """Normalization must not weaken the unique-resolvable-hierarchy
contract."""
+ failure = normalize_chart_query_result(
+ {"queries": [{"data": ADHOC_ROWS}]}, {**ADHOC_FORM_DATA, "groupby":
groupby}
+ )
+ assert isinstance(failure, ChartError)
+ assert failure.error_type == "InvalidTreemapFormData"
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
index 792e7a3fdac..463894cb1c5 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
@@ -1938,6 +1938,161 @@ class TestSavedChartExtraFormDataFilters:
assert payload["queries"][0]["data"] is rows
assert payload["queries"][0]["rowcount"] == source_rowcount
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("data_path", ["saved", "saved_cache",
"unsaved_cache"])
+ @pytest.mark.parametrize("export_format", ["json", "csv", "excel"])
+ @pytest.mark.parametrize("has_finite", [True, False])
+ async def test_treemap_fastmcp_validates_results_and_exports(
+ self,
+ mcp_server: Any,
+ mock_auth: Any,
+ data_path: str,
+ export_format: str,
+ has_finite: bool,
+ ) -> None:
+ """Treemap validates hierarchy metrics on saved and cached export
paths."""
+ from decimal import Decimal
+
+ module = importlib.import_module(
+ "superset.mcp_service.chart.tool.get_chart_data"
+ )
+ chart = SimpleNamespace(
+ id=10,
+ slice_name="SLA",
+ viz_type="treemap_v2",
+ datasource_id=1,
+ datasource_type="table",
+ query_context=json.dumps(
+ {
+ "datasource": {"id": 1, "type": "table"},
+ "queries": [
+ {
+ "columns": ["team"],
+ "metrics": ["saved_sla"],
+ "row_limit": 10,
+ }
+ ],
+ }
+ ),
+ params=json.dumps(
+ {
+ "viz_type": "treemap_v2",
+ "metric": "saved_sla",
+ "groupby": ["team"],
+ }
+ ),
+ )
+
+ def fake_load(self: Any, data: dict[str, Any]) -> Any:
+ return SimpleNamespace(
+ queries=[
+ SimpleNamespace(
+ filter=[],
+ time_range=None,
+ to_dict=lambda: dict(data["queries"][0]),
+ )
+ ],
+ form_data={},
+ )
+
+ rows: list[dict[str, Any]] = [
+ {"team": "Blue", "saved_sla": Decimal("42") if has_finite else
None}
+ ]
+ source_rowcount = len(rows) + 7
+ payload = {
+ "queries": [
+ {
+ "data": rows,
+ "rowcount": source_rowcount,
+ "colnames": ["team", "saved_sla"],
+ }
+ ]
+ }
+
+ class Command:
+ def __init__(self, query_context: Any) -> None: ...
+ def validate(self) -> None: ...
+ def run(self) -> dict[str, Any]:
+ return payload
+
+ cached_form_data = {
+ "viz_type": "treemap_v2",
+ "datasource": "1__table",
+ "metric": "saved_sla",
+ "groupby": ["team"],
+ "slice_name": "SLA",
+ }
+ query = {"columns": ["team"], "metrics": ["saved_sla"]}
+ with (
+ patch.object(
+ module,
+ "get_cached_form_data",
+ return_value=json.dumps(cached_form_data),
+ ),
+ patch.object(
+ module, "build_query_dicts_from_form_data",
return_value=[query]
+ ),
+ patch.object(
+ module,
+ "build_query_context_from_form_data",
+ return_value=fake_load(None, {"queries": [query]}),
+ ),
+ patch.object(module, "find_chart_by_identifier",
return_value=chart),
+ patch.object(
+ module,
+ "validate_chart_dataset",
+ return_value=SimpleNamespace(is_valid=True, warnings=[],
error=None),
+ ),
+ patch(
+ "superset.charts.schemas.ChartDataQueryContextSchema.load",
fake_load
+ ),
+ patch(
+
"superset.commands.chart.data.get_data_command.ChartDataCommand",
+ Command,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ request = {"format": export_format}
+ if data_path != "unsaved_cache":
+ request["identifier"] = "10"
+ if data_path != "saved":
+ request["form_data_key"] = "raw-treemap-cache"
+ result = await client.call_tool("get_chart_data", {"request":
request})
+
+ data = json.loads(result.content[0].text)
+ if not has_finite:
+ assert data["error_type"] == "InvalidTreemapMetric"
+ return
+ assert data["row_count"] == len(rows)
+ assert data["total_rows"] == (
+ source_rowcount if export_format == "json" else len(rows)
+ )
+ expected_groups = [row["team"] for row in rows]
+ if export_format == "json":
+ assert [row["team"] for row in data["data"]] == expected_groups
+ assert data["data_quality"]["completeness"] == pytest.approx(1)
+ assert data.get("query_results") is None
+ elif export_format == "csv":
+ import csv
+ from io import StringIO
+
+ exported = list(csv.DictReader(StringIO(data["csv_data"])))
+ assert [row["team"] for row in exported] == expected_groups
+ else:
+ import base64
+ from io import BytesIO
+
+ from openpyxl import load_workbook
+
+ workbook =
load_workbook(BytesIO(base64.b64decode(data["excel_data"])))
+ assert list(workbook.active.values)[0] == ("team", "saved_sla")
+ assert [
+ row[0] for row in list(workbook.active.values)[1:]
+ ] == expected_groups
+ assert [row[1] for row in list(workbook.active.values)[1:]] == [42]
+ assert payload["queries"][0]["data"] is rows
+ assert payload["queries"][0]["rowcount"] == source_rowcount
+
class TestOAuthErrorRouting:
"""Query-time OAuth errors must reach the dedicated OAuth handlers.
diff --git
a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_type_schema.py
b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_type_schema.py
index 8a2ab50f6bf..27b22ae76ce 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_type_schema.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_type_schema.py
@@ -17,12 +17,15 @@
"""Tests for get_chart_type_schema tool logic."""
+from typing import Any
+
import pytest
from superset.extensions import feature_flag_manager
from superset.mcp_service.chart.tool.get_chart_type_schema import (
_CHART_EXAMPLES,
_CHART_TYPE_ADAPTERS,
+ _compiled_chart_schema,
_get_chart_type_schema_impl as _call_schema,
VALID_CHART_TYPES,
)
@@ -148,3 +151,26 @@ class TestGetChartTypeSchema:
for chart_type in VALID_CHART_TYPES:
assert chart_type in _CHART_EXAMPLES
assert len(_CHART_EXAMPLES[chart_type]) >= 1
+
+
+def test_schema_cache_compiles_once_and_isolates_response_mutation(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """Repeated discovery avoids compilation without sharing mutable response
data."""
+ _compiled_chart_schema.cache_clear()
+ adapter = _CHART_TYPE_ADAPTERS["treemap_v2"]
+ original = adapter.json_schema
+ calls = 0
+
+ def counted_schema() -> dict[str, Any]:
+ nonlocal calls
+ calls += 1
+ return original()
+
+ monkeypatch.setattr(adapter, "json_schema", counted_schema)
+ first = _call_schema("treemap_v2")
+ first["schema"]["properties"].clear()
+ second = _call_schema("treemap_v2")
+ assert "groupby" in second["schema"]["properties"]
+ assert calls == 1
+ _compiled_chart_schema.cache_clear()
diff --git
a/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py
b/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py
index dba58131399..3b8b79b88d7 100644
--- a/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py
+++ b/tests/unit_tests/mcp_service/explore/tool/test_generate_explore_link.py
@@ -97,6 +97,55 @@ async def
test_gauge_fastmcp_entry_compiles_and_returns_native_form_data(
assert mock_validate.call_args.kwargs["run_compile_check"] is True
[email protected]("time_range", ["", " ", "No filter"])
[email protected]("comparator", ["Last week", "", " "])
[email protected](generate_explore_link_module, "validate_and_compile")
+@patch("superset.daos.dataset.DatasetDAO.find_by_id")
[email protected]
+async def test_treemap_fastmcp_entry_compiles_native_temporal_form_data(
+ mock_find_dataset, mock_validate, mcp_server, time_range: str, comparator:
str
+) -> None:
+ """The public Treemap request reaches compile and a native Explore
payload."""
+ from superset.mcp_service.chart.compile import CompileResult
+
+ mock_find_dataset.return_value = _mock_dataset(id=3)
+ mock_validate.return_value = CompileResult(success=True)
+ request = {
+ "dataset_id": "3",
+ "config": {
+ "chart_type": "treemap_v2",
+ "metric": {"name": "num", "aggregate": "AVG"},
+ "groupby": ["name"],
+ "show_labels": False,
+ "number_format": ",.1f",
+ "time_range": time_range,
+ "adhoc_filters": [
+ {
+ "subject": "event_time",
+ "operator": "TEMPORAL_RANGE",
+ "comparator": comparator,
+ }
+ ],
+ },
+ }
+
+ async with Client(mcp_server) as client:
+ result = await client.call_tool("generate_explore_link", {"request":
request})
+
+ assert result.structured_content["success"] is True
+ form_data = result.structured_content["form_data"]
+ assert form_data["viz_type"] == "treemap_v2"
+ assert form_data["metric"]["label"] == "AVG(num)"
+ assert form_data["number_format"] == ",.1f"
+ if comparator == "Last week":
+ assert form_data["time_range"] == "Last week"
+ else:
+ assert form_data.get("time_range") in (None, "No filter")
+ assert form_data["adhoc_filters"][0]["subject"] == "event_time"
+ assert form_data["adhoc_filters"][0]["comparator"] == "No filter"
+ assert mock_validate.call_args.kwargs["run_compile_check"] is True
+
+
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
