fix: Row limit support for chart mcp tools (#38717)

This commit is contained in:
Kamil Gabryjelski
2026-03-18 13:40:47 +01:00
committed by GitHub
parent e02ca8871d
commit a314e5b35e
5 changed files with 266 additions and 64 deletions

View File

@@ -299,6 +299,78 @@ class TestXYChartConfig:
assert config.orientation == "horizontal"
class TestRowLimit:
"""Test row_limit field on chart configs."""
def test_xy_chart_default_row_limit(self) -> None:
"""Test that XYChartConfig has default row_limit of 10000."""
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
)
assert config.row_limit == 10000
def test_xy_chart_custom_row_limit(self) -> None:
"""Test that XYChartConfig accepts custom row_limit."""
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
row_limit=100,
)
assert config.row_limit == 100
def test_xy_chart_row_limit_validation(self) -> None:
"""Test that XYChartConfig rejects invalid row_limit."""
with pytest.raises(ValidationError):
XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
row_limit=0,
)
with pytest.raises(ValidationError):
XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
row_limit=100000,
)
def test_table_chart_default_row_limit(self) -> None:
"""Test that TableChartConfig has default row_limit of 1000."""
config = TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product")],
)
assert config.row_limit == 1000
def test_table_chart_custom_row_limit(self) -> None:
"""Test that TableChartConfig accepts custom row_limit."""
config = TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product")],
row_limit=500,
)
assert config.row_limit == 500
def test_table_chart_row_limit_validation(self) -> None:
"""Test that TableChartConfig rejects invalid row_limit."""
with pytest.raises(ValidationError):
TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product")],
row_limit=0,
)
with pytest.raises(ValidationError):
TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product")],
row_limit=100000,
)
class TestTableChartConfigExtraFields:
"""Test TableChartConfig rejects unknown fields."""

View File

@@ -23,6 +23,7 @@ from unittest.mock import MagicMock, patch
import pytest
from superset.mcp_service.chart.chart_utils import (
_add_adhoc_filters,
configure_temporal_handling,
create_metric_object,
generate_chart_name,
@@ -313,6 +314,66 @@ class TestMapTableConfig:
assert result["viz_type"] == "table"
def test_map_table_config_row_limit(self) -> None:
"""Test that row_limit is mapped to form_data."""
config = TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product")],
row_limit=500,
)
result = map_table_config(config)
assert result["row_limit"] == 500
def test_map_table_config_default_row_limit(self) -> None:
"""Test that default row_limit is mapped to form_data."""
config = TableChartConfig(
chart_type="table",
columns=[ColumnRef(name="product", aggregate="SUM")],
)
result = map_table_config(config)
assert result["row_limit"] == 1000
class TestAddAdhocFilters:
"""Test _add_adhoc_filters helper function"""
def test_adds_filters_to_form_data(self) -> None:
"""Test that filters are correctly added to form_data."""
form_data: dict[str, Any] = {}
filters = [
FilterConfig(column="region", op="=", value="US"),
FilterConfig(column="year", op=">", value=2020),
]
_add_adhoc_filters(form_data, filters)
assert "adhoc_filters" in form_data
assert len(form_data["adhoc_filters"]) == 2
assert form_data["adhoc_filters"][0]["subject"] == "region"
assert form_data["adhoc_filters"][0]["operator"] == "=="
assert form_data["adhoc_filters"][1]["subject"] == "year"
assert form_data["adhoc_filters"][1]["operator"] == ">"
def test_no_filters_does_nothing(self) -> None:
"""Test that None filters leave form_data unchanged."""
form_data: dict[str, Any] = {"viz_type": "table"}
_add_adhoc_filters(form_data, None)
assert "adhoc_filters" not in form_data
def test_empty_list_does_nothing(self) -> None:
"""Test that empty filter list leaves form_data unchanged."""
form_data: dict[str, Any] = {"viz_type": "table"}
_add_adhoc_filters(form_data, [])
assert "adhoc_filters" not in form_data
class TestMapXYConfig:
"""Test map_xy_config function"""
@@ -538,6 +599,57 @@ class TestMapXYConfig:
assert result["stack"] == "Stack"
assert result["groupby"] == ["level"]
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_map_xy_config_with_filters(self, mock_is_temporal) -> None:
"""Test that filters are mapped to adhoc_filters in XY form_data."""
mock_is_temporal.return_value = True
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
kind="line",
filters=[FilterConfig(column="region", op="=", value="US")],
)
result = map_xy_config(config)
assert "adhoc_filters" in result
assert len(result["adhoc_filters"]) == 1
assert result["adhoc_filters"][0]["subject"] == "region"
assert result["adhoc_filters"][0]["operator"] == "=="
assert result["adhoc_filters"][0]["comparator"] == "US"
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_map_xy_config_row_limit(self, mock_is_temporal) -> None:
"""Test that row_limit is mapped to form_data."""
mock_is_temporal.return_value = True
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
kind="line",
row_limit=250,
)
result = map_xy_config(config)
assert result["row_limit"] == 250
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_map_xy_config_default_row_limit(self, mock_is_temporal) -> None:
"""Test that default row_limit is mapped to form_data."""
mock_is_temporal.return_value = True
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
kind="bar",
)
result = map_xy_config(config)
assert result["row_limit"] == 10000
class TestMapConfigToFormData:
"""Test map_config_to_form_data function"""

View File

@@ -512,6 +512,25 @@ class TestMixedTimeseriesChartConfigSchema:
unknown_field="bad",
)
def test_mixed_timeseries_default_row_limit(self) -> None:
config = MixedTimeseriesChartConfig(
chart_type="mixed_timeseries",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
y_secondary=[ColumnRef(name="orders", aggregate="COUNT")],
)
assert config.row_limit == 10000
def test_mixed_timeseries_custom_row_limit(self) -> None:
config = MixedTimeseriesChartConfig(
chart_type="mixed_timeseries",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
y_secondary=[ColumnRef(name="orders", aggregate="COUNT")],
row_limit=500,
)
assert config.row_limit == 500
# ============================================================
# Mixed Timeseries Form Data Mapping Tests
@@ -636,6 +655,35 @@ class TestMapMixedTimeseriesConfig:
assert result["y_axis_format_secondary"] == ",d"
assert result["logAxisSecondary"] is True
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_mixed_form_data_row_limit(self, mock_is_temporal) -> None:
mock_is_temporal.return_value = True
config = MixedTimeseriesChartConfig(
chart_type="mixed_timeseries",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
y_secondary=[ColumnRef(name="orders", aggregate="COUNT")],
row_limit=300,
)
result = map_mixed_timeseries_config(config, dataset_id=1)
assert result["row_limit"] == 300
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_mixed_form_data_default_row_limit(self, mock_is_temporal) -> None:
mock_is_temporal.return_value = True
config = MixedTimeseriesChartConfig(
chart_type="mixed_timeseries",
x=ColumnRef(name="date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
y_secondary=[ColumnRef(name="orders", aggregate="COUNT")],
)
result = map_mixed_timeseries_config(config, dataset_id=1)
assert result["row_limit"] == 10000
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_mixed_form_data_with_filters(self, mock_is_temporal) -> None:
mock_is_temporal.return_value = True