Skip to content

Commit

Permalink
feat: simplify SQLite time grain (#19745)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored May 3, 2022
1 parent 99f1f9e commit 90130e4
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 13 deletions.
26 changes: 14 additions & 12 deletions superset/db_engine_specs/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]] = {
Expand Down
61 changes: 60 additions & 1 deletion tests/unit_tests/db_engine_specs/test_sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
Expand Down Expand Up @@ -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

0 comments on commit 90130e4

Please sign in to comment.