This is an automated email from the ASF dual-hosted git repository.
richardfogaca 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 51e708da6bb fix(mcp): accept Superset vocabulary in chart configs and
clarify query_dataset metric errors (#40972)
51e708da6bb is described below
commit 51e708da6bb52ae4838613e5c3a1d851d725f0dc
Author: Richard Fogaca Nienkotter
<[email protected]>
AuthorDate: Mon Jul 6 16:01:14 2026 -0300
fix(mcp): accept Superset vocabulary in chart configs and clarify
query_dataset metric errors (#40972)
---
superset/mcp_service/chart/schemas.py | 99 +++++++-
superset/mcp_service/dataset/tool/query_dataset.py | 46 +++-
.../chart/test_chart_config_coercion.py | 281 +++++++++++++++++++++
3 files changed, 415 insertions(+), 11 deletions(-)
diff --git a/superset/mcp_service/chart/schemas.py
b/superset/mcp_service/chart/schemas.py
index dc9414e08c8..469e28ba1f8 100644
--- a/superset/mcp_service/chart/schemas.py
+++ b/superset/mcp_service/chart/schemas.py
@@ -1678,7 +1678,9 @@ class XYChartConfig(UnknownFieldCheckMixin):
validation_alias=AliasChoices("time_grain", "time_grain_sqla"),
)
orientation: Literal["vertical", "horizontal"] | None = Field(
- None, description="Bar orientation (only for kind='bar')"
+ None,
+ description="Bar orientation (only for kind='bar')",
+ validation_alias=AliasChoices("orientation", "chart_orientation"),
)
stacked: bool = Field(
False,
@@ -1693,7 +1695,10 @@ class XYChartConfig(UnknownFieldCheckMixin):
)
x_axis: AxisConfig | None = None
y_axis: AxisConfig | None = None
- legend: LegendConfig | None = None
+ legend: LegendConfig | None = Field(
+ None,
+ validation_alias=AliasChoices("legend", "show_legend"),
+ )
x_axis_time_format: str | None = Field(
None,
description=(
@@ -1738,6 +1743,24 @@ class XYChartConfig(UnknownFieldCheckMixin):
def wrap_single_group_by(cls, v: Any) -> Any:
return _normalize_group_by_input(v)
+ @field_validator("x", mode="before")
+ @classmethod
+ def coerce_x_column_name(cls, v: Any) -> Any:
+ """Accept a bare column name string for the x-axis."""
+ return {"name": v} if isinstance(v, str) else v
+
+ @field_validator("y", mode="before")
+ @classmethod
+ def coerce_y_column_names(cls, v: Any) -> Any:
+ """Accept bare strings or a single entry for the y-axis metrics."""
+ return _normalize_group_by_input(v)
+
+ @field_validator("legend", mode="before")
+ @classmethod
+ def coerce_legend_flag(cls, v: Any) -> Any:
+ """Accept the form_data-style show_legend boolean."""
+ return {"show": v} if isinstance(v, bool) else v
+
@model_validator(mode="after")
def reject_sql_expression_on_dimensions(self) -> "XYChartConfig":
"""sql_expression is metric-only; reject it on x and group_by."""
@@ -1825,6 +1848,70 @@ ChartConfig = Annotated[
# giving LLMs enough context to construct valid configs.
+# Superset viz_type values that LLM clients routinely send where this API
+# expects its chart_type discriminator. Extend this observed-pattern map when
+# recurring session traces show additional unambiguous aliases. Each maps to
+# (chart_type, kind); kind is only meaningful for the xy family.
+_VIZ_TYPE_TO_CHART_TYPE: dict[str, tuple[str, str | None]] = {
+ "bar": ("xy", "bar"),
+ "dist_bar": ("xy", "bar"),
+ "echarts_timeseries_bar": ("xy", "bar"),
+ "line": ("xy", "line"),
+ "echarts_timeseries_line": ("xy", "line"),
+ "echarts_timeseries_smooth": ("xy", "line"),
+ "area": ("xy", "area"),
+ "echarts_area": ("xy", "area"),
+ "scatter": ("xy", "scatter"),
+ "echarts_timeseries_scatter": ("xy", "scatter"),
+ "ag-grid-table": ("table", None),
+ "big_number_total": ("big_number", None),
+ "pivot_table_v2": ("pivot_table", None),
+}
+
+
+def _normalize_chart_request_input(data: Any) -> Any:
+ """Accept common Superset REST/form_data vocabulary in chart requests.
+
+ LLM clients reliably reach for Superset's public field names —
+ ``datasource_id``, ``viz_type``, and concrete viz plugin names such as
+ ``echarts_timeseries_bar`` — before consulting this API's schema. Each
+ rejection costs the client a model round trip, so the unambiguous
+ synonyms are translated instead of refused.
+ """
+ if not isinstance(data, dict):
+ return data
+ if "dataset_id" not in data and "datasource_id" in data:
+ data["dataset_id"] = data.pop("datasource_id")
+ config = data.get("config")
+ if isinstance(config, dict):
+ viz_type = config.get("viz_type")
+ if isinstance(viz_type, str) and "chart_type" not in config:
+ config["chart_type"] = viz_type
+ chart_type = config.get("chart_type")
+ if isinstance(chart_type, str) and chart_type in
_VIZ_TYPE_TO_CHART_TYPE:
+ mapped_type, kind = _VIZ_TYPE_TO_CHART_TYPE[chart_type]
+ config["chart_type"] = mapped_type
+ if kind is not None:
+ config.setdefault("kind", kind)
+ if mapped_type == "table":
+ config.setdefault("viz_type", chart_type)
+ else:
+ config.pop("viz_type", None)
+ elif config.get("chart_type") != "table":
+ config.pop("viz_type", None)
+ return data
+
+
+class ChartRequestNormalizerMixin(BaseModel):
+ """Mixin translating Superset-vocabulary synonyms before validation."""
+
+ @model_validator(mode="before")
+ @classmethod
+ def _normalize_request_vocabulary(cls, data: Any) -> Any:
+ """Normalize Superset-style request keys before Pydantic validation."""
+ return _normalize_chart_request_input(data)
+
+
class ListChartsRequest(OwnedByMeMixin, CreatedByMeMixin,
MetadataCacheControl):
"""Request schema for list_charts with clear, unambiguous types."""
@@ -1919,7 +2006,7 @@ class ListChartsRequest(OwnedByMeMixin, CreatedByMeMixin,
MetadataCacheControl):
# The tool input models
-class GenerateChartRequest(QueryCacheControl):
+class GenerateChartRequest(ChartRequestNormalizerMixin, QueryCacheControl):
model_config = ConfigDict(populate_by_name=True)
dataset_id: int | str = Field(..., description="Dataset identifier (ID,
UUID)")
@@ -2019,7 +2106,7 @@ class GenerateChartRequest(QueryCacheControl):
return self
-class GenerateExploreLinkRequest(FormDataCacheControl):
+class GenerateExploreLinkRequest(ChartRequestNormalizerMixin,
FormDataCacheControl):
dataset_id: int | str = Field(..., description="Dataset identifier (ID,
UUID)")
config: ChartConfig | None = Field(
None,
@@ -2031,7 +2118,7 @@ class GenerateExploreLinkRequest(FormDataCacheControl):
)
-class UpdateChartRequest(QueryCacheControl):
+class UpdateChartRequest(ChartRequestNormalizerMixin, QueryCacheControl):
model_config = ConfigDict(populate_by_name=True)
identifier: int | str = Field(..., description="Chart ID or UUID")
@@ -2078,7 +2165,7 @@ class UpdateChartRequest(QueryCacheControl):
return sanitize_user_input(v, "Chart name", max_length=255,
allow_empty=True)
-class UpdateChartPreviewRequest(FormDataCacheControl):
+class UpdateChartPreviewRequest(ChartRequestNormalizerMixin,
FormDataCacheControl):
form_data_key: str | None = Field(
None,
description=(
diff --git a/superset/mcp_service/dataset/tool/query_dataset.py
b/superset/mcp_service/dataset/tool/query_dataset.py
index da7f2d227f6..c87c7a0945e 100644
--- a/superset/mcp_service/dataset/tool/query_dataset.py
+++ b/superset/mcp_service/dataset/tool/query_dataset.py
@@ -84,22 +84,48 @@ def _validate_names(
requested: list[str],
valid: set[str],
kind: str,
+ *,
+ empty_hint: str | None = None,
+ list_valid_on_miss: bool = False,
) -> list[str]:
"""Return list of error messages for names not found in *valid*.
- Includes close-match suggestions when available.
+ Includes close-match suggestions when available. When *valid* is empty,
+ appends *empty_hint* instead of a useless fuzzy match. When no close
+ match exists and *list_valid_on_miss* is set, lists the valid names so
+ the caller does not have to guess again.
"""
errors: list[str] = []
for name in requested:
if name not in valid:
- suggestions = difflib.get_close_matches(name, valid, n=3,
cutoff=0.6)
msg = f"Unknown {kind}: '{name}'"
- if suggestions:
- msg += f". Did you mean: {', '.join(suggestions)}?"
+ if not valid:
+ if empty_hint:
+ msg += f". {empty_hint}"
+ else:
+ suggestions = difflib.get_close_matches(name, valid, n=3,
cutoff=0.6)
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ elif list_valid_on_miss:
+ shown = sorted(valid)[:10]
+ more = len(valid) - len(shown)
+ suffix = (
+ f" (and {more} more; call get_dataset_info for the
full list)"
+ if more > 0
+ else ""
+ )
+ msg += f". Valid {kind}s: {', '.join(shown)}{suffix}"
errors.append(msg)
return errors
+_NO_SAVED_METRICS_HINT = (
+ "This dataset has no saved metrics, and query_dataset accepts saved "
+ "metric names only (not ad-hoc expressions). Use execute_sql for "
+ "ad-hoc aggregates, or add a saved metric to the dataset."
+)
+
+
@requires_data_model_metadata_access
@tool(
tags=["data"],
@@ -119,6 +145,10 @@ async def query_dataset( # noqa: C901
to compute saved metrics, group by dimensions, or apply filters directly
against a dataset's curated semantic layer.
+ Metrics must be saved metric names from get_dataset_info; ad-hoc
+ expressions such as "SUM(col)" are not accepted. When the dataset has no
+ saved metric for the aggregate you need, use execute_sql instead.
+
Workflow:
1. list_datasets -> find a dataset
2. get_dataset_info -> discover available columns and metrics
@@ -213,7 +243,13 @@ async def query_dataset( # noqa: C901
_validate_names(request.columns, valid_columns, "column")
)
validation_errors.extend(
- _validate_names(request.metrics, valid_metrics, "metric")
+ _validate_names(
+ request.metrics,
+ valid_metrics,
+ "metric",
+ empty_hint=_NO_SAVED_METRICS_HINT,
+ list_valid_on_miss=True,
+ )
)
# Validate filter column names against dataset columns
filter_cols = [f.col for f in request.filters]
diff --git a/tests/unit_tests/mcp_service/chart/test_chart_config_coercion.py
b/tests/unit_tests/mcp_service/chart/test_chart_config_coercion.py
new file mode 100644
index 00000000000..bc1c8696b9c
--- /dev/null
+++ b/tests/unit_tests/mcp_service/chart/test_chart_config_coercion.py
@@ -0,0 +1,281 @@
+# 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.
+
+"""
+Unit tests for Superset-vocabulary coercion in MCP chart request schemas.
+
+LLM clients reliably send Superset's public field names (datasource_id,
+viz_type, show_legend, plugin viz names) before consulting the MCP tool
+schema. These tests pin the translations that absorb those inputs instead
+of rejecting them.
+"""
+
+import pytest
+from pydantic import ValidationError
+
+from superset.mcp_service.chart.schemas import (
+ GenerateChartRequest,
+ GenerateExploreLinkRequest,
+ TableChartConfig,
+ UpdateChartPreviewRequest,
+ UpdateChartRequest,
+ XYChartConfig,
+)
+from superset.mcp_service.dataset.tool.query_dataset import (
+ _NO_SAVED_METRICS_HINT,
+ _validate_names,
+)
+
+XY_BAR_CONFIG = {
+ "chart_type": "xy",
+ "kind": "bar",
+ "x": {"name": "genre"},
+ "y": [{"name": "global_sales", "aggregate": "SUM"}],
+}
+
+
+class TestDatasourceIdAlias:
+ """Datasource alias normalization for chart request models."""
+
+ def test_generate_chart_accepts_datasource_id(self) -> None:
+ request = GenerateChartRequest.model_validate(
+ {"datasource_id": 23, "config": dict(XY_BAR_CONFIG)}
+ )
+ assert request.dataset_id == 23
+
+ def test_dataset_id_wins_over_datasource_id(self) -> None:
+ request = GenerateChartRequest.model_validate(
+ {"dataset_id": 7, "datasource_id": 23, "config":
dict(XY_BAR_CONFIG)}
+ )
+ assert request.dataset_id == 7
+
+
+class TestVizTypeTranslation:
+ """Superset viz plugin names map to MCP chart config discriminators."""
+
+ @pytest.mark.parametrize(
+ "viz_type,expected_kind",
+ [
+ ("bar", "bar"),
+ ("dist_bar", "bar"),
+ ("echarts_timeseries_bar", "bar"),
+ ("echarts_timeseries_line", "line"),
+ ("area", "area"),
+ ("echarts_timeseries_scatter", "scatter"),
+ ],
+ )
+ def test_xy_family_viz_types_map_to_xy(
+ self, viz_type: str, expected_kind: str
+ ) -> None:
+ config = dict(XY_BAR_CONFIG)
+ del config["kind"]
+ config["chart_type"] = viz_type
+ request = GenerateChartRequest.model_validate(
+ {"dataset_id": 23, "config": config}
+ )
+ assert request.config.chart_type == "xy"
+ assert request.config.kind == expected_kind
+
+ def test_viz_type_key_is_accepted_as_chart_type(self) -> None:
+ config = dict(XY_BAR_CONFIG)
+ del config["chart_type"]
+ del config["kind"]
+ config["viz_type"] = "bar"
+ request = GenerateChartRequest.model_validate(
+ {"dataset_id": 23, "config": config}
+ )
+ assert request.config.chart_type == "xy"
+ assert request.config.kind == "bar"
+
+ def test_explicit_chart_type_wins_over_viz_type(self) -> None:
+ config = dict(XY_BAR_CONFIG)
+ config["viz_type"] = "pie"
+ request = GenerateChartRequest.model_validate(
+ {"dataset_id": 23, "config": config}
+ )
+ assert request.config.chart_type == "xy"
+
+ def test_table_viz_type_is_preserved_with_explicit_chart_type(self) ->
None:
+ request = GenerateExploreLinkRequest.model_validate(
+ {
+ "dataset_id": 23,
+ "config": {
+ "chart_type": "table",
+ "viz_type": "ag-grid-table",
+ "columns": [{"name": "genre"}],
+ },
+ }
+ )
+ assert isinstance(request.config, TableChartConfig)
+ assert request.config.viz_type == "ag-grid-table"
+
+ def test_table_viz_type_key_maps_to_table_config(self) -> None:
+ request = GenerateExploreLinkRequest.model_validate(
+ {
+ "dataset_id": 23,
+ "config": {
+ "viz_type": "ag-grid-table",
+ "columns": [{"name": "genre"}],
+ },
+ }
+ )
+ assert isinstance(request.config, TableChartConfig)
+ assert request.config.chart_type == "table"
+ assert request.config.viz_type == "ag-grid-table"
+
+ def test_explicit_invalid_chart_type_wins_over_viz_type(self) -> None:
+ with pytest.raises(ValidationError, match="xyy"):
+ GenerateChartRequest.model_validate(
+ {
+ "dataset_id": 23,
+ "config": {
+ **XY_BAR_CONFIG,
+ "chart_type": "xyy",
+ "viz_type": "echarts_timeseries_bar",
+ },
+ }
+ )
+
+ def test_explicit_kind_is_preserved(self) -> None:
+ config = dict(XY_BAR_CONFIG)
+ config["chart_type"] = "echarts_timeseries_bar"
+ config["kind"] = "line"
+ request = GenerateChartRequest.model_validate(
+ {"dataset_id": 23, "config": config}
+ )
+ assert request.config.kind == "line"
+
+ def test_big_number_total_maps_to_big_number(self) -> None:
+ request = GenerateChartRequest.model_validate(
+ {
+ "dataset_id": 23,
+ "config": {
+ "chart_type": "big_number_total",
+ "metric": {"name": "global_sales", "aggregate": "SUM"},
+ },
+ }
+ )
+ assert request.config.chart_type == "big_number"
+
+
+class TestRequestModelsShareNormalization:
+ """Request models share chart vocabulary normalization behavior."""
+
+ def test_update_chart_request(self) -> None:
+ request = UpdateChartRequest.model_validate(
+ {"identifier": 409, "config": {**XY_BAR_CONFIG, "chart_type":
"bar"}}
+ )
+ assert request.config is not None
+ assert request.config.chart_type == "xy"
+
+ def test_update_chart_preview_request(self) -> None:
+ request = UpdateChartPreviewRequest.model_validate(
+ {"datasource_id": 23, "config": {**XY_BAR_CONFIG, "chart_type":
"bar"}}
+ )
+ assert request.dataset_id == 23
+ assert request.config.chart_type == "xy"
+
+ def test_generate_explore_link_request(self) -> None:
+ request = GenerateExploreLinkRequest.model_validate(
+ {"datasource_id": 23, "config": {**XY_BAR_CONFIG, "chart_type":
"bar"}}
+ )
+ assert request.dataset_id == 23
+ assert request.config is not None
+ assert request.config.chart_type == "xy"
+
+
+class TestXYFieldCoercion:
+ """XY chart config accepts common Superset form-data field shapes."""
+
+ def test_x_accepts_bare_column_name(self) -> None:
+ config = XYChartConfig.model_validate({**XY_BAR_CONFIG, "x": "genre"})
+ assert config.x is not None
+ assert config.x.name == "genre"
+
+ def test_y_accepts_bare_metric_names(self) -> None:
+ config = XYChartConfig.model_validate({**XY_BAR_CONFIG, "y":
["count"]})
+ assert config.y[0].name == "count"
+
+ def test_show_legend_bool_becomes_legend_config(self) -> None:
+ config = XYChartConfig.model_validate({**XY_BAR_CONFIG, "show_legend":
True})
+ assert config.legend is not None
+ assert config.legend.show is True
+
+ def test_show_legend_false(self) -> None:
+ config = XYChartConfig.model_validate({**XY_BAR_CONFIG, "show_legend":
False})
+ assert config.legend is not None
+ assert config.legend.show is False
+
+ def test_chart_orientation_alias(self) -> None:
+ config = XYChartConfig.model_validate(
+ {**XY_BAR_CONFIG, "chart_orientation": "horizontal"}
+ )
+ assert config.orientation == "horizontal"
+
+ def test_unknown_fields_still_rejected_with_suggestion(self) -> None:
+ with pytest.raises(ValidationError, match="order_desc"):
+ XYChartConfig.model_validate({**XY_BAR_CONFIG, "order_desc": True})
+
+
+class TestQueryDatasetMetricGuidance:
+ """Metric validation errors guide callers toward saved metric names."""
+
+ def test_no_saved_metrics_hint(self) -> None:
+ errors = _validate_names(
+ ["sum__global_sales"],
+ set(),
+ "metric",
+ empty_hint=_NO_SAVED_METRICS_HINT,
+ )
+ assert len(errors) == 1
+ assert _NO_SAVED_METRICS_HINT in errors[0]
+ assert "Did you mean" not in errors[0]
+
+ def test_valid_metrics_listed_when_no_close_match(self) -> None:
+ errors = _validate_names(
+ ["zzz_nonexistent"],
+ {"count", "revenue_total"},
+ "metric",
+ list_valid_on_miss=True,
+ )
+ assert len(errors) == 1
+ assert "Valid metrics: count, revenue_total" in errors[0]
+
+ def test_truncated_valid_metrics_report_remaining_count(self) -> None:
+ valid_metrics = {f"metric_{idx:02d}" for idx in range(12)}
+
+ errors = _validate_names(
+ ["zzz_nonexistent"],
+ valid_metrics,
+ "metric",
+ list_valid_on_miss=True,
+ )
+
+ assert len(errors) == 1
+ assert "metric_00, metric_01" in errors[0]
+ assert "metric_10" not in errors[0]
+ assert "and 2 more; call get_dataset_info for the full list" in
errors[0]
+
+ def test_close_match_suggestion_unchanged(self) -> None:
+ errors = _validate_names(
+ ["revenue_totl"],
+ {"count", "revenue_total"},
+ "metric",
+ list_valid_on_miss=True,
+ )
+ assert len(errors) == 1
+ assert "Did you mean: revenue_total?" in errors[0]
