fix(mcp): add TEMPORAL_RANGE filter for temporal x-axis in generate_chart (#38978)

This commit is contained in:
Amin Ghadersohi
2026-03-31 18:39:08 +02:00
committed by GitHub
parent 4245720851
commit c37a3ec292
3 changed files with 166 additions and 7 deletions

View File

@@ -22,8 +22,10 @@ from unittest.mock import MagicMock, patch
import pytest
from superset.constants import NO_TIME_RANGE
from superset.mcp_service.chart.chart_utils import (
_add_adhoc_filters,
_ensure_temporal_adhoc_filter,
adhoc_filters_to_query_filters,
configure_temporal_handling,
create_metric_object,
@@ -44,7 +46,7 @@ from superset.mcp_service.chart.schemas import (
TableChartConfig,
XYChartConfig,
)
from superset.utils.core import GenericDataType
from superset.utils.core import FilterOperator, GenericDataType
class TestCreateMetricObject:
@@ -665,10 +667,13 @@ class TestMapXYConfig:
result = map_xy_config(config)
assert "adhoc_filters" in result
assert len(result["adhoc_filters"]) == 1
# User filter + auto-added TEMPORAL_RANGE filter for temporal x-axis
assert len(result["adhoc_filters"]) == 2
assert result["adhoc_filters"][0]["subject"] == "region"
assert result["adhoc_filters"][0]["operator"] == "=="
assert result["adhoc_filters"][0]["comparator"] == "US"
assert result["adhoc_filters"][1]["operator"] == "TEMPORAL_RANGE"
assert result["adhoc_filters"][1]["subject"] == "date"
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_map_xy_config_row_limit(self, mock_is_temporal) -> None:
@@ -1284,16 +1289,18 @@ class TestConfigureTemporalHandling:
"""Test configure_temporal_handling function"""
def test_temporal_column_with_time_grain(self) -> None:
"""Test temporal column sets time_grain_sqla"""
form_data: dict[str, Any] = {}
"""Test temporal column sets time_grain_sqla and granularity_sqla"""
form_data: dict[str, Any] = {"x_axis": "order_date"}
configure_temporal_handling(form_data, x_is_temporal=True, time_grain="P1M")
assert form_data["time_grain_sqla"] == "P1M"
assert form_data["granularity_sqla"] == "order_date"
def test_temporal_column_without_time_grain(self) -> None:
"""Test temporal column without time_grain doesn't set time_grain_sqla"""
form_data: dict[str, Any] = {}
"""Test temporal column sets granularity_sqla but not time_grain_sqla"""
form_data: dict[str, Any] = {"x_axis": "order_date"}
configure_temporal_handling(form_data, x_is_temporal=True, time_grain=None)
assert "time_grain_sqla" not in form_data
assert form_data["granularity_sqla"] == "order_date"
def test_non_temporal_column_sets_categorical_config(self) -> None:
"""Test non-temporal column sets categorical configuration"""
@@ -1355,7 +1362,16 @@ class TestMapXYConfigWithNonTemporalColumn:
assert result["x_axis"] == "created_at"
assert result["time_grain_sqla"] == "P1W"
assert result["granularity_sqla"] == "created_at"
assert "x_axis_sort_series_type" not in result
# Temporal x-axis should have a TEMPORAL_RANGE adhoc filter
temporal_filters = [
f
for f in result.get("adhoc_filters", [])
if f.get("operator") == "TEMPORAL_RANGE"
]
assert len(temporal_filters) == 1
assert temporal_filters[0]["subject"] == "created_at"
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_non_temporal_ignores_time_grain_param(self, mock_is_temporal) -> None:
@@ -1377,6 +1393,112 @@ class TestMapXYConfigWithNonTemporalColumn:
assert result["x_axis_sort_series_type"] == "name"
class TestEnsureTemporalAdhocFilter:
"""Test _ensure_temporal_adhoc_filter helper and its integration in map_xy_config"""
def test_adds_filter_to_empty_form_data(self) -> None:
"""Test adds TEMPORAL_RANGE filter when no adhoc_filters exist"""
form_data: dict[str, Any] = {}
_ensure_temporal_adhoc_filter(form_data, "order_date")
assert len(form_data["adhoc_filters"]) == 1
f = form_data["adhoc_filters"][0]
assert f["operator"] == FilterOperator.TEMPORAL_RANGE.value
assert f["subject"] == "order_date"
assert f["comparator"] == NO_TIME_RANGE
assert f["expressionType"] == "SIMPLE"
assert f["clause"] == "WHERE"
def test_appends_to_existing_filters(self) -> None:
"""Test appends temporal filter after existing user filters"""
form_data: dict[str, Any] = {
"adhoc_filters": [
{"subject": "region", "operator": "==", "comparator": "US"}
]
}
_ensure_temporal_adhoc_filter(form_data, "order_date")
assert len(form_data["adhoc_filters"]) == 2
assert form_data["adhoc_filters"][0]["subject"] == "region"
assert (
form_data["adhoc_filters"][1]["operator"]
== FilterOperator.TEMPORAL_RANGE.value
)
def test_does_not_duplicate_existing_temporal_filter(self) -> None:
"""Test skips adding if a TEMPORAL_RANGE filter already exists for the column"""
form_data: dict[str, Any] = {
"adhoc_filters": [
{
"subject": "order_date",
"operator": FilterOperator.TEMPORAL_RANGE.value,
"comparator": "Last 7 days",
}
]
}
_ensure_temporal_adhoc_filter(form_data, "order_date")
# Should still be just 1 filter (no duplicate)
assert len(form_data["adhoc_filters"]) == 1
def test_adds_filter_for_different_column(self) -> None:
"""Test adds filter when existing temporal filter is on a different column"""
form_data: dict[str, Any] = {
"adhoc_filters": [
{
"subject": "created_at",
"operator": FilterOperator.TEMPORAL_RANGE.value,
"comparator": NO_TIME_RANGE,
}
]
}
_ensure_temporal_adhoc_filter(form_data, "order_date")
assert len(form_data["adhoc_filters"]) == 2
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_temporal_x_axis_adds_filter_in_map_xy(self, mock_is_temporal) -> None:
"""Test map_xy_config adds TEMPORAL_RANGE filter for temporal x-axis"""
mock_is_temporal.return_value = True
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="order_date"),
y=[ColumnRef(name="revenue", aggregate="SUM")],
kind="bar",
)
result = map_xy_config(config, dataset_id=123)
temporal_filters = [
f
for f in result.get("adhoc_filters", [])
if f.get("operator") == FilterOperator.TEMPORAL_RANGE.value
]
assert len(temporal_filters) == 1
assert temporal_filters[0]["subject"] == "order_date"
assert temporal_filters[0]["comparator"] == NO_TIME_RANGE
@patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
def test_non_temporal_x_axis_no_temporal_filter(self, mock_is_temporal) -> None:
"""Test non-temporal x-axis skips TEMPORAL_RANGE filter"""
mock_is_temporal.return_value = False
config = XYChartConfig(
chart_type="xy",
x=ColumnRef(name="year"),
y=[ColumnRef(name="sales", aggregate="SUM")],
kind="bar",
)
result = map_xy_config(config, dataset_id=123)
temporal_filters = [
f
for f in result.get("adhoc_filters", [])
if f.get("operator") == FilterOperator.TEMPORAL_RANGE.value
]
assert len(temporal_filters) == 0
class TestFilterConfigValidation:
"""Test FilterConfig validation for new operators"""