From de0d78f2be0dac08f395f153ef12f2590c740fb6 Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Sat, 16 Apr 2022 20:57:47 -0700 Subject: [PATCH] feat: simplify SQLite time grain --- superset/db_engine_specs/sqlite.py | 26 ++++---- .../unit_tests/db_engine_specs/test_sqlite.py | 61 ++++++++++++++++++- 2 files changed, 74 insertions(+), 13 deletions(-) diff --git a/superset/db_engine_specs/sqlite.py b/superset/db_engine_specs/sqlite.py index 23512b3cb492f..06646f23039b8 100644 --- a/superset/db_engine_specs/sqlite.py +++ b/superset/db_engine_specs/sqlite.py @@ -42,20 +42,22 @@ class SqliteEngineSpec(BaseEngineSpec): "PT1S": "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:%S', {col}))", "PT1M": "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}))", "PT1H": "DATETIME(STRFTIME('%Y-%m-%dT%H:00:00', {col}))", - "P1D": "DATE({col})", - "P1W": "DATE({col}, -strftime('%w', {col}) || ' days')", - "P1M": "DATE({col}, -strftime('%d', {col}) || ' days', '+1 day')", + "P1D": "DATETIME({col}, 'start of day')", + "P1W": "DATETIME({col}, 'start of day', -strftime('%w', {col}) || ' days')", + "P1M": "DATETIME({col}, 'start of month')", "P3M": ( - "DATETIME(STRFTIME('%Y-', {col}) || " # year - "SUBSTR('00' || " # pad with zeros to 2 chars - "((CAST(STRFTIME('%m', {col}) AS INTEGER)) - " # month as integer - "(((CAST(STRFTIME('%m', {col}) AS INTEGER)) - 1) % 3)), " # month in quarter - "-2) || " # close pad - "'-01T00:00:00')" + "DATETIME({col}, 'start of month', " + "printf('-%d month', (strftime('%m', {col}) - 1) % 3))" + ), + "P1Y": "DATETIME({col}, 'start of year')", + "P1W/1970-01-03T00:00:00Z": "DATETIME({col}, 'start of day', 'weekday 6')", + "P1W/1970-01-04T00:00:00Z": "DATETIME({col}, 'start of day', 'weekday 0')", + "1969-12-28T00:00:00Z/P1W": ( + "DATETIME({col}, 'start of day', 'weekday 0', '-7 days')" + ), + "1969-12-29T00:00:00Z/P1W": ( + "DATETIME({col}, 'start of day', 'weekday 1', '-7 days')" ), - "P1Y": "DATETIME(STRFTIME('%Y-01-01T00:00:00', {col}))", - "P1W/1970-01-03T00:00:00Z": "DATE({col}, 'weekday 6')", - "1969-12-28T00:00:00Z/P1W": "DATE({col}, 'weekday 0', '-7 days')", } custom_errors: Dict[Pattern[str], Tuple[str, SupersetErrorType, Dict[str, Any]]] = { diff --git a/tests/unit_tests/db_engine_specs/test_sqlite.py b/tests/unit_tests/db_engine_specs/test_sqlite.py index 2ee8ea9a2c603..a8c2293ac7dab 100644 --- a/tests/unit_tests/db_engine_specs/test_sqlite.py +++ b/tests/unit_tests/db_engine_specs/test_sqlite.py @@ -14,11 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel, redefined-outer-name from datetime import datetime from unittest import mock import pytest from flask.ctx import AppContext +from sqlalchemy.engine import create_engine from tests.unit_tests.fixtures.common import dttm @@ -38,7 +40,7 @@ def test_convert_dttm_lower(app_context: AppContext, dttm: datetime) -> None: def test_convert_dttm_invalid_type(app_context: AppContext, dttm: datetime) -> None: from superset.db_engine_specs.sqlite import SqliteEngineSpec - assert SqliteEngineSpec.convert_dttm("other", dttm) == None + assert SqliteEngineSpec.convert_dttm("other", dttm) is None def test_get_all_datasource_names_table(app_context: AppContext) -> None: @@ -88,3 +90,60 @@ def test_get_all_datasource_names_invalid_type(app_context: AppContext) -> None: with pytest.raises(Exception): SqliteEngineSpec.get_all_datasource_names(database, invalid_type) + + +@pytest.mark.parametrize( + "dttm,grain,expected", + [ + ("2022-05-04T05:06:07.89Z", "PT1S", "2022-05-04 05:06:07"), + ("2022-05-04T05:06:07.89Z", "PT1M", "2022-05-04 05:06:00"), + ("2022-05-04T05:06:07.89Z", "PT1H", "2022-05-04 05:00:00"), + ("2022-05-04T05:06:07.89Z", "P1D", "2022-05-04 00:00:00"), + ("2022-05-04T05:06:07.89Z", "P1W", "2022-05-01 00:00:00"), + ("2022-05-04T05:06:07.89Z", "P1M", "2022-05-01 00:00:00"), + ("2022-05-04T05:06:07.89Z", "P1Y", "2022-01-01 00:00:00"), + # ___________________________ + # | May 2022 | + # |---------------------------| + # | S | M | T | W | T | F | S | + # |---+---+---+---+---+---+---| + # | 1 | 2 | 3 | 4 | 5 | 6 | 7 | + # --------------------------- + # week ending Saturday + ("2022-05-04T05:06:07.89Z", "P1W/1970-01-03T00:00:00Z", "2022-05-07 00:00:00"), + # week ending Sunday + ("2022-05-04T05:06:07.89Z", "P1W/1970-01-04T00:00:00Z", "2022-05-08 00:00:00"), + # week starting Sunday + ("2022-05-04T05:06:07.89Z", "1969-12-28T00:00:00Z/P1W", "2022-05-01 00:00:00"), + # week starting Monday + ("2022-05-04T05:06:07.89Z", "1969-12-29T00:00:00Z/P1W", "2022-05-02 00:00:00"), + # tests for quarter + ("2022-01-04T05:06:07.89Z", "P3M", "2022-01-01 00:00:00"), + ("2022-02-04T05:06:07.89Z", "P3M", "2022-01-01 00:00:00"), + ("2022-03-04T05:06:07.89Z", "P3M", "2022-01-01 00:00:00"), + ("2022-04-04T05:06:07.89Z", "P3M", "2022-04-01 00:00:00"), + ("2022-05-04T05:06:07.89Z", "P3M", "2022-04-01 00:00:00"), + ("2022-06-04T05:06:07.89Z", "P3M", "2022-04-01 00:00:00"), + ("2022-07-04T05:06:07.89Z", "P3M", "2022-07-01 00:00:00"), + ("2022-08-04T05:06:07.89Z", "P3M", "2022-07-01 00:00:00"), + ("2022-09-04T05:06:07.89Z", "P3M", "2022-07-01 00:00:00"), + ("2022-10-04T05:06:07.89Z", "P3M", "2022-10-01 00:00:00"), + ("2022-11-04T05:06:07.89Z", "P3M", "2022-10-01 00:00:00"), + ("2022-12-04T05:06:07.89Z", "P3M", "2022-10-01 00:00:00"), + ], +) +def test_time_grain_expressions( + dttm: str, grain: str, expected: str, app_context: AppContext +) -> None: + from superset.db_engine_specs.sqlite import SqliteEngineSpec + + engine = create_engine("sqlite://") + connection = engine.connect() + connection.execute("CREATE TABLE t (dttm DATETIME)") + connection.execute("INSERT INTO t VALUES (?)", dttm) + + # pylint: disable=protected-access + expression = SqliteEngineSpec._time_grain_expressions[grain].format(col="dttm") + sql = f"SELECT {expression} FROM t" + result = connection.execute(sql).scalar() + assert result == expected