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 <[email protected]> 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 + + [email protected] [email protected]( + ("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) + + [email protected]( + "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) + + [email protected]("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)
