You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/09/24 09:14:43 UTC
(superset) 01/01: fix(mcp): return actionable authorized column suggestions
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch sc-121713-column-suggestions-placeholder
in repository https://gitbox.apache.org/repos/asf/superset.git
commit b354b1ef0b523d768712f80781f9c248cc1a7a10
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Thu Sep 24 09:14:26 2026 +0000
fix(mcp): return actionable authorized column suggestions
---
docs/admin_docs/configuration/mcp-server.mdx | 9 +
.../chart/validation/dataset_validator.py | 68 +++---
superset/mcp_service/utils/error_builder.py | 7 +-
tests/unit_tests/mcp_service/chart/test_compile.py | 16 +-
.../chart/tool/test_column_suggestions.py | 229 +++++++++++++++++++++
5 files changed, 296 insertions(+), 33 deletions(-)
diff --git a/docs/admin_docs/configuration/mcp-server.mdx b/docs/admin_docs/configuration/mcp-server.mdx
index a8bb4d6a92c..9ace2d79685 100644
--- a/docs/admin_docs/configuration/mcp-server.mdx
+++ b/docs/admin_docs/configuration/mcp-server.mdx
@@ -1361,3 +1361,12 @@ predicate, whose behavior varied by database. It is not an indexed
full-text search: large catalogs and long descriptions can increase scan time.
Benchmark with your metadata database and catalog size when tuning discovery
latency.
+
+### Missing chart columns
+
+When a chart references a missing column, `generate_chart` returns a structured
+`column_not_found` error with up to three similar physical column names. If no
+similar columns exist, it returns explicit no-match guidance instead. Column
+errors include a bounded `dataset_context` with up to ten available column names
+from the authorized dataset; names are escaped and length-limited. Saved-metric
+expressions are not included. Use `get_dataset_info` for the complete schema.
diff --git a/superset/mcp_service/chart/validation/dataset_validator.py b/superset/mcp_service/chart/validation/dataset_validator.py
index d7cb7158b56..376c5864330 100644
--- a/superset/mcp_service/chart/validation/dataset_validator.py
+++ b/superset/mcp_service/chart/validation/dataset_validator.py
@@ -435,6 +435,7 @@ class DatasetValidator:
"""Fetch the ORM dataset by ID/UUID and build a :class:`DatasetContext`."""
try:
from superset.daos.dataset import DatasetDAO
+ from superset.mcp_service.auth import has_dataset_access
if isinstance(dataset_id, int) or (
isinstance(dataset_id, str) and dataset_id.isdigit()
@@ -443,6 +444,9 @@ class DatasetValidator:
else:
dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
+ if dataset is None or not has_dataset_access(dataset):
+ return None
+
return build_dataset_context_from_orm(dataset)
except Exception as e:
@@ -646,9 +650,6 @@ class DatasetValidator:
for col in dataset_context.available_columns:
all_names.append((col["name"], "column", col.get("type", "UNKNOWN")))
- for metric in dataset_context.available_metrics:
- all_names.append((metric["name"], "metric", "METRIC"))
-
# Find close matches
column_lower = column_name.lower()
candidate_lookup = [name[0].lower() for name in all_names]
@@ -690,34 +691,43 @@ class DatasetValidator:
)
if len(invalid_columns) == 1:
- col = invalid_columns[0]
- col_name = col.name or "<unknown column>"
- suggestions = suggestions_map.get(col_name, [])
-
- if suggestions:
- return ChartErrorBuilder.column_not_found_error(
- col_name, [s.name for s in suggestions]
- )
- else:
- return ChartErrorBuilder.column_not_found_error(col_name)
+ col_name = invalid_columns[0].name or "<unknown column>"
+ error = ChartErrorBuilder.column_not_found_error(
+ col_name, [s.name for s in suggestions_map.get(col_name, [])]
+ )
else:
- # Multiple invalid columns
- invalid_names: list[str] = [col.name for col in invalid_columns if col.name]
- return ChartErrorBuilder.build_error(
- error_type="multiple_invalid_columns",
- template_key="column_not_found",
- template_vars={
- "column": ", ".join(invalid_names[:3])
- + ("..." if len(invalid_names) > 3 else ""),
- "suggestions": "Use get_dataset_info to see all available columns",
- },
- custom_suggestions=[
- f"Invalid columns: {', '.join(invalid_names)}",
- "Check spelling and case sensitivity",
- "Use get_dataset_info to list available columns",
- ],
- error_code="MULTIPLE_INVALID_COLUMNS",
+ candidates = list(
+ dict.fromkeys(
+ suggestion.name
+ for suggestions in suggestions_map.values()
+ for suggestion in suggestions
+ )
+ )
+ error = ChartErrorBuilder.column_not_found_error(
+ "multiple requested columns", candidates
)
+ error.error_type = "multiple_invalid_columns"
+ error.error_code = "MULTIPLE_INVALID_COLUMNS"
+
+ # Return names only, not SQL expressions or unbounded dataset metadata.
+ # Reuse the error builder's escaping and per-value length limit.
+ from superset.mcp_service.utils.error_builder import _sanitize_user_input
+
+ error.dataset_context = DatasetContext(
+ id=dataset_context.id,
+ table_name=_sanitize_user_input(dataset_context.table_name),
+ schema=(
+ _sanitize_user_input(dataset_context.schema_name)
+ if dataset_context.schema_name is not None
+ else None
+ ),
+ database_name=_sanitize_user_input(dataset_context.database_name),
+ available_columns=[
+ {"name": _sanitize_user_input(col["name"])}
+ for col in dataset_context.available_columns[:10]
+ ],
+ )
+ return error
@staticmethod
def _validate_saved_metrics(
diff --git a/superset/mcp_service/utils/error_builder.py b/superset/mcp_service/utils/error_builder.py
index f019092dcc6..2c991b2bfc3 100644
--- a/superset/mcp_service/utils/error_builder.py
+++ b/superset/mcp_service/utils/error_builder.py
@@ -134,7 +134,7 @@ class ChartErrorBuilder:
"suggestions": [
"Check column name spelling and case sensitivity",
"Use get_dataset_info to see available columns",
- "Did you mean: {suggestions}?",
+ "{suggestions}",
],
},
# Runtime errors
@@ -363,7 +363,10 @@ class ChartErrorBuilder:
) -> ChartGenerationError:
"""Build a column not found error."""
suggestion_text = (
- ", ".join(suggestions[:3]) if suggestions else "Check available columns"
+ f"Did you mean: {', '.join(suggestions[:3])}?"
+ if suggestions
+ else "No matching columns found. "
+ "Use get_dataset_info to see available columns."
)
return cls.build_error(
error_type="column_not_found",
diff --git a/tests/unit_tests/mcp_service/chart/test_compile.py b/tests/unit_tests/mcp_service/chart/test_compile.py
index 16eeeb8a673..4db5e7c23be 100644
--- a/tests/unit_tests/mcp_service/chart/test_compile.py
+++ b/tests/unit_tests/mcp_service/chart/test_compile.py
@@ -127,7 +127,13 @@ class TestValidateAndCompileChartTypeCoverage:
assert not result.success
assert result.tier == "validation"
assert result.error_obj is not None
- assert any("sum_boys" in s for s in (result.error_obj.suggestions or []))
+ assert result.error_obj.error_type == "column_not_found"
+ assert result.error_obj.suggestions[-1].startswith("No matching columns found.")
+ assert all("sum_boys" not in s for s in result.error_obj.suggestions)
+ assert result.error_obj.dataset_context is not None
+ assert {
+ c["name"] for c in result.error_obj.dataset_context.available_columns
+ } == {"ds", "gender", "name", "num"}
def test_pie_bad_metric_column_rejected(self):
ds = _orm_dataset()
@@ -139,7 +145,13 @@ class TestValidateAndCompileChartTypeCoverage:
assert not result.success, "Pie chart with bad metric column should fail"
assert result.tier == "validation"
assert result.error_obj is not None
- assert any("sum_boys" in s for s in (result.error_obj.suggestions or []))
+ assert result.error_obj.error_type == "column_not_found"
+ assert result.error_obj.suggestions[-1].startswith("No matching columns found.")
+ assert all("sum_boys" not in s for s in result.error_obj.suggestions)
+ assert result.error_obj.dataset_context is not None
+ assert {
+ c["name"] for c in result.error_obj.dataset_context.available_columns
+ } == {"ds", "gender", "name", "num"}
def test_pie_valid_dimension_and_saved_metric_passes(self):
ds = _orm_dataset()
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_column_suggestions.py b/tests/unit_tests/mcp_service/chart/tool/test_column_suggestions.py
new file mode 100644
index 00000000000..1d3c0453fec
--- /dev/null
+++ b/tests/unit_tests/mcp_service/chart/tool/test_column_suggestions.py
@@ -0,0 +1,229 @@
+# 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.
+
+"""Column-error guidance uses only bounded, authorized dataset metadata."""
+
+from contextlib import nullcontext
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from superset.mcp_service.chart.schemas import GenerateChartRequest
+from superset.mcp_service.chart.tool.generate_chart import generate_chart
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("column", "names", "allowed", "expected_type", "guidance"),
+ [
+ (
+ "no_such_col",
+ ["month", "category", "revenue"],
+ True,
+ "column_not_found",
+ "No matching columns found",
+ ),
+ (
+ "revnue",
+ ["month", "category", "revenue"],
+ True,
+ "column_not_found",
+ "Did you mean: revenue?",
+ ),
+ (
+ "no_such_col",
+ [],
+ True,
+ "multiple_invalid_columns",
+ "No matching columns found",
+ ),
+ (
+ "no_such_col",
+ ["month", "category", "revenue"],
+ False,
+ "dataset_not_found",
+ None,
+ ),
+ (
+ "gross_margi",
+ ["month", "category", "revenue"],
+ True,
+ "column_not_found",
+ "No matching columns found",
+ ),
+ ],
+)
+async def test_generate_chart_column_guidance(
+ column: str,
+ names: list[str],
+ allowed: bool,
+ expected_type: str,
+ guidance: str | None,
+) -> None:
+ """Exercise the reported unsaved bar request through the real pipeline."""
+ dataset = Mock(
+ id=268,
+ table_name="sales_fixture",
+ schema=None,
+ database=Mock(database_name="fixture", db_engine_spec=None),
+ columns=[
+ Mock(
+ column_name=name,
+ type="VARCHAR",
+ is_temporal=False,
+ is_numeric=name == "revenue",
+ )
+ for name in names
+ ],
+ metrics=[
+ Mock(
+ metric_name="gross_margin",
+ expression="SUM(revenue)",
+ description="not a physical column",
+ )
+ ],
+ )
+ request = GenerateChartRequest.model_validate(
+ {
+ "dataset_id": 268,
+ "save_chart": False,
+ "config": {
+ "chart_type": "xy",
+ "kind": "bar",
+ "x": {"name": column},
+ "y": [{"name": "revenue", "aggregate": "SUM"}],
+ },
+ }
+ )
+ ctx = Mock(
+ info=AsyncMock(),
+ debug=AsyncMock(),
+ warning=AsyncMock(),
+ error=AsyncMock(),
+ report_progress=AsyncMock(),
+ )
+ with (
+ patch(
+ "superset.mcp_service.chart.tool.generate_chart.event_logger.log_context",
+ side_effect=lambda **kwargs: nullcontext(),
+ ),
+ patch(
+ "superset.mcp_service.auth.get_user_from_request",
+ return_value=Mock(id=1, username="fixture_user", roles=[], groups=[]),
+ ),
+ patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=dataset),
+ patch(
+ "superset.mcp_service.auth.security_manager.can_access_datasource",
+ return_value=allowed,
+ ) as access,
+ ):
+ result = await generate_chart(request, ctx=ctx)
+
+ assert result.success is False
+ assert result.error is not None
+ error = result.error
+ assert error.error_type == expected_type
+ suggestions = " ".join(error.suggestions)
+ assert column not in suggestions
+ assert "Check available columns?" not in suggestions
+ assert "gross_margin" not in suggestions
+ if guidance:
+ assert guidance in suggestions
+ assert error.dataset_context is not None
+ assert error.dataset_context.available_columns == [
+ {"name": name} for name in names
+ ]
+ assert error.dataset_context.available_metrics == []
+ else:
+ assert error.dataset_context is None
+ for name in names:
+ assert name not in error.model_dump_json()
+ access.assert_called_with(datasource=dataset)
+
+
+@pytest.mark.parametrize(
+ "missing", [["private_input"], ["private_input", "other_input"]]
+)
+def test_column_context_is_bounded_and_sanitized(missing: list[str]) -> None:
+ """Neither raw references nor unrestricted metadata become suggestions."""
+ from superset.mcp_service.chart.schemas import ColumnRef
+ from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator
+ from superset.mcp_service.common.error_schemas import DatasetContext
+
+ names = ["<fixture>", "x" * 1000] + [f"column_{i}" for i in range(20)]
+ context = DatasetContext(
+ id=268,
+ table_name="fixture",
+ database_name="fixture",
+ available_columns=[
+ {"name": name, "expression": "PRIVATE SQL"} for name in names
+ ],
+ available_metrics=[{"name": "metric", "expression": "PRIVATE SQL"}],
+ )
+ error = DatasetValidator._validate_columns_exist(
+ [ColumnRef(name=name) for name in missing], context
+ )
+ assert error is not None
+ assert error.dataset_context is not None
+ columns = error.dataset_context.available_columns
+ assert len(columns) == 10
+ assert columns[0] == {"name": "<fixture>"}
+ assert len(columns[1]["name"]) < 220
+ assert "PRIVATE SQL" not in error.model_dump_json()
+ assert len(error.suggestions) <= 10
+ for name in missing:
+ assert name not in " ".join(error.suggestions)
+
+
+def test_column_candidates_are_bounded_and_sanitized() -> None:
+ """Candidate guidance retains escaping and the three-candidate cap."""
+ from superset.mcp_service.utils.error_builder import ChartErrorBuilder
+
+ error = ChartErrorBuilder.column_not_found_error(
+ "private_input", ["<fixture>", "revenue", "category", "excluded"]
+ )
+ assert error.error_type == "column_not_found"
+ assert error.error_code == "CHART_COLUMN_NOT_FOUND"
+ assert error.suggestions[-1] == "Did you mean: <fixture>, revenue, category?"
+ assert "private_input" not in " ".join(error.suggestions)
+
+
+@pytest.mark.parametrize("names", [[], ["revenue", "category"]])
+def test_multiple_column_guidance_uses_real_candidates(names: list[str]) -> None:
+ """Multiple errors share the same real-candidate/no-match behavior."""
+ from superset.mcp_service.chart.schemas import ColumnRef
+ from superset.mcp_service.chart.validation.dataset_validator import DatasetValidator
+ from superset.mcp_service.common.error_schemas import DatasetContext
+
+ context = DatasetContext(
+ id=268,
+ table_name="fixture",
+ database_name="fixture",
+ available_columns=[{"name": name} for name in names],
+ )
+ error = DatasetValidator._validate_columns_exist(
+ [ColumnRef(name="revnue"), ColumnRef(name="categry")], context
+ )
+ assert error is not None
+ assert error.error_type == "multiple_invalid_columns"
+ assert error.error_code == "MULTIPLE_INVALID_COLUMNS"
+ if names:
+ assert error.suggestions[-1] == "Did you mean: revenue, category?"
+ else:
+ assert error.suggestions[-1].startswith("No matching columns found.")
+ assert "revnue" not in " ".join(error.suggestions)
+ assert "categry" not in " ".join(error.suggestions)