From 2e0d482ccf5301149d18bdbf3c7e64e8d20fbc3c Mon Sep 17 00:00:00 2001 From: Gabriel Torres Ruiz Date: Thu, 16 Apr 2026 18:53:25 -0300 Subject: [PATCH] fix(mcp): support explicit query_mode in TableChartConfig (#39412) --- superset/mcp_service/chart/chart_utils.py | 103 ++++++++++-------- superset/mcp_service/chart/schemas.py | 11 ++ .../mcp_service/chart/test_chart_schemas.py | 38 +++++++ .../mcp_service/chart/test_chart_utils.py | 56 ++++++++++ 4 files changed, 165 insertions(+), 43 deletions(-) diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py index b5e2d6b0a09..3f2f9e2cb31 100644 --- a/superset/mcp_service/chart/chart_utils.py +++ b/superset/mcp_service/chart/chart_utils.py @@ -388,63 +388,80 @@ def map_table_config(config: TableChartConfig) -> Dict[str, Any]: if not config.columns: raise ValueError("Table chart must have at least one column") - # Separate columns with aggregates from raw columns - raw_columns = [] - aggregated_metrics = [] - - for col in config.columns: - if col.is_metric: - # Saved metric or column with aggregation - treat as metric - aggregated_metrics.append(create_metric_object(col)) - else: - # No aggregation - treat as raw column - raw_columns.append(col.name) - - # Final validation - ensure we have some data to display - if not raw_columns and not aggregated_metrics: - raise ValueError("Table chart configuration resulted in no displayable columns") - # Use the viz_type from config (defaults to "table", can be "ag-grid-table") form_data: Dict[str, Any] = { "viz_type": config.viz_type, } - # Handle raw columns (no aggregation) - if raw_columns and not aggregated_metrics: - # Pure raw columns - show individual rows - # Include both "all_columns" (Superset table viz) and "columns" - # (QueryContextFactory validation) to avoid "Empty query?" errors + # When query_mode is explicitly set to "raw", force raw mode for all columns. + # Aggregate settings on individual columns are ignored in this case. + if config.query_mode == "raw": + column_names = [col.name for col in config.columns] form_data.update( { - "all_columns": raw_columns, - "columns": raw_columns, + "all_columns": column_names, + "columns": column_names, "query_mode": "raw", "include_time": False, "order_desc": True, } ) + else: + # Auto-detect or explicit "aggregate": separate columns with aggregates + # from raw columns and build the appropriate form_data. + raw_columns = [] + aggregated_metrics = [] - # Handle aggregated columns only - elif aggregated_metrics and not raw_columns: - # Pure aggregation - show totals - form_data.update( - { - "metrics": aggregated_metrics, - "query_mode": "aggregate", - } - ) + for col in config.columns: + if col.is_metric: + # Saved metric or column with aggregation - treat as metric + aggregated_metrics.append(create_metric_object(col)) + else: + # No aggregation - treat as raw column + raw_columns.append(col.name) - # Handle mixed columns (raw + aggregated) - elif raw_columns and aggregated_metrics: - # Mixed mode - group by raw columns, aggregate metrics - form_data.update( - { - "all_columns": raw_columns, - "metrics": aggregated_metrics, - "groupby": raw_columns, - "query_mode": "aggregate", - } - ) + # Final validation - ensure we have some data to display + if not raw_columns and not aggregated_metrics: + raise ValueError( + "Table chart configuration resulted in no displayable columns" + ) + + # Handle raw columns (no aggregation) + if raw_columns and not aggregated_metrics: + # Pure raw columns - show individual rows + # Include both "all_columns" (Superset table viz) and "columns" + # (QueryContextFactory validation) to avoid "Empty query?" errors + form_data.update( + { + "all_columns": raw_columns, + "columns": raw_columns, + "query_mode": "raw", + "include_time": False, + "order_desc": True, + } + ) + + # Handle aggregated columns only + elif aggregated_metrics and not raw_columns: + # Pure aggregation - show totals + form_data.update( + { + "metrics": aggregated_metrics, + "query_mode": "aggregate", + } + ) + + # Handle mixed columns (raw + aggregated) + else: + # Mixed mode - group by raw columns, aggregate metrics + form_data.update( + { + "all_columns": raw_columns, + "metrics": aggregated_metrics, + "groupby": raw_columns, + "query_mode": "aggregate", + } + ) _add_adhoc_filters(form_data, config.filters) diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index c03599bb141..8e5e13d0c44 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -996,6 +996,17 @@ class TableChartConfig(UnknownFieldCheckMixin): viz_type: Literal["table", "ag-grid-table"] = Field( "table", description="'ag-grid-table' for interactive features" ) + query_mode: Literal["aggregate", "raw"] | None = Field( + None, + description=( + "Query mode: 'raw' returns individual rows without aggregation, " + "'aggregate' groups data using metrics. " + "When set to 'raw', all columns are treated as plain columns regardless " + "of any aggregate settings. " + "Defaults to auto-detection: 'raw' if no column has an aggregate " + "function, 'aggregate' otherwise." + ), + ) columns: List[ColumnRef] = Field( ..., min_length=1, diff --git a/tests/unit_tests/mcp_service/chart/test_chart_schemas.py b/tests/unit_tests/mcp_service/chart/test_chart_schemas.py index 682471b4d4a..1d3b671fa7a 100644 --- a/tests/unit_tests/mcp_service/chart/test_chart_schemas.py +++ b/tests/unit_tests/mcp_service/chart/test_chart_schemas.py @@ -109,6 +109,44 @@ class TestTableChartConfig: columns=[ColumnRef(name="product")], ) + def test_explicit_raw_query_mode_accepted(self) -> None: + """Test that TableChartConfig accepts explicit query_mode='raw'.""" + config = TableChartConfig( + chart_type="table", + query_mode="raw", + columns=[ColumnRef(name="product"), ColumnRef(name="category")], + ) + assert config.query_mode == "raw" + assert len(config.columns) == 2 + + def test_explicit_aggregate_query_mode_accepted(self) -> None: + """Test that TableChartConfig accepts explicit query_mode='aggregate'.""" + config = TableChartConfig( + chart_type="table", + query_mode="aggregate", + columns=[ColumnRef(name="sales", aggregate="SUM")], + ) + assert config.query_mode == "aggregate" + + def test_default_query_mode_is_none(self) -> None: + """Test that default query_mode is None (auto-detect).""" + config = TableChartConfig( + chart_type="table", + columns=[ColumnRef(name="product")], + ) + assert config.query_mode is None + + def test_invalid_query_mode_rejected(self) -> None: + """Test that invalid query_mode values are rejected.""" + from pydantic import ValidationError + + with pytest.raises(ValidationError): + TableChartConfig( + chart_type="table", + query_mode="invalid", + columns=[ColumnRef(name="product")], + ) + class TestXYChartConfig: """Test XYChartConfig validation.""" diff --git a/tests/unit_tests/mcp_service/chart/test_chart_utils.py b/tests/unit_tests/mcp_service/chart/test_chart_utils.py index 293adfb09c3..894440d4888 100644 --- a/tests/unit_tests/mcp_service/chart/test_chart_utils.py +++ b/tests/unit_tests/mcp_service/chart/test_chart_utils.py @@ -390,6 +390,62 @@ class TestMapTableConfig: assert result["metrics"] == ["total_revenue", "avg_order_value"] assert "all_columns" not in result + def test_map_table_config_explicit_raw_mode(self) -> None: + """Test that explicit query_mode='raw' forces raw mode.""" + config = TableChartConfig( + chart_type="table", + query_mode="raw", + columns=[ + ColumnRef(name="product"), + ColumnRef(name="category"), + ], + ) + + result = map_table_config(config) + + assert result["viz_type"] == "table" + assert result["query_mode"] == "raw" + assert result["all_columns"] == ["product", "category"] + assert result["columns"] == ["product", "category"] + assert "metrics" not in result + + def test_map_table_config_explicit_raw_mode_ignores_aggregates(self) -> None: + """Test that explicit query_mode='raw' ignores aggregate settings on columns.""" + config = TableChartConfig( + chart_type="table", + query_mode="raw", + columns=[ + ColumnRef(name="product"), + ColumnRef(name="sales", aggregate="SUM"), + ], + ) + + result = map_table_config(config) + + assert result["query_mode"] == "raw" + # Both columns treated as raw; aggregate setting on "sales" is ignored + assert result["all_columns"] == ["product", "sales"] + assert result["columns"] == ["product", "sales"] + assert "metrics" not in result + + def test_map_table_config_explicit_aggregate_mode(self) -> None: + """Test that explicit query_mode='aggregate' uses inference logic.""" + config = TableChartConfig( + chart_type="table", + query_mode="aggregate", + columns=[ + ColumnRef(name="product"), + ColumnRef(name="revenue", aggregate="SUM"), + ], + ) + + result = map_table_config(config) + + assert result["query_mode"] == "aggregate" + assert len(result["metrics"]) == 1 + assert result["metrics"][0]["aggregate"] == "SUM" + assert "product" in result["groupby"] + class TestAddAdhocFilters: """Test _add_adhoc_filters helper function"""