diff --git a/superset/db_engine_specs/db2.py b/superset/db_engine_specs/db2.py index 8a09f5af3c13..eec1e2a0f2a4 100644 --- a/superset/db_engine_specs/db2.py +++ b/superset/db_engine_specs/db2.py @@ -48,11 +48,7 @@ class Db2EngineSpec(BaseEngineSpec): " - MINUTE({col}) MINUTES" " - SECOND({col}) SECONDS" " - MICROSECOND({col}) MICROSECONDS ", - TimeGrain.DAY: "CAST({col} as TIMESTAMP)" - " - HOUR({col}) HOURS" - " - MINUTE({col}) MINUTES" - " - SECOND({col}) SECONDS" - " - MICROSECOND({col}) MICROSECONDS", + TimeGrain.DAY: "DATE({col})", TimeGrain.WEEK: "{col} - (DAYOFWEEK({col})) DAYS", TimeGrain.MONTH: "{col} - (DAY({col})-1) DAYS", TimeGrain.QUARTER: "{col} - (DAY({col})-1) DAYS" diff --git a/tests/unit_tests/db_engine_specs/test_db2.py b/tests/unit_tests/db_engine_specs/test_db2.py index 9e500ddc0c53..1ae4471f9ff7 100644 --- a/tests/unit_tests/db_engine_specs/test_db2.py +++ b/tests/unit_tests/db_engine_specs/test_db2.py @@ -17,7 +17,10 @@ import pytest # noqa: F401 from pytest_mock import MockerFixture +from sqlglot import parse_one +from sqlglot.errors import ParseError +from superset.constants import TimeGrain from superset.sql.parse import Table @@ -78,3 +81,70 @@ def test_get_prequeries(mocker: MockerFixture) -> None: assert Db2EngineSpec.get_prequeries(database, schema="my_schema") == [ 'set current_schema "my_schema"' ] + + +@pytest.mark.parametrize( + ("grain", "expected_expression"), + [ + (None, "my_col"), + ( + TimeGrain.SECOND, + "CAST(my_col as TIMESTAMP) - MICROSECOND(my_col) MICROSECONDS", + ), + ( + TimeGrain.MINUTE, + "CAST(my_col as TIMESTAMP)" + " - SECOND(my_col) SECONDS - MICROSECOND(my_col) MICROSECONDS", + ), + ( + TimeGrain.HOUR, + "CAST(my_col as TIMESTAMP)" + " - MINUTE(my_col) MINUTES" + " - SECOND(my_col) SECONDS - MICROSECOND(my_col) MICROSECONDS ", + ), + (TimeGrain.DAY, "DATE(my_col)"), + (TimeGrain.WEEK, "my_col - (DAYOFWEEK(my_col)) DAYS"), + (TimeGrain.MONTH, "my_col - (DAY(my_col)-1) DAYS"), + ( + TimeGrain.QUARTER, + "my_col - (DAY(my_col)-1) DAYS" + " - (MONTH(my_col)-1) MONTHS + ((QUARTER(my_col)-1) * 3) MONTHS", + ), + ( + TimeGrain.YEAR, + "my_col - (DAY(my_col)-1) DAYS - (MONTH(my_col)-1) MONTHS", + ), + ], +) +def test_time_grain_expressions(grain: TimeGrain, expected_expression: str) -> None: + """ + Test that time grain expressions generate the expected SQL. + """ + from superset.db_engine_specs.db2 import Db2EngineSpec + + actual = Db2EngineSpec._time_grain_expressions[grain].format(col="my_col") + assert actual == expected_expression + + +def test_time_grain_day_parseable() -> None: + """ + Test that the DAY time grain expression generates valid SQL + that can be parsed by sqlglot. + + This test addresses the bug where the previous expression + "CAST({col} as TIMESTAMP) - HOUR({col}) HOURS - ..." + could not be parsed by sqlglot. + """ + from superset.db_engine_specs.db2 import Db2EngineSpec + + expression = Db2EngineSpec._time_grain_expressions[TimeGrain.DAY].format( + col="my_timestamp_col", + ) + sql = f"SELECT {expression} FROM my_table" # noqa: S608 + + # This should not raise a ParseError + try: + parsed = parse_one(sql) + assert parsed is not None + except ParseError as e: + pytest.fail(f"Failed to parse DAY time grain SQL: {e}")