From db5424a75f1009a61ed5c95338c647f5c3efa712 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Fri, 31 Jul 2026 18:38:53 +0000 Subject: [PATCH 1/4] fix(mcp): preserve table chart state during updates --- superset/mcp_service/chart/compile.py | 22 +++- superset/mcp_service/chart/schemas.py | 21 ++++ .../mcp_service/chart/tool/update_chart.py | 114 ++++++++++++++++-- .../mcp_service/chart/test_compile.py | 43 +++++-- .../chart/tool/test_update_chart.py | 66 ++++++++++ 5 files changed, 238 insertions(+), 28 deletions(-) diff --git a/superset/mcp_service/chart/compile.py b/superset/mcp_service/chart/compile.py index 2711dc7f0612..33c62f4789d6 100644 --- a/superset/mcp_service/chart/compile.py +++ b/superset/mcp_service/chart/compile.py @@ -242,11 +242,9 @@ def _validate_adhoc_filter_columns( represented on the new config — those would otherwise bypass validation and surface only when Explore tries to run the query. """ - adhoc_filters = form_data.get("adhoc_filters") or [] + adhoc_filters = _active_adhoc_filters(form_data.get("adhoc_filters") or []) invalid: List[str] = [] for f in adhoc_filters: - if not isinstance(f, dict): - continue # SIMPLE filters expose the column via "subject"; SQL-expression # filters carry a free-form ``sqlExpression`` we can't safely parse, # so skip those. @@ -282,13 +280,29 @@ def _validate_adhoc_filter_columns( details=( "Adhoc filter columns must exist on the dataset. " "If these filters were preserved from a previous chart preview, " - "remove them or pass an explicit ``filters`` list on the new config." + "pass an explicit 'filters' list on the new config; use " + "'filters': [] to clear them." ), suggestions=suggestions, error_code="CHART_VALIDATION_FAILED", ) +def _is_inert_adhoc_filter(filter_: dict[str, Any]) -> bool: + """Whether a saved filter is Superset's non-filtering placeholder.""" + comparator = filter_.get("comparator", filter_.get("val")) + return isinstance(comparator, str) and comparator.casefold() == "no filter" + + +def _active_adhoc_filters(filters: list[Any]) -> list[dict[str, Any]]: + """Return structurally valid filters that can produce a predicate.""" + return [ + filter_ + for filter_ in filters + if isinstance(filter_, dict) and not _is_inert_adhoc_filter(filter_) + ] + + def _classify_as_database_error(exc: BaseException, dataset_id: int) -> bool: """Use the dataset's DB engine spec to classify the error. diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index 12bf88e769ac..c4bf7ca9353f 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -2351,6 +2351,14 @@ class UpdateChartRequest(ChartRequestNormalizerMixin, QueryCacheControl): None, description="Chart configuration. Optional; omit to only update chart_name.", ) + add_columns: List[ColumnRef] | None = Field( + None, + description=( + "Table columns or metrics to append while preserving every existing " + "column and metric. Use this instead of config.columns when adding " + "columns to an existing table chart." + ), + ) chart_name: str | None = Field( None, description="Auto-generates if omitted", @@ -2383,6 +2391,19 @@ class UpdateChartRequest(ChartRequestNormalizerMixin, QueryCacheControl): ), ) + @model_validator(mode="after") + def validate_column_patch(self) -> "UpdateChartRequest": + """Keep full-config replacement and additive table updates unambiguous.""" + if self.config is not None and self.add_columns is not None: + raise ValueError( + "Use either 'config' for a full visualization replacement or " + "'add_columns' to append table columns while preserving the existing " + "configuration, not both." + ) + if self.add_columns == []: + raise ValueError("'add_columns' must contain at least one column") + return self + @field_validator("chart_name") @classmethod def sanitize_chart_name(cls, v: str | None) -> str | None: diff --git a/superset/mcp_service/chart/tool/update_chart.py b/superset/mcp_service/chart/tool/update_chart.py index c5070719d99e..1556bb1c657b 100644 --- a/superset/mcp_service/chart/tool/update_chart.py +++ b/superset/mcp_service/chart/tool/update_chart.py @@ -43,8 +43,10 @@ from superset.mcp_service.chart.compile import validate_and_compile from superset.mcp_service.chart.schemas import ( AccessibilityMetadata, + ColumnRef, GenerateChartResponse, PerformanceMetadata, + TableChartConfig, UpdateChartRequest, wrap_sql_adhoc_metrics, ) @@ -77,11 +79,14 @@ def _validation_error_response(message: str, details: str) -> GenerateChartRespo def _missing_config_or_name_error() -> GenerateChartResponse: return _validation_error_response( - message="Either 'config', 'chart_name', or 'dataset_id' must be provided.", + message=( + "Either 'config', 'add_columns', 'chart_name', or 'dataset_id' must be " + "provided." + ), details=( - "Either 'config', 'chart_name', or 'dataset_id' must be provided. " - "Use config for visualization changes, chart_name for renaming, " - "dataset_id to rebind the chart to a different dataset." + "Use config for full visualization changes, add_columns to append table " + "columns without replacing existing columns, chart_name for renaming, " + "or dataset_id to rebind the chart to a different dataset." ), ) @@ -95,6 +100,48 @@ def _wrapped_form_data_for_response( return payload +def _append_table_columns( + existing_form_data: dict[str, Any], + columns: list[ColumnRef], +) -> dict[str, Any] | GenerateChartResponse: + """Append table columns/metrics without replacing the saved column lists.""" + if existing_form_data.get("viz_type") not in {"table", "ag-grid-table"}: + return _validation_error_response( + message="'add_columns' is only supported for table charts.", + details=( + "Use 'config' to replace the full configuration of a non-table chart." + ), + ) + + patch = map_config_to_form_data(TableChartConfig(columns=columns)) + merged = dict(existing_form_data) + query_mode = existing_form_data.get("query_mode") + if query_mode == "raw": + merged["all_columns"] = list(existing_form_data.get("all_columns") or []) + merged["all_columns"].extend( + column.name for column in columns if column.name is not None + ) + return merged + + merged["groupby"] = list(existing_form_data.get("groupby") or []) + merged["metrics"] = list(existing_form_data.get("metrics") or []) + merged["groupby"].extend(patch.get("groupby") or []) + merged["metrics"].extend(patch.get("metrics") or []) + return merged + + +def _merge_replacement_config( + existing_form_data: dict[str, Any], + new_form_data: dict[str, Any], + parsed_config: Any, +) -> dict[str, Any]: + """Merge a replacement config, honoring an explicit empty filter list.""" + merged = {**existing_form_data, **new_form_data} + if getattr(parsed_config, "filters", None) == []: + merged.pop("adhoc_filters", None) + return merged + + def _build_update_payload( request: UpdateChartRequest, chart: Any, @@ -136,6 +183,22 @@ def _build_update_payload( payload["datasource_type"] = "table" return payload + if request.add_columns is not None: + try: + existing_form_data = json.loads(chart.params) if chart.params else {} + except (ValueError, TypeError): + existing_form_data = {} + patched = _append_table_columns(existing_form_data, request.add_columns) + if isinstance(patched, GenerateChartResponse): + return patched + chart_name = request.chart_name or chart.slice_name + return { + "slice_name": chart_name, + "viz_type": patched["viz_type"], + "params": json.dumps(patched), + "query_context": None, + } + # Dataset-only update: rebind chart to a different dataset without changing viz if request.dataset_id is not None: payload = { @@ -185,7 +248,16 @@ def _build_preview_form_data( parsed_config, dataset_id=effective_dataset_id ) new_form_data.pop("_mcp_warnings", None) - merged = {**existing_form_data, **new_form_data} + # An explicit filters list, including [], replaces saved filters. An + # omitted filters field preserves them through the shallow merge. + merged = _merge_replacement_config( + existing_form_data, new_form_data, parsed_config + ) + elif request.add_columns is not None: + patched = _append_table_columns(existing_form_data, request.add_columns) + if isinstance(patched, GenerateChartResponse): + return patched + merged = patched else: if not request.chart_name and request.dataset_id is None: return _missing_config_or_name_error() @@ -352,6 +424,7 @@ async def update_chart( # noqa: C901 - LLM clients MUST display the returned explore URL to users. - Use numeric ID or UUID string to identify the chart (NOT chart name). - config is optional — omit it to rename a chart without changing its visualization + - To append table columns without restating the existing list, use add_columns Example usage (preview, default): ```json @@ -383,6 +456,14 @@ async def update_chart( # noqa: C901 } ``` + Add a table column while preserving existing columns and metrics: + ```json + { + "identifier": 123, + "add_columns": [{"name": "go_live_date", "aggregate": "MIN"}] + } + ``` + Example usage with a custom SQL metric (ratios, conditional aggregations, unit conversions). Pass 'sql_expression' instead of 'name'+'aggregate'. A 'label' is required: @@ -470,6 +551,9 @@ async def update_chart( # noqa: C901 # config is already a typed ChartConfig | None (validated by Pydantic) parsed_config = request.config + validation_config = parsed_config + if request.add_columns is not None: + validation_config = TableChartConfig(columns=request.add_columns) # Normalize column case to match dataset canonical names # (mirrors generate_chart pipeline layer 4) @@ -481,16 +565,22 @@ async def update_chart( # noqa: C901 if request.dataset_id is not None else getattr(chart, "datasource_id", None) ) - if parsed_config is not None and effective_norm_dataset_id is not None: + if validation_config is not None and effective_norm_dataset_id is not None: from superset.mcp_service.chart.validation.dataset_validator import ( DatasetValidator, NORMALIZATION_EXCEPTIONS, ) try: - parsed_config = DatasetValidator.normalize_column_names( - parsed_config, effective_norm_dataset_id + validation_config = DatasetValidator.normalize_column_names( + validation_config, effective_norm_dataset_id ) + if parsed_config is not None: + parsed_config = validation_config + else: + request = request.model_copy( + update={"add_columns": validation_config.columns} + ) except NORMALIZATION_EXCEPTIONS as e: logger.warning( "Column normalization failed for chart %s: %s", chart.id, e @@ -511,10 +601,10 @@ async def update_chart( # noqa: C901 # SQL errors so we don't commit a chart that can't be queried. # Renames (no parsed_config and no dataset_id) skip validation since # form_data is untouched and no rebind is requested. - if parsed_config is not None and new_form_data is not None: + if validation_config is not None and new_form_data is not None: with event_logger.log_context(action="mcp.update_chart.validation"): validation_error = _validate_update_against_dataset( - parsed_config, + validation_config, new_form_data, chart, dataset_id=request.dataset_id, @@ -549,10 +639,10 @@ async def update_chart( # noqa: C901 return preview_or_error # Validate before caching the form_data — same rationale as above. - if parsed_config is not None: + if validation_config is not None: with event_logger.log_context(action="mcp.update_chart.validation"): validation_error = _validate_update_against_dataset( - parsed_config, + validation_config, preview_or_error, chart, dataset_id=request.dataset_id, diff --git a/tests/unit_tests/mcp_service/chart/test_compile.py b/tests/unit_tests/mcp_service/chart/test_compile.py index e33c1d8dde00..53837e5372fb 100644 --- a/tests/unit_tests/mcp_service/chart/test_compile.py +++ b/tests/unit_tests/mcp_service/chart/test_compile.py @@ -210,9 +210,9 @@ def test_pivot_table_min_on_non_numeric_column_passes(self): metrics=[ColumnRef(name="name", aggregate="MIN")], ) result = validate_and_compile(config, {}, ds, run_compile_check=False) - assert result.success, ( - "MIN on a text column should not be rejected by Tier-1 validation" - ) + assert ( + result.success + ), "MIN on a text column should not be rejected by Tier-1 validation" def test_table_with_invalid_filter_column_rejected(self): ds = _orm_dataset() @@ -225,6 +225,25 @@ def test_table_with_invalid_filter_column_rejected(self): assert not result.success assert result.error_obj is not None + def test_inert_stale_filter_column_is_ignored(self): + """A No filter placeholder produces no predicate and cannot block edits.""" + ds = _orm_dataset() + config = TableChartConfig(columns=[ColumnRef(name="gender")]) + form_data = { + "adhoc_filters": [ + { + "expressionType": "SIMPLE", + "subject": "dropped_column", + "operator": "TEMPORAL_RANGE", + "comparator": "No filter", + } + ] + } + + result = validate_and_compile(config, form_data, ds, run_compile_check=False) + + assert result.success + class TestSavedMetricNotMarked: """A non-saved-metric ColumnRef whose name matches a saved metric is a @@ -245,9 +264,9 @@ def test_table_metric_name_without_saved_metric_flag_rejected(self): ], ) result = validate_and_compile(config, {}, ds, run_compile_check=False) - assert not result.success, ( - "ref.name matches a saved metric but saved_metric=False -> reject" - ) + assert ( + not result.success + ), "ref.name matches a saved metric but saved_metric=False -> reject" assert result.error_obj is not None assert result.error_obj.error_code == "SAVED_METRIC_NOT_MARKED" # Suggestion should point the LLM at the right correction. @@ -360,9 +379,9 @@ def test_where_filter_with_metric_name_rejected(self): ] } result = validate_and_compile(config, form_data, ds, run_compile_check=False) - assert not result.success, ( - "A saved-metric name used in a WHERE filter must not pass Tier-1" - ) + assert ( + not result.success + ), "A saved-metric name used in a WHERE filter must not pass Tier-1" assert result.error_obj is not None assert "sum_boys" in (result.error_obj.message or "") @@ -388,9 +407,9 @@ def test_having_filter_with_metric_name_passes(self): ] } result = validate_and_compile(config, form_data, ds, run_compile_check=False) - assert result.success, ( - "A saved-metric name in a HAVING filter should pass Tier-1 validation" - ) + assert ( + result.success + ), "A saved-metric name in a HAVING filter should pass Tier-1 validation" class TestValidateAndCompileTier2: diff --git a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py index a4d1da7b2693..eaeac348a732 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py @@ -43,6 +43,7 @@ _build_preview_form_data, _build_update_payload, ) +from superset.utils import json # The __init__.py re-exports the update_chart *function*, so a plain # `from ... import update_chart` gives the function, not the module. @@ -76,6 +77,10 @@ async def test_update_chart_request_structure(self): assert table_request.config.columns[0].name == "region" assert table_request.config.columns[1].aggregate == "SUM" + # Dataset column names may contain punctuation supported by Superset. + parenthesized = ColumnRef(name="New Draft Apps (This Week)") + assert parenthesized.name == "New Draft Apps (This Week)" + # XY chart update with UUID xy_config = XYChartConfig( chart_type="xy", @@ -638,6 +643,36 @@ def test_config_update_keeps_existing_name(self): # query_context must be cleared so get_chart_data uses updated params assert result["query_context"] is None + def test_add_columns_preserves_existing_columns_and_metrics(self): + """An additive update does not require reconstructing the table.""" + request = UpdateChartRequest( + identifier=1, + add_columns=[ + ColumnRef( + name="go_live_date", aggregate="MIN", label="Earliest Go Live Date" + ) + ], + ) + chart = Mock() + chart.slice_name = "Existing" + chart.params = json.dumps( + { + "viz_type": "ag-grid-table", + "query_mode": "aggregate", + "groupby": ["employer"], + "metrics": ["count"], + } + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + params = json.loads(result["params"]) + assert params["groupby"] == ["employer"] + assert params["metrics"][0] == "count" + assert params["metrics"][1]["label"] == "Earliest Go Live Date" + assert params["metrics"][1]["aggregate"] == "MIN" + class TestUpdateChartNameOnly: """Integration-style tests for name-only update via MCP tool.""" @@ -936,6 +971,37 @@ def test_handles_invalid_existing_params(self): assert result["slice_id"] == 9 assert result["slice_name"] == "Broken" + def test_explicit_empty_filters_clear_saved_filters(self): + """filters=[] must honor the remedy advertised by validation errors.""" + config = TableChartConfig( + columns=[ColumnRef(name="region")], + filters=[], + ) + request = UpdateChartRequest(identifier=1, config=config) + chart = Mock( + id=9, + datasource_id=4, + slice_name="Filtered", + params=json.dumps( + { + "viz_type": "table", + "adhoc_filters": [ + { + "expressionType": "SIMPLE", + "subject": "dropped_column", + "operator": "TEMPORAL_RANGE", + "comparator": "No filter", + } + ], + } + ), + ) + + result = _build_preview_form_data(request, chart, parsed_config=config) + + assert isinstance(result, dict) + assert "adhoc_filters" not in result + class TestUpdateChartSaveWithConfig: """Save-path integration tests for update_chart with a full config payload.""" From c6f19c0b1f6f6ed3f733574d16c2e18d09ec8abb Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Tue, 4 Aug 2026 19:45:18 +0000 Subject: [PATCH 2/4] fix(mcp): validate additive table updates --- superset/mcp_service/chart/compile.py | 8 ++- .../mcp_service/chart/tool/update_chart.py | 16 ++++- .../mcp_service/chart/test_compile.py | 21 +++++++ .../chart/tool/test_update_chart.py | 59 +++++++++++++++++++ 4 files changed, 102 insertions(+), 2 deletions(-) diff --git a/superset/mcp_service/chart/compile.py b/superset/mcp_service/chart/compile.py index 33c62f4789d6..062be8ef6cb7 100644 --- a/superset/mcp_service/chart/compile.py +++ b/superset/mcp_service/chart/compile.py @@ -290,8 +290,14 @@ def _validate_adhoc_filter_columns( def _is_inert_adhoc_filter(filter_: dict[str, Any]) -> bool: """Whether a saved filter is Superset's non-filtering placeholder.""" + operator = filter_.get("operator", filter_.get("op")) comparator = filter_.get("comparator", filter_.get("val")) - return isinstance(comparator, str) and comparator.casefold() == "no filter" + return ( + isinstance(operator, str) + and operator.casefold() == "temporal_range" + and isinstance(comparator, str) + and comparator.casefold() == "no filter" + ) def _active_adhoc_filters(filters: list[Any]) -> list[dict[str, Any]]: diff --git a/superset/mcp_service/chart/tool/update_chart.py b/superset/mcp_service/chart/tool/update_chart.py index 1556bb1c657b..8d0478b15822 100644 --- a/superset/mcp_service/chart/tool/update_chart.py +++ b/superset/mcp_service/chart/tool/update_chart.py @@ -117,6 +117,16 @@ def _append_table_columns( merged = dict(existing_form_data) query_mode = existing_form_data.get("query_mode") if query_mode == "raw": + metric_columns = [column for column in columns if column.is_metric] + if metric_columns: + return _validation_error_response( + message="Cannot add metrics to a table in raw query mode.", + details=( + "Raw tables accept only unaggregated columns in 'add_columns'. " + "Use 'config' with query_mode='aggregate' and the complete table " + "configuration to convert the chart before adding metrics." + ), + ) merged["all_columns"] = list(existing_form_data.get("all_columns") or []) merged["all_columns"].extend( column.name for column in columns if column.name is not None @@ -192,12 +202,16 @@ def _build_update_payload( if isinstance(patched, GenerateChartResponse): return patched chart_name = request.chart_name or chart.slice_name - return { + additive_payload: dict[str, Any] = { "slice_name": chart_name, "viz_type": patched["viz_type"], "params": json.dumps(patched), "query_context": None, } + if request.dataset_id is not None: + additive_payload["datasource_id"] = request.dataset_id + additive_payload["datasource_type"] = "table" + return additive_payload # Dataset-only update: rebind chart to a different dataset without changing viz if request.dataset_id is not None: diff --git a/tests/unit_tests/mcp_service/chart/test_compile.py b/tests/unit_tests/mcp_service/chart/test_compile.py index 53837e5372fb..81d8709f8afa 100644 --- a/tests/unit_tests/mcp_service/chart/test_compile.py +++ b/tests/unit_tests/mcp_service/chart/test_compile.py @@ -244,6 +244,27 @@ def test_inert_stale_filter_column_is_ignored(self): assert result.success + def test_no_filter_literal_with_non_temporal_operator_is_validated(self): + """A literal value of No filter is not generally an inert predicate.""" + ds = _orm_dataset() + config = TableChartConfig(columns=[ColumnRef(name="gender")]) + form_data = { + "adhoc_filters": [ + { + "expressionType": "SIMPLE", + "subject": "dropped_column", + "operator": "==", + "comparator": "No filter", + } + ] + } + + result = validate_and_compile(config, form_data, ds, run_compile_check=False) + + assert not result.success + assert result.error_obj is not None + assert result.error_obj.error_type == "invalid_column" + class TestSavedMetricNotMarked: """A non-saved-metric ColumnRef whose name matches a saved metric is a diff --git a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py index eaeac348a732..54f6fce6c2c8 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py @@ -673,6 +673,65 @@ def test_add_columns_preserves_existing_columns_and_metrics(self): assert params["metrics"][1]["label"] == "Earliest Go Live Date" assert params["metrics"][1]["aggregate"] == "MIN" + def test_add_columns_rebinds_requested_dataset(self): + """Additive saves validate and persist against the same dataset.""" + request = UpdateChartRequest( + identifier=1, + dataset_id=22, + add_columns=[ColumnRef(name="region")], + ) + chart = Mock( + slice_name="Existing", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "aggregate", + "groupby": ["employer"], + "metrics": [], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + assert result["datasource_id"] == 22 + assert result["datasource_type"] == "table" + + @pytest.mark.parametrize( + "metric", + [ + ColumnRef(name="go_live_date", aggregate="MIN"), + ColumnRef(name="saved_count", saved_metric=True), + ColumnRef(sql_expression="COUNT(*)", label="Count"), + ], + ) + def test_add_metric_to_raw_table_returns_actionable_error( + self, metric: ColumnRef + ) -> None: + """Raw tables must not silently discard aggregate semantics.""" + request = UpdateChartRequest(identifier=1, add_columns=[metric]) + chart = Mock( + slice_name="Raw table", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "raw", + "all_columns": ["employer"], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, GenerateChartResponse) + assert result.success is False + assert result.error is not None + assert ( + result.error.message == "Cannot add metrics to a table in raw query mode." + ) + assert "query_mode='aggregate'" in result.error.details + class TestUpdateChartNameOnly: """Integration-style tests for name-only update via MCP tool.""" From 3472ea554d7a7f9884c5fa03d93b32603fcbe4ef Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 5 Aug 2026 00:02:48 +0000 Subject: [PATCH 3/4] style: apply ruff-format to test_compile.py Co-Authored-By: Claude --- .../mcp_service/chart/test_compile.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/unit_tests/mcp_service/chart/test_compile.py b/tests/unit_tests/mcp_service/chart/test_compile.py index 81d8709f8afa..60608f5e4b5b 100644 --- a/tests/unit_tests/mcp_service/chart/test_compile.py +++ b/tests/unit_tests/mcp_service/chart/test_compile.py @@ -210,9 +210,9 @@ def test_pivot_table_min_on_non_numeric_column_passes(self): metrics=[ColumnRef(name="name", aggregate="MIN")], ) result = validate_and_compile(config, {}, ds, run_compile_check=False) - assert ( - result.success - ), "MIN on a text column should not be rejected by Tier-1 validation" + assert result.success, ( + "MIN on a text column should not be rejected by Tier-1 validation" + ) def test_table_with_invalid_filter_column_rejected(self): ds = _orm_dataset() @@ -285,9 +285,9 @@ def test_table_metric_name_without_saved_metric_flag_rejected(self): ], ) result = validate_and_compile(config, {}, ds, run_compile_check=False) - assert ( - not result.success - ), "ref.name matches a saved metric but saved_metric=False -> reject" + assert not result.success, ( + "ref.name matches a saved metric but saved_metric=False -> reject" + ) assert result.error_obj is not None assert result.error_obj.error_code == "SAVED_METRIC_NOT_MARKED" # Suggestion should point the LLM at the right correction. @@ -400,9 +400,9 @@ def test_where_filter_with_metric_name_rejected(self): ] } result = validate_and_compile(config, form_data, ds, run_compile_check=False) - assert ( - not result.success - ), "A saved-metric name used in a WHERE filter must not pass Tier-1" + assert not result.success, ( + "A saved-metric name used in a WHERE filter must not pass Tier-1" + ) assert result.error_obj is not None assert "sum_boys" in (result.error_obj.message or "") @@ -428,9 +428,9 @@ def test_having_filter_with_metric_name_passes(self): ] } result = validate_and_compile(config, form_data, ds, run_compile_check=False) - assert ( - result.success - ), "A saved-metric name in a HAVING filter should pass Tier-1 validation" + assert result.success, ( + "A saved-metric name in a HAVING filter should pass Tier-1 validation" + ) class TestValidateAndCompileTier2: From 8f9a96a24666488d1a0b1c25c0b607771724754d Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 5 Aug 2026 04:27:56 +0000 Subject: [PATCH 4/4] fix(mcp): route added table columns by kind and skip duplicates Addresses review feedback on #42655. _append_table_columns() built its patch with map_config_to_form_data(TableChartConfig(columns=...)), which infers query_mode from the columns it is handed. A dimension-only append compiles to a raw table whose groupby is empty, so appending a plain column to an aggregate table extended groupby with nothing and the column was silently dropped. Route columns by ColumnRef.is_metric instead, so dimensions reach groupby and metrics reach metrics regardless of the mix. Appends also blindly extended the saved lists, so re-adding a column already on the chart duplicated it. Skip entries that are already present, keyed on a stable serialization -- metric lists mix saved names with adhoc dicts, so they are not reliably hashable and a set would both crash and lose the ordering that drives table layout. Cover the config/add_columns validator error branches, which had no tests, and document that add_columns combines with chart_name. Co-Authored-By: Claude --- .../mcp_service/chart/tool/update_chart.py | 68 +++++++- .../chart/tool/test_update_chart.py | 164 ++++++++++++++++++ 2 files changed, 223 insertions(+), 9 deletions(-) diff --git a/superset/mcp_service/chart/tool/update_chart.py b/superset/mcp_service/chart/tool/update_chart.py index 8d0478b15822..e1eea366a84c 100644 --- a/superset/mcp_service/chart/tool/update_chart.py +++ b/superset/mcp_service/chart/tool/update_chart.py @@ -100,6 +100,33 @@ def _wrapped_form_data_for_response( return payload +def _entry_key(entry: Any) -> str: + """Stable identity for a form_data column or metric entry. + + Metric lists mix saved-metric names with adhoc metric dicts, so entries + are not reliably hashable. Serializing gives every shape a comparable + key without dropping the unhashable ones. + """ + return json.dumps(entry, sort_keys=True, default=str) + + +def _extend_without_duplicates(existing: list[Any], additions: Any) -> list[Any]: + """Append entries that are not already present, preserving order. + + Column and metric order drives the rendered table layout, so appending + keeps first-occurrence position rather than rebuilding from a set. + """ + merged = list(existing) + seen = {_entry_key(entry) for entry in merged} + for entry in additions: + key = _entry_key(entry) + if key in seen: + continue + seen.add(key) + merged.append(entry) + return merged + + def _append_table_columns( existing_form_data: dict[str, Any], columns: list[ColumnRef], @@ -113,11 +140,11 @@ def _append_table_columns( ), ) - patch = map_config_to_form_data(TableChartConfig(columns=columns)) merged = dict(existing_form_data) + metric_columns = [column for column in columns if column.is_metric] + dimension_columns = [column for column in columns if not column.is_metric] query_mode = existing_form_data.get("query_mode") if query_mode == "raw": - metric_columns = [column for column in columns if column.is_metric] if metric_columns: return _validation_error_response( message="Cannot add metrics to a table in raw query mode.", @@ -127,16 +154,29 @@ def _append_table_columns( "configuration to convert the chart before adding metrics." ), ) - merged["all_columns"] = list(existing_form_data.get("all_columns") or []) - merged["all_columns"].extend( - column.name for column in columns if column.name is not None + merged["all_columns"] = _extend_without_duplicates( + list(existing_form_data.get("all_columns") or []), + (column.name for column in columns if column.name is not None), ) return merged - merged["groupby"] = list(existing_form_data.get("groupby") or []) - merged["metrics"] = list(existing_form_data.get("metrics") or []) - merged["groupby"].extend(patch.get("groupby") or []) - merged["metrics"].extend(patch.get("metrics") or []) + # map_config_to_form_data() infers query_mode from the columns handed to + # it, so a dimension-only patch compiles to a raw table whose groupby is + # empty. Route each kind by is_metric instead, to keep an aggregate chart + # aggregate no matter which mix of columns is appended. + metric_patch = ( + map_config_to_form_data(TableChartConfig(columns=metric_columns)) + if metric_columns + else {} + ) + merged["groupby"] = _extend_without_duplicates( + list(existing_form_data.get("groupby") or []), + (column.name for column in dimension_columns if column.name is not None), + ) + merged["metrics"] = _extend_without_duplicates( + list(existing_form_data.get("metrics") or []), + metric_patch.get("metrics") or [], + ) return merged @@ -478,6 +518,16 @@ async def update_chart( # noqa: C901 } ``` + add_columns combines with chart_name to append and rename in one call. + Columns already on the chart are ignored, so repeating a column is safe: + ```json + { + "identifier": 123, + "chart_name": "Go-Live Tracker", + "add_columns": [{"name": "region"}] + } + ``` + Example usage with a custom SQL metric (ratios, conditional aggregations, unit conversions). Pass 'sql_expression' instead of 'name'+'aggregate'. A 'label' is required: diff --git a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py index 54f6fce6c2c8..d3643192ca1c 100644 --- a/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py +++ b/tests/unit_tests/mcp_service/chart/tool/test_update_chart.py @@ -25,6 +25,7 @@ import pytest from fastmcp import Client +from pydantic import ValidationError from superset.mcp_service.app import mcp from superset.mcp_service.chart.chart_helpers import find_chart_by_identifier @@ -732,6 +733,169 @@ def test_add_metric_to_raw_table_returns_actionable_error( ) assert "query_mode='aggregate'" in result.error.details + def test_add_dimension_only_column_to_aggregate_table(self): + """A dimension-only append compiles to a raw patch on its own, so it + must be routed by is_metric rather than the patch's inferred mode.""" + request = UpdateChartRequest( + identifier=1, add_columns=[ColumnRef(name="region")] + ) + chart = Mock( + slice_name="Existing", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "aggregate", + "groupby": ["employer"], + "metrics": ["count"], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + params = json.loads(result["params"]) + assert params["groupby"] == ["employer", "region"] + assert params["metrics"] == ["count"] + # The chart must not be flipped to raw by the append. + assert params["query_mode"] == "aggregate" + + def test_add_mixed_dimension_and_metric_to_aggregate_table(self): + """Dimensions land in groupby and metrics in metrics, in one call.""" + request = UpdateChartRequest( + identifier=1, + add_columns=[ + ColumnRef(name="region"), + ColumnRef(name="go_live_date", aggregate="MIN", label="Earliest"), + ], + ) + chart = Mock( + slice_name="Existing", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "aggregate", + "groupby": ["employer"], + "metrics": ["count"], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + params = json.loads(result["params"]) + assert params["groupby"] == ["employer", "region"] + assert params["metrics"][0] == "count" + assert params["metrics"][1]["label"] == "Earliest" + + def test_add_columns_skips_entries_already_on_the_chart(self): + """Re-adding a saved column must not duplicate it in the table.""" + request = UpdateChartRequest( + identifier=1, + add_columns=[ColumnRef(name="employer"), ColumnRef(name="region")], + ) + chart = Mock( + slice_name="Existing", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "aggregate", + "groupby": ["employer"], + "metrics": ["count"], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + params = json.loads(result["params"]) + # "employer" is already present, so only "region" is appended. + assert params["groupby"] == ["employer", "region"] + assert params["metrics"] == ["count"] + + def test_add_columns_dedupes_adhoc_metric_dicts(self): + """Adhoc metrics are dicts -- unhashable, so identity is structural.""" + request = UpdateChartRequest( + identifier=1, + add_columns=[ + ColumnRef( + name="go_live_date", aggregate="MIN", label="Earliest Go Live Date" + ) + ], + ) + chart = Mock(slice_name="Existing") + chart.params = json.dumps( + { + "viz_type": "table", + "query_mode": "aggregate", + "groupby": [], + "metrics": [], + } + ) + + first = _build_update_payload(request, chart) + assert isinstance(first, dict) + # Feed the result back in to emulate the same additive call twice. + chart.params = first["params"] + second = _build_update_payload(request, chart) + + assert isinstance(second, dict) + metrics = json.loads(second["params"])["metrics"] + assert len(metrics) == 1 + assert metrics[0]["label"] == "Earliest Go Live Date" + + def test_add_columns_dedupes_raw_mode_all_columns(self): + """Raw tables append into all_columns and must not duplicate either.""" + request = UpdateChartRequest( + identifier=1, + add_columns=[ColumnRef(name="employer"), ColumnRef(name="region")], + ) + chart = Mock( + slice_name="Raw table", + params=json.dumps( + { + "viz_type": "table", + "query_mode": "raw", + "all_columns": ["employer"], + } + ), + ) + + result = _build_update_payload(request, chart) + + assert isinstance(result, dict) + assert json.loads(result["params"])["all_columns"] == ["employer", "region"] + + +class TestUpdateChartRequestColumnPatchValidation: + """The config/add_columns validator on UpdateChartRequest.""" + + def test_config_and_add_columns_together_rejected(self): + """Full replacement and additive append are mutually exclusive.""" + with pytest.raises(ValidationError, match="not both"): + UpdateChartRequest( + identifier=1, + config=TableChartConfig(columns=[ColumnRef(name="region")]), + add_columns=[ColumnRef(name="employer")], + ) + + def test_empty_add_columns_rejected(self): + """An empty append is a no-op the caller almost certainly didn't mean.""" + with pytest.raises(ValidationError, match="at least one column"): + UpdateChartRequest(identifier=1, add_columns=[]) + + def test_add_columns_alone_accepted(self): + """The valid additive shape still passes validation.""" + request = UpdateChartRequest( + identifier=1, add_columns=[ColumnRef(name="region")] + ) + + assert request.config is None + assert request.add_columns is not None + assert request.add_columns[0].name == "region" + class TestUpdateChartNameOnly: """Integration-style tests for name-only update via MCP tool."""