From d041eb6fcd5cebffed042c138008a5d0c4b32f1f Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Fri, 12 Mar 2021 10:43:50 +0100 Subject: [PATCH 1/8] add __time column type mapping --- superset/db_engine_specs/druid.py | 47 ++++++++++++++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 3ab2dedb4c0d1..57d00abdbac16 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -16,12 +16,26 @@ # under the License. import json import logging +import re from datetime import datetime -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import ( + Any, + Callable, + Dict, + Match, + Optional, + Pattern, + Tuple, + TYPE_CHECKING, + Union, +) + +from sqlalchemy.types import TIMESTAMP, TypeEngine from superset.db_engine_specs.base import BaseEngineSpec from superset.exceptions import SupersetException from superset.utils import core as utils +from superset.utils.core import ColumnSpec, GenericDataType if TYPE_CHECKING: from superset.connectors.sqla.models import TableColumn @@ -38,6 +52,14 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method allows_joins = False allows_subqueries = True + column_type_mappings = ( + ( + re.compile(r"^__time", re.IGNORECASE), + TIMESTAMP(), + utils.GenericDataType.TEMPORAL, + ), + ) + _time_grain_expressions = { None: "{col}", "PT1S": "FLOOR({col} TO SECOND)", @@ -91,3 +113,26 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: if tt in (utils.TemporalType.DATETIME, utils.TemporalType.TIMESTAMP): return f"""TIME_PARSE('{dttm.isoformat(timespec="seconds")}')""" return None + + @classmethod + def get_column_spec( # type: ignore + cls, + native_type: Optional[str], + source: utils.ColumnTypeSource = utils.ColumnTypeSource.GET_TABLE, + column_type_mappings: Tuple[ + Tuple[ + Pattern[str], + Union[TypeEngine, Callable[[Match[str]], TypeEngine]], + GenericDataType, + ], + ..., + ] = column_type_mappings, + ) -> Union[ColumnSpec, None]: + + column_spec = super().get_column_spec(native_type) + if column_spec: + return column_spec + + return super().get_column_spec( + native_type, column_type_mappings=column_type_mappings + ) From 536be6149c3bbfea66f50f1a0824fc7f2a3fa405 Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Fri, 12 Mar 2021 12:14:47 +0100 Subject: [PATCH 2/8] add __time column type mapping --- superset/db_engine_specs/druid.py | 26 ++++---------------------- 1 file changed, 4 insertions(+), 22 deletions(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 57d00abdbac16..decaafc191eb9 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -52,14 +52,6 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method allows_joins = False allows_subqueries = True - column_type_mappings = ( - ( - re.compile(r"^__time", re.IGNORECASE), - TIMESTAMP(), - utils.GenericDataType.TEMPORAL, - ), - ) - _time_grain_expressions = { None: "{col}", "PT1S": "FLOOR({col} TO SECOND)", @@ -118,21 +110,11 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: def get_column_spec( # type: ignore cls, native_type: Optional[str], + column_name: Optional[str], source: utils.ColumnTypeSource = utils.ColumnTypeSource.GET_TABLE, - column_type_mappings: Tuple[ - Tuple[ - Pattern[str], - Union[TypeEngine, Callable[[Match[str]], TypeEngine]], - GenericDataType, - ], - ..., - ] = column_type_mappings, ) -> Union[ColumnSpec, None]: - column_spec = super().get_column_spec(native_type) - if column_spec: - return column_spec + if column_name == "__time": + return ColumnSpec(TIMESTAMP, GenericDataType.TEMPORAL, True) - return super().get_column_spec( - native_type, column_type_mappings=column_type_mappings - ) + return super().get_column_spec(native_type) From f4cb1755f288f2bc2f82593d0f3dfd82f266d823 Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Fri, 12 Mar 2021 14:33:34 +0100 Subject: [PATCH 3/8] replacing column_name __time type --- superset/db_engine_specs/druid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index decaafc191eb9..0c7adeeeedbd3 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -110,11 +110,11 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: def get_column_spec( # type: ignore cls, native_type: Optional[str], - column_name: Optional[str], + column_name: Optional[str] = "", source: utils.ColumnTypeSource = utils.ColumnTypeSource.GET_TABLE, ) -> Union[ColumnSpec, None]: if column_name == "__time": - return ColumnSpec(TIMESTAMP, GenericDataType.TEMPORAL, True) + return ColumnSpec(TIMESTAMP(), GenericDataType.TEMPORAL, True) return super().get_column_spec(native_type) From 8bc18152ca0b0bfe3393a1a0d99bcf63fb209afa Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Fri, 12 Mar 2021 15:40:44 +0100 Subject: [PATCH 4/8] replacing column_name __time type --- superset/db_engine_specs/druid.py | 14 ++++++++------ tests/sqla_models_tests.py | 8 ++++---- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 0c7adeeeedbd3..9308df5d38c52 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -16,7 +16,6 @@ # under the License. import json import logging -import re from datetime import datetime from typing import ( Any, @@ -68,11 +67,6 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method "P1Y": "FLOOR({col} TO YEAR)", } - @classmethod - def alter_new_orm_column(cls, orm_col: "TableColumn") -> None: - if orm_col.column_name == "__time": - orm_col.is_dttm = True - @staticmethod def get_extra_params(database: "Database") -> Dict[str, Any]: """ @@ -110,6 +104,14 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: def get_column_spec( # type: ignore cls, native_type: Optional[str], + column_type_mappings: Tuple[ + Tuple[ + Pattern[str], + Union[TypeEngine, Callable[[Match[str]], TypeEngine]], + GenericDataType, + ], + ..., + ], column_name: Optional[str] = "", source: utils.ColumnTypeSource = utils.ColumnTypeSource.GET_TABLE, ) -> Union[ColumnSpec, None]: diff --git a/tests/sqla_models_tests.py b/tests/sqla_models_tests.py index cdd77c270b90d..c629769908936 100644 --- a/tests/sqla_models_tests.py +++ b/tests/sqla_models_tests.py @@ -54,10 +54,10 @@ def test_is_time_druid_time_col(self): database = Database(database_name="druid_db", sqlalchemy_uri="druid://db") tbl = SqlaTable(table_name="druid_tbl", database=database) - col = TableColumn(column_name="__time", type="INTEGER", table=tbl) - self.assertEqual(col.is_dttm, None) - DruidEngineSpec.alter_new_orm_column(col) - self.assertEqual(col.is_dttm, True) + column_spec = DruidEngineSpec.get_column_spec("INTEGER", "__time") + self.assertEqual(column_spec.is_dttm, True) + col = TableColumn(column_name="__time", type=column_spec.sqla_type, table=tbl) + self.assertTrue(col) col = TableColumn(column_name="__not_time", type="INTEGER", table=tbl) self.assertEqual(col.is_temporal, False) From 70b74564a72b625285d7edd2a82adba4a5f8b0eb Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Fri, 12 Mar 2021 15:45:51 +0100 Subject: [PATCH 5/8] replacing column_name __time type --- superset/db_engine_specs/druid.py | 1 + 1 file changed, 1 insertion(+) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 9308df5d38c52..fb81a45754094 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -102,6 +102,7 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: @classmethod def get_column_spec( # type: ignore + # pylint: disable=arguments-differ cls, native_type: Optional[str], column_type_mappings: Tuple[ From 309f5f749c05f723d713e557396bafd7f7ad0d08 Mon Sep 17 00:00:00 2001 From: Nikola Gigic Date: Mon, 15 Mar 2021 13:46:46 +0100 Subject: [PATCH 6/8] fix tests --- superset/db_engine_specs/druid.py | 7 ++++--- tests/sqla_models_tests.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index fb81a45754094..419882b9eb41a 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -104,7 +104,7 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: def get_column_spec( # type: ignore # pylint: disable=arguments-differ cls, - native_type: Optional[str], + column_name: Optional[str] = "", column_type_mappings: Tuple[ Tuple[ Pattern[str], @@ -112,12 +112,13 @@ def get_column_spec( # type: ignore GenericDataType, ], ..., - ], - column_name: Optional[str] = "", + ] = (), + native_type: Optional[str] = "", source: utils.ColumnTypeSource = utils.ColumnTypeSource.GET_TABLE, ) -> Union[ColumnSpec, None]: if column_name == "__time": + print(">>> hello!") return ColumnSpec(TIMESTAMP(), GenericDataType.TEMPORAL, True) return super().get_column_spec(native_type) diff --git a/tests/sqla_models_tests.py b/tests/sqla_models_tests.py index c629769908936..608feec8ca911 100644 --- a/tests/sqla_models_tests.py +++ b/tests/sqla_models_tests.py @@ -54,8 +54,9 @@ def test_is_time_druid_time_col(self): database = Database(database_name="druid_db", sqlalchemy_uri="druid://db") tbl = SqlaTable(table_name="druid_tbl", database=database) - column_spec = DruidEngineSpec.get_column_spec("INTEGER", "__time") - self.assertEqual(column_spec.is_dttm, True) + column_spec = DruidEngineSpec.get_column_spec(column_name="__time") + print(">>> {}".format(column_spec)) + self.assertTrue(column_spec.is_dttm) col = TableColumn(column_name="__time", type=column_spec.sqla_type, table=tbl) self.assertTrue(col) From ee8b1d1e7dbcb97ae988ce39a129b3901779200d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nikola=20Gigi=C4=87?= Date: Tue, 23 Mar 2021 16:49:31 +0100 Subject: [PATCH 7/8] Update superset/db_engine_specs/druid.py Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com> --- superset/db_engine_specs/druid.py | 1 - 1 file changed, 1 deletion(-) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 419882b9eb41a..b2db66bd6a717 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -118,7 +118,6 @@ def get_column_spec( # type: ignore ) -> Union[ColumnSpec, None]: if column_name == "__time": - print(">>> hello!") return ColumnSpec(TIMESTAMP(), GenericDataType.TEMPORAL, True) return super().get_column_spec(native_type) From 2bf41c9a87335a2a4a22b8b1b2cd621e38c47ec8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nikola=20Gigi=C4=87?= Date: Tue, 23 Mar 2021 16:49:38 +0100 Subject: [PATCH 8/8] Update tests/sqla_models_tests.py Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com> --- tests/sqla_models_tests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/sqla_models_tests.py b/tests/sqla_models_tests.py index 608feec8ca911..acddcde7b03c6 100644 --- a/tests/sqla_models_tests.py +++ b/tests/sqla_models_tests.py @@ -55,7 +55,6 @@ def test_is_time_druid_time_col(self): database = Database(database_name="druid_db", sqlalchemy_uri="druid://db") tbl = SqlaTable(table_name="druid_tbl", database=database) column_spec = DruidEngineSpec.get_column_spec(column_name="__time") - print(">>> {}".format(column_spec)) self.assertTrue(column_spec.is_dttm) col = TableColumn(column_name="__time", type=column_spec.sqla_type, table=tbl) self.assertTrue(col)