diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py index d3801e66885..8c3ba134288 100644 --- a/superset/mcp_service/chart/chart_utils.py +++ b/superset/mcp_service/chart/chart_utils.py @@ -290,6 +290,10 @@ def map_xy_config(config: XYChartConfig) -> Dict[str, Any]: "metrics": metrics, } + # Add time grain if specified (for temporal x-axis columns) + if config.time_grain: + form_data["time_grain_sqla"] = config.time_grain + # CRITICAL FIX: For time series charts, handle groupby carefully to avoid duplicates # The x_axis field already tells Superset which column to use for time grouping groupby_columns = [] diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py index f28d7e75f69..48f3e45ba16 100644 --- a/superset/mcp_service/chart/schemas.py +++ b/superset/mcp_service/chart/schemas.py @@ -36,6 +36,7 @@ from pydantic import ( PositiveInt, ) +from superset.constants import TimeGrain from superset.daos.base import ColumnOperator, ColumnOperatorEnum from superset.mcp_service.common.cache_schemas import ( CacheStatus, @@ -678,6 +679,15 @@ class XYChartConfig(BaseModel): kind: Literal["line", "bar", "area", "scatter"] = Field( "line", description="Chart visualization type" ) + time_grain: TimeGrain | None = Field( + None, + description=( + "Time granularity for the x-axis when it's a temporal column. " + "Common values: PT1S (second), PT1M (minute), PT1H (hour), " + "P1D (day), P1W (week), P1M (month), P3M (quarter), P1Y (year). " + "If not specified, Superset will use its default behavior." + ), + ) group_by: ColumnRef | None = Field(None, description="Column to group by") x_axis: AxisConfig | None = Field(None, description="X-axis configuration") y_axis: AxisConfig | None = Field(None, description="Y-axis configuration") 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 4f4f3a3857b..f2e9d966c89 100644 --- a/tests/unit_tests/mcp_service/chart/test_chart_utils.py +++ b/tests/unit_tests/mcp_service/chart/test_chart_utils.py @@ -275,6 +275,82 @@ class TestMapXYConfig: assert result["show_legend"] is False assert result["legend_orientation"] == "top" + def test_map_xy_config_with_time_grain_month(self) -> None: + """Test XY config mapping with monthly time grain""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="order_date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + kind="bar", + time_grain="P1M", + ) + + result = map_xy_config(config) + + assert result["viz_type"] == "echarts_timeseries_bar" + assert result["x_axis"] == "order_date" + assert result["time_grain_sqla"] == "P1M" + + def test_map_xy_config_with_time_grain_day(self) -> None: + """Test XY config mapping with daily time grain""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="created_at"), + y=[ColumnRef(name="count", aggregate="COUNT")], + kind="line", + time_grain="P1D", + ) + + result = map_xy_config(config) + + assert result["viz_type"] == "echarts_timeseries_line" + assert result["x_axis"] == "created_at" + assert result["time_grain_sqla"] == "P1D" + + def test_map_xy_config_with_time_grain_hour(self) -> None: + """Test XY config mapping with hourly time grain""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="timestamp"), + y=[ColumnRef(name="requests", aggregate="SUM")], + kind="area", + time_grain="PT1H", + ) + + result = map_xy_config(config) + + assert result["time_grain_sqla"] == "PT1H" + + def test_map_xy_config_without_time_grain(self) -> None: + """Test XY config mapping without time grain (should not set time_grain_sqla)""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="date"), + y=[ColumnRef(name="revenue")], + kind="line", + ) + + result = map_xy_config(config) + + assert "time_grain_sqla" not in result + + def test_map_xy_config_with_time_grain_and_groupby(self) -> None: + """Test XY config mapping with time grain and group by""" + config = XYChartConfig( + chart_type="xy", + x=ColumnRef(name="order_date"), + y=[ColumnRef(name="revenue", aggregate="SUM")], + kind="line", + time_grain="P1W", + group_by=ColumnRef(name="category"), + ) + + result = map_xy_config(config) + + assert result["time_grain_sqla"] == "P1W" + assert result["groupby"] == ["category"] + assert result["x_axis"] == "order_date" + class TestMapConfigToFormData: """Test map_config_to_form_data function"""