fix(mcp): support explicit query_mode in TableChartConfig (#39412)

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)
 

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,

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."""

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"""