You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/07/10 20:48:16 UTC
(superset) branch master updated: fix(mcp): bind Big Number chart to a temporal column for dashboard time filters (#41895)
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new bd61e09b4ea fix(mcp): bind Big Number chart to a temporal column for dashboard time filters (#41895)
bd61e09b4ea is described below
commit bd61e09b4ea1b0391ec67d62311fe874e8231542
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Fri Jul 10 16:48:09 2026 -0400
fix(mcp): bind Big Number chart to a temporal column for dashboard time filters (#41895)
---
superset/mcp_service/chart/chart_utils.py | 112 ++++++++++++++-----
superset/mcp_service/chart/plugins/big_number.py | 2 +-
superset/mcp_service/chart/schemas.py | 6 +-
.../mcp_service/chart/test_big_number_chart.py | 121 +++++++++++++++++++++
.../mcp_service/chart/test_chart_utils.py | 84 ++++++++++----
5 files changed, 272 insertions(+), 53 deletions(-)
diff --git a/superset/mcp_service/chart/chart_utils.py b/superset/mcp_service/chart/chart_utils.py
index 62a3f35f544..55e07b798e8 100644
--- a/superset/mcp_service/chart/chart_utils.py
+++ b/superset/mcp_service/chart/chart_utils.py
@@ -25,7 +25,10 @@ generation that can be used by both generate_chart and generate_explore_link too
import hashlib
import logging
from dataclasses import dataclass
-from typing import Any, Dict
+from typing import Any, Dict, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from superset.connectors.sqla.models import SqlaTable
from superset.constants import NO_TIME_RANGE
from superset.mcp_service.chart.schemas import (
@@ -266,7 +269,28 @@ def generate_explore_link(
return f"{base_url}/explore/?datasource_type=table&datasource_id={dataset_id}"
-def is_column_truly_temporal(column_name: str, dataset_id: int | str | None) -> bool:
+def _find_dataset_by_id_or_uuid(dataset_id: int | str | None) -> "SqlaTable | None":
+ """Look up a dataset by numeric ID or UUID string.
+
+ Shared by callers that resolve a dataset from the ``dataset_id | str | None``
+ shape accepted throughout this module; each caller decides its own behavior
+ for a missing dataset_id or dataset (raise vs. default vs. None).
+
+ Delegates to ``DatasetDAO.find_by_id_or_uuid`` (also used by the dataset
+ API) instead of reimplementing the id/uuid dispatch here.
+ """
+ if not dataset_id:
+ return None
+ from superset.daos.dataset import DatasetDAO # avoid circular import
+
+ return DatasetDAO.find_by_id_or_uuid(str(dataset_id))
+
+
+def is_column_truly_temporal(
+ column_name: str,
+ dataset_id: int | str | None,
+ dataset: "SqlaTable | None" = None,
+) -> bool:
"""
Check if a column is truly temporal based on its SQL data type.
@@ -280,24 +304,20 @@ def is_column_truly_temporal(column_name: str, dataset_id: int | str | None) ->
Args:
column_name: Name of the column to check
dataset_id: Dataset ID to look up column metadata
+ dataset: Optional pre-fetched dataset, reused as-is to avoid a
+ redundant DAO lookup when the caller already resolved it.
Returns:
True if the column has a real temporal SQL type, False otherwise
"""
- from superset.daos.dataset import DatasetDAO
from superset.utils.core import GenericDataType
- if not dataset_id:
+ if not dataset_id and dataset is None:
return True # Default to temporal if we can't check (backward compatible)
try:
- # Find dataset
- if isinstance(dataset_id, int) or (
- isinstance(dataset_id, str) and dataset_id.isdigit()
- ):
- dataset = DatasetDAO.find_by_id(int(dataset_id))
- else:
- dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
+ if dataset is None:
+ dataset = _find_dataset_by_id_or_uuid(dataset_id)
if not dataset:
return True # Default to temporal if dataset not found
@@ -717,21 +737,21 @@ def _ensure_temporal_adhoc_filter(form_data: Dict[str, Any], column: str) -> Non
def _resolve_default_x_axis(
config: XYChartConfig, dataset_id: int | str | None
-) -> XYChartConfig:
- """Resolve x-axis to the dataset's main_dttm_col when x is omitted."""
+) -> tuple[XYChartConfig, "SqlaTable | None"]:
+ """Resolve x-axis to the dataset's main_dttm_col when x is omitted.
+
+ Returns the (possibly updated) config alongside the dataset fetched while
+ resolving the default, if any, so callers can reuse it (e.g. passing it
+ into is_column_truly_temporal) instead of re-querying DatasetDAO for the
+ same dataset_id.
+ """
if config.x is not None:
- return config
+ return config, None
if not dataset_id:
raise ValueError("x-axis column is required when dataset_id is not provided")
- from superset.daos.dataset import DatasetDAO
- if isinstance(dataset_id, int) or (
- isinstance(dataset_id, str) and dataset_id.isdigit()
- ):
- dataset = DatasetDAO.find_by_id(int(dataset_id))
- else:
- dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
+ dataset = _find_dataset_by_id_or_uuid(dataset_id)
if not dataset or not dataset.main_dttm_col:
raise ValueError(
@@ -741,7 +761,10 @@ def _resolve_default_x_axis(
)
from superset.mcp_service.chart.schemas import ColumnRef
- return config.model_copy(update={"x": ColumnRef(name=dataset.main_dttm_col)})
+ return (
+ config.model_copy(update={"x": ColumnRef(name=dataset.main_dttm_col)}),
+ dataset,
+ )
def _add_xy_limits(form_data: Dict[str, Any], config: XYChartConfig) -> None:
@@ -759,14 +782,17 @@ def map_xy_config( # noqa: C901
raise ValueError("XY chart must have at least one Y-axis metric")
# Resolve x-axis default: use dataset's main_dttm_col when x is omitted.
- config = _resolve_default_x_axis(config, dataset_id)
+ config, resolved_dataset = _resolve_default_x_axis(config, dataset_id)
# ``_resolve_default_x_axis`` guarantees x is set.
if config.x is None or config.x.name is None:
raise ValueError("XY chart requires an x-axis with a resolvable column name")
- # Check if x-axis column is truly temporal (based on actual SQL type)
- x_is_temporal = is_column_truly_temporal(config.x.name, dataset_id)
+ # Check if x-axis column is truly temporal (based on actual SQL type).
+ # Reuse the dataset fetched above (if any) to avoid a second DAO lookup.
+ x_is_temporal = is_column_truly_temporal(
+ config.x.name, dataset_id, dataset=resolved_dataset
+ )
# Map chart kind to viz_type - always use the same viz types
# The temporal vs non-temporal handling is done via form_data configuration
@@ -865,7 +891,9 @@ def map_pie_config(config: PieChartConfig) -> Dict[str, Any]:
return form_data
-def map_big_number_config(config: BigNumberChartConfig) -> Dict[str, Any]:
+def map_big_number_config(
+ config: BigNumberChartConfig, dataset_id: int | str | None = None
+) -> Dict[str, Any]:
"""Map big number chart config to Superset form_data."""
# Determine viz_type: big_number (with trendline) or big_number_total
if config.show_trendline and config.temporal_column:
@@ -911,9 +939,41 @@ def map_big_number_config(config: BigNumberChartConfig) -> Dict[str, Any]:
_add_adhoc_filters(form_data, config.filters)
+ # Bind a TEMPORAL_RANGE adhoc filter so dashboard time-range filters have
+ # a column to apply to — mirrors the Explore UI's `BigNumberTotal` control
+ # panel, which exposes an `adhoc_filters` control even though there's no
+ # dedicated time-column control for the total variant.
+ if temporal_column := _resolve_big_number_temporal_column(config, dataset_id):
+ _ensure_temporal_adhoc_filter(form_data, temporal_column)
+
return form_data
+def _resolve_big_number_temporal_column(
+ config: BigNumberChartConfig, dataset_id: int | str | None
+) -> str | None:
+ """Resolve the column to bind a Big Number's TEMPORAL_RANGE filter to.
+
+ Falls back to the dataset's main_dttm_col when the caller didn't specify
+ temporal_column, and guards the result with is_column_truly_temporal (same
+ check map_xy_config applies to its x-axis) so a non-temporal column never
+ gets a TEMPORAL_RANGE filter. The dataset is fetched at most once here and
+ reused by is_column_truly_temporal instead of letting it re-query by
+ dataset_id.
+ """
+ dataset = None
+ if not config.temporal_column:
+ dataset = _find_dataset_by_id_or_uuid(dataset_id)
+ temporal_column = config.temporal_column or (
+ dataset.main_dttm_col if dataset else None
+ )
+ if temporal_column and is_column_truly_temporal(
+ temporal_column, dataset_id, dataset=dataset
+ ):
+ return temporal_column
+ return None
+
+
def map_handlebars_config(config: HandlebarsChartConfig) -> Dict[str, Any]:
"""Map handlebars chart config to Superset form_data."""
form_data: Dict[str, Any] = {
diff --git a/superset/mcp_service/chart/plugins/big_number.py b/superset/mcp_service/chart/plugins/big_number.py
index 8c847c6ba8f..9a861883d22 100644
--- a/superset/mcp_service/chart/plugins/big_number.py
+++ b/superset/mcp_service/chart/plugins/big_number.py
@@ -151,7 +151,7 @@ class BigNumberChartPlugin(BaseChartPlugin):
def to_form_data(
self, config: Any, dataset_id: int | str | None = None
) -> dict[str, Any]:
- return map_big_number_config(config)
+ return map_big_number_config(config, dataset_id=dataset_id)
def post_map_validate(
self,
diff --git a/superset/mcp_service/chart/schemas.py b/superset/mcp_service/chart/schemas.py
index 623bec8aedf..1b5e3b22ddb 100644
--- a/superset/mcp_service/chart/schemas.py
+++ b/superset/mcp_service/chart/schemas.py
@@ -1405,8 +1405,10 @@ class BigNumberChartConfig(UnknownFieldCheckMixin):
temporal_column: str | None = Field(
None,
description=(
- "Temporal column for the trendline x-axis. "
- "Required when show_trendline is True."
+ "Temporal column for the trendline x-axis. Required when "
+ "show_trendline is True. Also used (whether or not a trendline is "
+ "shown) to bind the chart's dashboard time-range filter; when "
+ "omitted, the dataset's main temporal column is used instead."
),
min_length=1,
max_length=255,
diff --git a/tests/unit_tests/mcp_service/chart/test_big_number_chart.py b/tests/unit_tests/mcp_service/chart/test_big_number_chart.py
index a19f22282e4..cbf255ac784 100644
--- a/tests/unit_tests/mcp_service/chart/test_big_number_chart.py
+++ b/tests/unit_tests/mcp_service/chart/test_big_number_chart.py
@@ -17,6 +17,8 @@
"""Tests for Big Number chart type support in MCP service."""
+from unittest.mock import MagicMock, patch
+
import pytest
from pydantic import ValidationError
@@ -357,6 +359,104 @@ class TestMapBigNumberConfig:
assert "time_grain_sqla" not in form_data
assert "start_y_axis_at_zero" not in form_data
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
+ def test_total_binds_dataset_main_dttm_col_for_dashboard_filters(
+ self, mock_find_by_id_or_uuid: MagicMock
+ ) -> None:
+ """Big Number Total falls back to dataset.main_dttm_col so dashboard
+ time-range filters have a column to bind to (regression test for
+ charts created via MCP not responding to dashboard filters)."""
+ mock_dataset = MagicMock()
+ mock_dataset.main_dttm_col = "order_date"
+ mock_dataset.columns = []
+ mock_find_by_id_or_uuid.return_value = mock_dataset
+
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ )
+ form_data = map_big_number_config(config, dataset_id=42)
+
+ assert form_data["viz_type"] == "big_number_total"
+ assert "adhoc_filters" in form_data
+ assert len(form_data["adhoc_filters"]) == 1
+ assert form_data["adhoc_filters"][0]["subject"] == "order_date"
+ assert form_data["adhoc_filters"][0]["operator"] == "TEMPORAL_RANGE"
+ # Regression test: the main_dttm_col fallback and the
+ # is_column_truly_temporal guard must share a single dataset lookup
+ # rather than each re-querying DatasetDAO for the same dataset_id.
+ mock_find_by_id_or_uuid.assert_called_once_with("42")
+
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
+ def test_total_no_dataset_main_dttm_col_skips_temporal_filter(
+ self, mock_find_by_id_or_uuid: MagicMock
+ ) -> None:
+ """When the dataset has no temporal column, no TEMPORAL_RANGE filter
+ can be added — there's nothing for a dashboard filter to bind to."""
+ mock_dataset = MagicMock()
+ mock_dataset.main_dttm_col = None
+ mock_find_by_id_or_uuid.return_value = mock_dataset
+
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ )
+ form_data = map_big_number_config(config, dataset_id=42)
+
+ assert "adhoc_filters" not in form_data
+
+ @patch("superset.mcp_service.chart.chart_utils.is_column_truly_temporal")
+ def test_total_non_temporal_column_skips_temporal_filter(
+ self, mock_is_temporal: MagicMock
+ ) -> None:
+ """Mirrors TestMapXYConfig.test_non_temporal_x_axis_no_temporal_filter:
+ when is_column_truly_temporal says an explicit temporal_column isn't
+ really temporal, Big Number Total must not bind a TEMPORAL_RANGE
+ filter to it either.
+ test_total_binds_dataset_main_dttm_col_for_dashboard_filters only
+ exercises the "column not found -> default True" branch (empty
+ columns list); this covers the "found a non-temporal column -> False"
+ branch that actually suppresses the filter."""
+ mock_is_temporal.return_value = False
+
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ temporal_column="year",
+ )
+ form_data = map_big_number_config(config)
+
+ assert "adhoc_filters" not in form_data
+
+ def test_explicit_temporal_column_used_without_trendline(self) -> None:
+ """An explicit temporal_column binds the dashboard filter even when
+ show_trendline is False, without needing a dataset lookup."""
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ temporal_column="order_date",
+ )
+ form_data = map_big_number_config(config)
+
+ assert form_data["viz_type"] == "big_number_total"
+ assert form_data["adhoc_filters"][0]["subject"] == "order_date"
+ assert form_data["adhoc_filters"][0]["operator"] == "TEMPORAL_RANGE"
+
+ def test_trendline_also_gets_temporal_adhoc_filter(self) -> None:
+ """Trendline charts get both granularity_sqla and a TEMPORAL_RANGE
+ adhoc filter, matching either binding mechanism the dashboard uses."""
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ temporal_column="order_date",
+ show_trendline=True,
+ )
+ form_data = map_big_number_config(config)
+
+ assert form_data["granularity_sqla"] == "order_date"
+ assert form_data["adhoc_filters"][0]["subject"] == "order_date"
+ assert form_data["adhoc_filters"][0]["operator"] == "TEMPORAL_RANGE"
+
def test_with_aggregation_sum(self) -> None:
"""aggregation='sum' is written to form_data for trendline charts."""
config = BigNumberChartConfig(
@@ -425,6 +525,27 @@ class TestMapConfigToFormDataBigNumber:
form_data = map_config_to_form_data(config)
assert form_data["viz_type"] == "big_number"
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
+ def test_dispatch_forwards_dataset_id_for_dashboard_filter_binding(
+ self, mock_find_by_id_or_uuid: MagicMock
+ ) -> None:
+ """Regression test for BigNumberChartPlugin.to_form_data() dropping
+ dataset_id: the dispatcher must forward it through so the dataset's
+ main_dttm_col fallback in map_big_number_config() actually runs."""
+ mock_dataset = MagicMock()
+ mock_dataset.main_dttm_col = "order_date"
+ mock_dataset.columns = []
+ mock_find_by_id_or_uuid.return_value = mock_dataset
+
+ config = BigNumberChartConfig(
+ chart_type="big_number",
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ )
+ form_data = map_config_to_form_data(config, dataset_id=42)
+
+ assert form_data["adhoc_filters"][0]["subject"] == "order_date"
+ assert form_data["adhoc_filters"][0]["operator"] == "TEMPORAL_RANGE"
+
class TestGenerateChartNameBigNumber:
"""Test generate_chart_name for big number configs."""
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 190b25b3e90..473a88a76eb 100644
--- a/tests/unit_tests/mcp_service/chart/test_chart_utils.py
+++ b/tests/unit_tests/mcp_service/chart/test_chart_utils.py
@@ -49,7 +49,7 @@ from superset.mcp_service.chart.schemas import (
TableChartConfig,
XYChartConfig,
)
-from superset.utils.core import FilterOperator, GenericDataType
+from superset.utils.core import ColumnSpec, FilterOperator, GenericDataType
class TestGetTableChartTypeLabel:
@@ -932,14 +932,54 @@ class TestMapConfigToFormData:
"superset.mcp_service.chart.chart_utils.is_column_truly_temporal",
return_value=True,
)
- @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
def test_map_xy_config_x_none_defaults_to_main_dttm_col(
- self, mock_find_by_id: Any, mock_is_temporal: Any
+ self, mock_find_by_id_or_uuid: Any, mock_is_temporal: Any
) -> None:
"""When x is None, map_xy_config resolves it from dataset.main_dttm_col."""
mock_dataset = MagicMock()
mock_dataset.main_dttm_col = "order_date"
- mock_find_by_id.return_value = mock_dataset
+ mock_find_by_id_or_uuid.return_value = mock_dataset
+
+ config = XYChartConfig(
+ chart_type="xy",
+ x=None,
+ y=[ColumnRef(name="revenue", aggregate="SUM")],
+ kind="bar",
+ )
+
+ result = map_xy_config(config, dataset_id=42)
+
+ assert result["x_axis"] == "order_date"
+ mock_find_by_id_or_uuid.assert_called_once_with("42")
+
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
+ def test_map_xy_config_x_none_reuses_dataset_for_temporal_check(
+ self, mock_find_by_id_or_uuid: Any
+ ) -> None:
+ """The main_dttm_col fallback and the is_column_truly_temporal guard
+ must share a single dataset lookup rather than each re-querying
+ DatasetDAO for the same dataset_id (mirrors the dedup fix applied to
+ map_big_number_config's _resolve_big_number_temporal_column)."""
+
+ mock_column = MagicMock()
+ mock_column.column_name = "order_date"
+ mock_column.type = "TIMESTAMP"
+
+ mock_db_engine_spec = MagicMock()
+ mock_db_engine_spec.get_column_spec.return_value = ColumnSpec(
+ sqla_type=MagicMock(),
+ generic_type=GenericDataType.TEMPORAL,
+ is_dttm=False,
+ )
+ mock_database = MagicMock()
+ mock_database.db_engine_spec = mock_db_engine_spec
+
+ mock_dataset = MagicMock()
+ mock_dataset.main_dttm_col = "order_date"
+ mock_dataset.columns = [mock_column]
+ mock_dataset.database = mock_database
+ mock_find_by_id_or_uuid.return_value = mock_dataset
config = XYChartConfig(
chart_type="xy",
@@ -951,7 +991,7 @@ class TestMapConfigToFormData:
result = map_xy_config(config, dataset_id=42)
assert result["x_axis"] == "order_date"
- mock_find_by_id.assert_called_once_with(42)
+ mock_find_by_id_or_uuid.assert_called_once_with("42")
def test_map_xy_config_x_none_no_dataset_id_raises(self) -> None:
"""When x is None and no dataset_id, raise ValueError."""
@@ -965,14 +1005,14 @@ class TestMapConfigToFormData:
with pytest.raises(ValueError, match="x-axis column is required"):
map_xy_config(config, dataset_id=None)
- @patch("superset.daos.dataset.DatasetDAO.find_by_id")
+ @patch("superset.daos.dataset.DatasetDAO.find_by_id_or_uuid")
def test_map_xy_config_x_none_no_main_dttm_col_raises(
- self, mock_find_by_id: Any
+ self, mock_find_by_id_or_uuid: Any
) -> None:
"""When x is None and dataset has no main_dttm_col, raise ValueError."""
mock_dataset = MagicMock()
mock_dataset.main_dttm_col = None
- mock_find_by_id.return_value = mock_dataset
+ mock_find_by_id_or_uuid.return_value = mock_dataset
config = XYChartConfig(
chart_type="xy",
@@ -1383,10 +1423,6 @@ class TestIsColumnTrulyTemporal:
generic_type: GenericDataType,
):
"""Helper to create a mock dataset with proper db_engine_spec"""
- from unittest.mock import MagicMock
-
- from superset.utils.core import ColumnSpec
-
mock_column = MagicMock()
mock_column.column_name = column_name
mock_column.type = column_type
@@ -1414,7 +1450,7 @@ class TestIsColumnTrulyTemporal:
@patch("superset.daos.dataset.DatasetDAO")
def test_returns_true_when_dataset_not_found(self, mock_dao) -> None:
"""Test returns True when dataset is not found"""
- mock_dao.find_by_id.return_value = None
+ mock_dao.find_by_id_or_uuid.return_value = None
result = is_column_truly_temporal("year", 123)
assert result is True
@@ -1424,7 +1460,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"year", "BIGINT", GenericDataType.NUMERIC
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("year", 123)
assert result is False
@@ -1435,7 +1471,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"month", "INTEGER", GenericDataType.NUMERIC
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("month", 123)
assert result is False
@@ -1446,7 +1482,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"created_at", "TIMESTAMP", GenericDataType.TEMPORAL
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("created_at", 123)
assert result is True
@@ -1457,7 +1493,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"order_date", "DATE", GenericDataType.TEMPORAL
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("order_date", 123)
assert result is True
@@ -1468,7 +1504,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"Year", "BIGINT", GenericDataType.NUMERIC
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("year", 123)
assert result is False
@@ -1476,14 +1512,14 @@ class TestIsColumnTrulyTemporal:
@patch("superset.daos.dataset.DatasetDAO")
def test_returns_true_on_value_error(self, mock_dao) -> None:
"""Test returns True (default) when ValueError occurs"""
- mock_dao.find_by_id.side_effect = ValueError("Invalid ID")
+ mock_dao.find_by_id_or_uuid.side_effect = ValueError("Invalid ID")
result = is_column_truly_temporal("year", 123)
assert result is True
@patch("superset.daos.dataset.DatasetDAO")
def test_returns_true_on_attribute_error(self, mock_dao) -> None:
"""Test returns True (default) when AttributeError occurs"""
- mock_dao.find_by_id.side_effect = AttributeError("Missing attribute")
+ mock_dao.find_by_id_or_uuid.side_effect = AttributeError("Missing attribute")
result = is_column_truly_temporal("year", 123)
assert result is True
@@ -1493,11 +1529,11 @@ class TestIsColumnTrulyTemporal:
mock_dataset = self._create_mock_dataset(
"year", "BIGINT", GenericDataType.NUMERIC
)
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("year", "abc-123-uuid")
assert result is False
- mock_dao.find_by_id.assert_called_with("abc-123-uuid", id_column="uuid")
+ mock_dao.find_by_id_or_uuid.assert_called_with("abc-123-uuid")
@patch("superset.daos.dataset.DatasetDAO")
def test_falls_back_to_is_dttm_when_no_column_spec(self, mock_dao) -> None:
@@ -1518,7 +1554,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = MagicMock()
mock_dataset.columns = [mock_column]
mock_dataset.database = mock_database
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("year", 123)
assert result is False
@@ -1535,7 +1571,7 @@ class TestIsColumnTrulyTemporal:
mock_dataset = MagicMock()
mock_dataset.columns = [mock_column]
- mock_dao.find_by_id.return_value = mock_dataset
+ mock_dao.find_by_id_or_uuid.return_value = mock_dataset
result = is_column_truly_temporal("year", 123)
assert result is True