sadpandajoe commented on code in PR #44152:
URL: https://github.com/apache/superset/pull/44152#discussion_r4041438764
##########
superset/mcp_service/chart/schemas.py:
##########
@@ -2925,7 +2991,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(
Review Comment:
Adding the defaulted partial Treemap model to this union makes a config
without `chart_type` validate as a Treemap. On `update_chart(...,
generate_preview=false)`, omitting the discriminator for an existing pie or
table chart can then overwrite its `viz_type` and params instead of rejecting
the request. Could `TreemapChartUpdateConfig.chart_type` be required here, or
could this remain a discriminated update union?
##########
superset/mcp_service/chart/query_result.py:
##########
@@ -223,3 +225,83 @@ 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 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 = form_data.get("groupby")
+ if (
+ not label
+ or not isinstance(hierarchy, list)
+ or not hierarchy
+ or not all(isinstance(column, str) and column for column in hierarchy)
Review Comment:
Saved Treemaps can contain an adhoc Custom SQL object in `groupby` via
Explore's `DndColumnSelect`, but this check accepts only strings.
`get_chart_data`, preview, and export now reject such otherwise-valid saved
charts with `InvalidTreemapFormData`. Could this normalize adhoc columns to
their output labels instead of rejecting them?
##########
.github/workflows/superset-python-unittest.yml:
##########
@@ -43,7 +43,10 @@ jobs:
needs: changes
if: needs.changes.outputs.python == 'true'
runs-on: ubuntu-26.04
- timeout-minutes: 30
+ # The full coverage run is followed by independent SQL and semantic-layer
+ # 100% gates. Allow all three coverage-preserving phases to finish under
+ # ordinary runner variance rather than cancelling after successful tests.
+ timeout-minutes: 40
Review Comment:
This hunk now conflicts with `master`, where #44178 runs this suite with
pytest-xdist and gives it 45 minutes. Resolving the conflict in favor of this
40-minute version would undo that newer guard and can time out
otherwise-passing required checks. Could you drop this workflow change when
rebasing?
##########
tests/unit_tests/mcp_service/chart/test_treemap_completeness.py:
##########
@@ -0,0 +1,1174 @@
+# 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."""
+
+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"):
Review Comment:
This is the only new test that actually compiles and renders the Vega-Lite
spec, but the Python-Unit workflow neither installs these Node packages nor
sets `NODE_PATH`, so both parametrizations skip in CI. A bad `calculate` field
or Vega expression could therefore pass the remaining dict-shape tests and fail
only at render time. Could CI run this compiler/scenegraph assertion, or an
equivalent schema/compile check?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
