fix(db2): update time grain expressions for DAY to use DATE function (#35848)

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}")