Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 1 addition & 5 deletions superset/db_engine_specs/db2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
69 changes: 69 additions & 0 deletions tests/unit_tests/db_engine_specs/test_db2.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@

import pytest # noqa: F401
from pytest_mock import MockerFixture
from sqlglot import parse_one

from superset.constants import TimeGrain
from superset.sql.parse import Table


Expand Down Expand Up @@ -78,3 +80,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 ",
Comment thread
ysinghc marked this conversation as resolved.
),
(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:
Comment thread
ysinghc marked this conversation as resolved.
"""
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"
)
Comment thread
ysinghc marked this conversation as resolved.
sql = f"SELECT {expression} FROM my_table" # noqa: S608

Copilot AI Oct 29, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] The # noqa: S608 comment suppresses a security warning about SQL injection, but this is a false positive since expression comes from a static dictionary value, not user input. However, consider using a regular string concatenation or format string without the noqa comment, as the warning suppression is unnecessary and may raise questions during code review. Alternatively, add a comment explaining why the suppression is safe here.

Copilot uses AI. Check for mistakes.

# This should not raise a ParseError
try:
parsed = parse_one(sql)
assert parsed is not None
except Exception as e:
pytest.fail(f"Failed to parse DAY time grain SQL: {e}")
Loading