diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py index 421e4bc37f4..f62d65df816 100644 --- a/superset/mcp_service/chart/chart_utils.py +++ b/superset/mcp_service/chart/chart_utils.py @@ -648,7 +648,7 @@ def _resolve_default_x_axis( return config.model_copy(update={"x": ColumnRef(name=dataset.main_dttm_col)}) -def map_xy_config( +def map_xy_config( # noqa: C901 config: XYChartConfig, dataset_id: int | str | None = None ) -> Dict[str, Any]: """Map XY chart config to form_data with defensive validation.""" @@ -714,6 +714,9 @@ def map_xy_config( form_data["row_limit"] = config.row_limit + if config.series_limit is not None: + form_data["series_limit"] = config.series_limit + # Add stacking configuration if getattr(config, "stacked", False): form_data["stack"] = "Stack" diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index aa3848141db..2e0532cecc8 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -1304,6 +1304,16 @@ class XYChartConfig(UnknownFieldCheckMixin): "Do NOT use adhoc_filters or raw SQL expressions.", ) row_limit: int = Field(10000, description="Max data points", ge=1, le=50000) + series_limit: int | None = Field( + None, + description=( + "Max number of series to show when group_by is set. " + "Limits the distinct values rendered as separate lines/bars. " + "Only applies when group_by is specified." + ), + ge=1, + le=10000, + ) @field_validator("group_by", mode="before") @classmethod 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 65ec838d16f..2d91cdd0048 100644 --- a/tests/unit_tests/mcp_service/chart/test_chart_schemas.py +++ b/tests/unit_tests/mcp_service/chart/test_chart_schemas.py @@ -452,6 +452,43 @@ class TestRowLimit: row_limit=100000, ) + def test_xy_chart_series_limit_default_none(self) -> None: + """Test that XYChartConfig series_limit defaults to None.""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + ) + assert config.series_limit is None + + def test_xy_chart_series_limit_custom(self) -> None: + """Test that XYChartConfig accepts a custom series_limit.""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + group_by=[ColumnRef(name="region")], + series_limit=5, + ) + assert config.series_limit == 5 + + def test_xy_chart_series_limit_validation(self) -> None: + """Test that XYChartConfig rejects invalid series_limit values.""" + with pytest.raises(ValidationError): + XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + series_limit=0, + ) + with pytest.raises(ValidationError): + XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + series_limit=10001, + ) + def test_table_chart_default_row_limit(self) -> None: """Test that TableChartConfig has default row_limit of 1000.""" config = TableChartConfig( 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 7a7567f48f2..4198dac3ea4 100644 --- a/tests/unit_tests/mcp_service/chart/test_chart_utils.py +++ b/tests/unit_tests/mcp_service/chart/test_chart_utils.py @@ -804,6 +804,38 @@ class TestMapXYConfig: assert result["row_limit"] == 10000 + @patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal") + def test_map_xy_config_series_limit(self, mock_is_temporal) -> None: + """Test that series_limit is mapped to form_data when set.""" + mock_is_temporal.return_value = True + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + kind="line", + group_by=[ColumnRef(name="region")], + series_limit=10, + ) + + result = map_xy_config(config) + + assert result["series_limit"] == 10 + + @patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal") + def test_map_xy_config_no_series_limit_by_default(self, mock_is_temporal) -> None: + """Test that series_limit is omitted from form_data when not set.""" + mock_is_temporal.return_value = True + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + kind="line", + ) + + result = map_xy_config(config) + + assert "series_limit" not in result + @patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal") def test_map_xy_config_saved_metric(self, mock_is_temporal: Any) -> None: """Test XY config with saved metric emits string in metrics list"""