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)
apache · Jun 25, 2020 · b205ce3 · b205ce3
1 parent ecb44a4
commit b205ce3
Show file tree

Hide file tree

Showing 19 changed files with 133 additions and 36 deletions.
diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py
@@ -18,6 +18,7 @@
 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
@@ -24,6 +24,7 @@
 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
@@ -18,6 +18,7 @@
 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
@@ -21,6 +21,7 @@
 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 @@ def epoch_ms_to_dttm(cls) -> str:
     @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
@@ -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 @@ def get_extra_params(database: "Database") -> Dict[str, Any]:
             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
@@ -18,6 +18,7 @@
 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
@@ -19,6 +19,7 @@
 
 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
@@ -198,9 +198,9 @@ def convert_to_hive_type(col_type: str) -> str:
     @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
@@ -20,6 +20,7 @@
 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 @@ def epoch_to_dttm(cls) -> str:
     @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
@@ -18,6 +18,7 @@
 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
@@ -22,6 +22,7 @@
 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 @@ def epoch_to_dttm(cls) -> str:
     @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
@@ -21,6 +21,7 @@
 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
@@ -18,6 +18,7 @@
 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
@@ -21,6 +21,7 @@
 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 @@ def get_table_names(
     @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
@@ -534,9 +534,9 @@ def adjust_database_uri(
     @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
@@ -22,6 +22,7 @@
 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 @@ def epoch_ms_to_dttm(cls) -> str:
     @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
@@ -75,7 +75,8 @@ def get_all_datasource_names(
 
     @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
@@ -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"