From b205ce32b0b5ead6ba98585ba22e0f9c6d3cebdd Mon Sep 17 00:00:00 2001 From: Ville Brofeldt <33317356+villebro@users.noreply.github.com> Date: Thu, 25 Jun 2020 12:18:37 +0300 Subject: [PATCH] feat: Add new timegrains and convert_dttm to Druid engine spec (#10160) * feat: Add new timegrains and convert_dttm to Druid engine spec * Add TemporalType enum and fix test case * Remove DATETIME for athena (original spec) --- superset/db_engine_specs/athena.py | 5 +- superset/db_engine_specs/bigquery.py | 9 ++-- superset/db_engine_specs/clickhouse.py | 5 +- superset/db_engine_specs/drill.py | 5 +- superset/db_engine_specs/druid.py | 16 ++++++- superset/db_engine_specs/elasticsearch.py | 3 +- superset/db_engine_specs/hana.py | 5 +- superset/db_engine_specs/hive.py | 4 +- superset/db_engine_specs/impala.py | 5 +- superset/db_engine_specs/kylin.py | 5 +- superset/db_engine_specs/mssql.py | 7 +-- superset/db_engine_specs/mysql.py | 5 +- superset/db_engine_specs/oracle.py | 7 +-- superset/db_engine_specs/postgres.py | 5 +- superset/db_engine_specs/presto.py | 4 +- superset/db_engine_specs/snowflake.py | 7 +-- superset/db_engine_specs/sqlite.py | 3 +- superset/utils/core.py | 13 ++++++ tests/db_engine_specs/druid_tests.py | 56 +++++++++++++++++++++++ 19 files changed, 133 insertions(+), 36 deletions(-) create mode 100644 tests/db_engine_specs/druid_tests.py diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py index bbdba14b3f..eaf76e858d 100644 --- a/superset/db_engine_specs/athena.py +++ b/superset/db_engine_specs/athena.py @@ -18,6 +18,7 @@ from datetime import datetime from typing import Optional from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class AthenaEngineSpec(BaseEngineSpec): @@ -42,9 +43,9 @@ class AthenaEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"from_iso8601_date('{dttm.date().isoformat()}')" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""from_iso8601_timestamp('{dttm.isoformat(timespec="microseconds")}')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 3091d65c44..f45145f3bd 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -24,6 +24,7 @@ from sqlalchemy import literal_column from sqlalchemy.sql.expression import ColumnClause from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils if TYPE_CHECKING: # pylint: disable=unused-import @@ -72,13 +73,13 @@ class BigQueryEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"CAST('{dttm.date().isoformat()}' AS DATE)" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS DATETIME)""" - if tt == "TIME": + if tt == utils.TemporalType.TIME: return f"""CAST('{dttm.strftime("%H:%M:%S.%f")}' AS TIME)""" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS TIMESTAMP)""" return None diff --git a/superset/db_engine_specs/clickhouse.py b/superset/db_engine_specs/clickhouse.py index 3b390530b6..f02ef51aba 100644 --- a/superset/db_engine_specs/clickhouse.py +++ b/superset/db_engine_specs/clickhouse.py @@ -18,6 +18,7 @@ from datetime import datetime from typing import Optional from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class ClickHouseEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @@ -46,8 +47,8 @@ class ClickHouseEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"toDate('{dttm.date().isoformat()}')" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""toDateTime('{dttm.isoformat(sep=" ", timespec="seconds")}')""" return None diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py index e43fe1def4..b1b0ceaf68 100644 --- a/superset/db_engine_specs/drill.py +++ b/superset/db_engine_specs/drill.py @@ -21,6 +21,7 @@ from urllib import parse from sqlalchemy.engine.url import URL from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class DrillEngineSpec(BaseEngineSpec): @@ -54,9 +55,9 @@ class DrillEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"TO_DATE('{dttm.date().isoformat()}', 'yyyy-MM-dd')" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""TO_TIMESTAMP('{dttm.isoformat(sep=" ", timespec="seconds")}', 'yyyy-MM-dd HH:mm:ss')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index ab4e36a522..5a81ea9dac 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -16,7 +16,8 @@ # under the License. import json import logging -from typing import Any, Dict, TYPE_CHECKING +from datetime import datetime +from typing import Any, Dict, Optional, TYPE_CHECKING from superset.db_engine_specs.base import BaseEngineSpec from superset.utils import core as utils @@ -41,6 +42,10 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method None: "{col}", "PT1S": "FLOOR({col} TO SECOND)", "PT1M": "FLOOR({col} TO MINUTE)", + "PT5M": "TIME_FLOOR({col}, 'PT5M')", + "PT10M": "TIME_FLOOR({col}, 'PT10M')", + "PT15M": "TIME_FLOOR({col}, 'PT15M')", + "PT0.5H": "TIME_FLOOR({col}, 'PT30M')", "PT1H": "FLOOR({col} TO HOUR)", "P1D": "FLOOR({col} TO DAY)", "P1W": "FLOOR({col} TO WEEK)", @@ -77,3 +82,12 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method engine_params["connect_args"] = connect_args extra["engine_params"] = engine_params return extra + + @classmethod + def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: + tt = target_type.upper() + if tt == utils.TemporalType.DATE: + return f"CAST(TIME_PARSE('{dttm.date().isoformat()}') AS DATE)" + if tt in (utils.TemporalType.DATETIME, utils.TemporalType.TIMESTAMP): + return f"""TIME_PARSE('{dttm.isoformat(timespec="seconds")}')""" + return None diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index 7ebc675a9b..f1cc561b0a 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/superset/db_engine_specs/elasticsearch.py @@ -18,6 +18,7 @@ from datetime import datetime from typing import Dict, Optional from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @@ -41,6 +42,6 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: - if target_type.upper() == "DATETIME": + if target_type.upper() == utils.TemporalType.DATETIME: return f"""CAST('{dttm.isoformat(timespec="seconds")}' AS DATETIME)""" return None diff --git a/superset/db_engine_specs/hana.py b/superset/db_engine_specs/hana.py index 0a80a02677..2514194f4f 100644 --- a/superset/db_engine_specs/hana.py +++ b/superset/db_engine_specs/hana.py @@ -19,6 +19,7 @@ from typing import Optional from superset.db_engine_specs.base import LimitMethod from superset.db_engine_specs.postgres import PostgresBaseEngineSpec +from superset.utils import core as utils class HanaEngineSpec(PostgresBaseEngineSpec): @@ -43,8 +44,8 @@ class HanaEngineSpec(PostgresBaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"TO_DATE('{dttm.date().isoformat()}', 'YYYY-MM-DD')" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""TO_TIMESTAMP('{dttm.isoformat(timespec="microseconds")}', 'YYYY-MM-DD"T"HH24:MI:SS.ff6')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 4cfcb4ce50..952f734f3e 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -198,9 +198,9 @@ class HiveEngineSpec(PrestoEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"CAST('{dttm.date().isoformat()}' AS DATE)" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""CAST('{dttm.isoformat(sep=" ", timespec="microseconds")}' AS TIMESTAMP)""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 59eaf0106e..c0c4f07430 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -20,6 +20,7 @@ from typing import List, Optional from sqlalchemy.engine.reflection import Inspector from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class ImpalaEngineSpec(BaseEngineSpec): @@ -45,9 +46,9 @@ class ImpalaEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"CAST('{dttm.date().isoformat()}' AS DATE)" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS TIMESTAMP)""" return None diff --git a/superset/db_engine_specs/kylin.py b/superset/db_engine_specs/kylin.py index 5e8bff5f2f..9f828d11d6 100644 --- a/superset/db_engine_specs/kylin.py +++ b/superset/db_engine_specs/kylin.py @@ -18,6 +18,7 @@ from datetime import datetime from typing import Optional from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class KylinEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @@ -42,8 +43,8 @@ class KylinEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"CAST('{dttm.date().isoformat()}' AS DATE)" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""CAST('{dttm.isoformat(sep=" ", timespec="seconds")}' AS TIMESTAMP)""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/mssql.py b/superset/db_engine_specs/mssql.py index 45c2f23a1b..a8a6b9a868 100644 --- a/superset/db_engine_specs/mssql.py +++ b/superset/db_engine_specs/mssql.py @@ -22,6 +22,7 @@ from typing import Any, List, Optional, Tuple, TYPE_CHECKING from sqlalchemy.types import String, TypeEngine, UnicodeText from superset.db_engine_specs.base import BaseEngineSpec, LimitMethod +from superset.utils import core as utils if TYPE_CHECKING: from superset.models.core import Database # pylint: disable=unused-import @@ -57,11 +58,11 @@ class MssqlEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"CONVERT(DATE, '{dttm.date().isoformat()}', 23)" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""CONVERT(DATETIME, '{dttm.isoformat(timespec="milliseconds")}', 126)""" # pylint: disable=line-too-long - if tt == "SMALLDATETIME": + if tt == utils.TemporalType.SMALLDATETIME: return f"""CONVERT(SMALLDATETIME, '{dttm.isoformat(sep=" ", timespec="seconds")}', 20)""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index abd7bd015e..7d750d3aae 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -21,6 +21,7 @@ from urllib import parse from sqlalchemy.engine.url import URL from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils class MySQLEngineSpec(BaseEngineSpec): @@ -51,9 +52,9 @@ class MySQLEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"STR_TO_DATE('{dttm.date().isoformat()}', '%Y-%m-%d')" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""STR_TO_DATE('{dttm.isoformat(sep=" ", timespec="microseconds")}', '%Y-%m-%d %H:%i:%s.%f')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/oracle.py b/superset/db_engine_specs/oracle.py index 0488aaedfc..813b1507a6 100644 --- a/superset/db_engine_specs/oracle.py +++ b/superset/db_engine_specs/oracle.py @@ -18,6 +18,7 @@ from datetime import datetime from typing import Optional from superset.db_engine_specs.base import BaseEngineSpec, LimitMethod +from superset.utils import core as utils class OracleEngineSpec(BaseEngineSpec): @@ -41,11 +42,11 @@ class OracleEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"TO_DATE('{dttm.date().isoformat()}', 'YYYY-MM-DD')" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""TO_DATE('{dttm.isoformat(timespec="seconds")}', 'YYYY-MM-DD"T"HH24:MI:SS')""" # pylint: disable=line-too-long - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""TO_TIMESTAMP('{dttm.isoformat(timespec="microseconds")}', 'YYYY-MM-DD"T"HH24:MI:SS.ff6')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index c5e4221c6f..dceac26b1c 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -21,6 +21,7 @@ from pytz import _FixedOffset # type: ignore from sqlalchemy.dialects.postgresql.base import PGInspector from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils if TYPE_CHECKING: # pylint: disable=unused-import @@ -79,8 +80,8 @@ class PostgresEngineSpec(PostgresBaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"TO_DATE('{dttm.date().isoformat()}', 'YYYY-MM-DD')" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""TO_TIMESTAMP('{dttm.isoformat(sep=" ", timespec="microseconds")}', 'YYYY-MM-DD HH24:MI:SS.US')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index ee474d53ae..6c014a5707 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -534,9 +534,9 @@ class PrestoEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"""from_iso8601_date('{dttm.date().isoformat()}')""" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""from_iso8601_timestamp('{dttm.isoformat(timespec="microseconds")}')""" # pylint: disable=line-too-long return None diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 0907443016..8b78b52487 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -22,6 +22,7 @@ from urllib import parse from sqlalchemy.engine.url import URL from superset.db_engine_specs.postgres import PostgresBaseEngineSpec +from superset.utils import core as utils if TYPE_CHECKING: from superset.models.core import Database # pylint: disable=unused-import @@ -74,11 +75,11 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: tt = target_type.upper() - if tt == "DATE": + if tt == utils.TemporalType.DATE: return f"TO_DATE('{dttm.date().isoformat()}')" - if tt == "DATETIME": + if tt == utils.TemporalType.DATETIME: return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS DATETIME)""" - if tt == "TIMESTAMP": + if tt == utils.TemporalType.TIMESTAMP: return f"""TO_TIMESTAMP('{dttm.isoformat(timespec="microseconds")}')""" return None diff --git a/superset/db_engine_specs/sqlite.py b/superset/db_engine_specs/sqlite.py index 863e5e624c..72f1f31871 100644 --- a/superset/db_engine_specs/sqlite.py +++ b/superset/db_engine_specs/sqlite.py @@ -75,7 +75,8 @@ class SqliteEngineSpec(BaseEngineSpec): @classmethod def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: - if target_type.upper() == "TEXT": + tt = target_type.upper() + if tt == utils.TemporalType.TEXT: return f"""'{dttm.isoformat(sep=" ", timespec="microseconds")}'""" return None diff --git a/superset/utils/core.py b/superset/utils/core.py index 1620542af4..c7fcd46532 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -1405,3 +1405,16 @@ class ChartDataResultFormat(str, Enum): CSV = "csv" JSON = "json" + + +class TemporalType(str, Enum): + """ + Supported temporal types + """ + + DATE = "DATE" + DATETIME = "DATETIME" + SMALLDATETIME = "SMALLDATETIME" + TEXT = "TEXT" + TIME = "TIME" + TIMESTAMP = "TIMESTAMP" diff --git a/tests/db_engine_specs/druid_tests.py b/tests/db_engine_specs/druid_tests.py new file mode 100644 index 0000000000..76f49e9567 --- /dev/null +++ b/tests/db_engine_specs/druid_tests.py @@ -0,0 +1,56 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from sqlalchemy import column + +from superset.db_engine_specs.druid import DruidEngineSpec +from tests.db_engine_specs.base_tests import DbEngineSpecTestCase + + +class DruidTestCase(DbEngineSpecTestCase): + def test_convert_dttm(self): + dttm = self.get_dttm() + + self.assertEqual( + DruidEngineSpec.convert_dttm("DATETIME", dttm), + "TIME_PARSE('2019-01-02T03:04:05')", + ) + + self.assertEqual( + DruidEngineSpec.convert_dttm("TIMESTAMP", dttm), + "TIME_PARSE('2019-01-02T03:04:05')", + ) + + self.assertEqual( + DruidEngineSpec.convert_dttm("DATE", dttm), + "CAST(TIME_PARSE('2019-01-02') AS DATE)", + ) + + def test_timegrain_expressions(self): + """ + DB Eng Specs (druid): Test time grain expressions + """ + col = "__time" + sqla_col = column(col) + test_cases = { + "PT1S": f"FLOOR({col} TO SECOND)", + "PT5M": f"TIME_FLOOR({col}, 'PT5M')", + } + for grain, expected in test_cases.items(): + actual = DruidEngineSpec.get_timestamp_expr( + col=sqla_col, pdf=None, time_grain=grain + ) + self.assertEqual(str(actual), expected)