diff --git a/superset/connectors/sqla/utils.py b/superset/connectors/sqla/utils.py index 14b9071d1da1..a2b54201d6b8 100644 --- a/superset/connectors/sqla/utils.py +++ b/superset/connectors/sqla/utils.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. from contextlib import closing +from datetime import date, datetime, time, timedelta from typing import Callable, Dict, List, Optional, Set, TYPE_CHECKING import sqlparse @@ -41,7 +42,7 @@ from superset.connectors.sqla.models import SqlaTable -TEMPORAL_TYPES = {"DATETIME", "DATE", "TIME", "TIMEDELTA"} +TEMPORAL_TYPES = {date, datetime, time, timedelta} def get_physical_table_metadata( @@ -172,6 +173,13 @@ def validate_adhoc_subquery( return ";\n".join(str(statement) for statement in statements) +def is_column_type_temporal(column_type: TypeEngine) -> bool: + try: + return column_type.python_type in TEMPORAL_TYPES + except NotImplementedError: + return False + + def load_or_create_tables( # pylint: disable=too-many-arguments session: Session, database_id: int, @@ -223,8 +231,7 @@ def load_or_create_tables( # pylint: disable=too-many-arguments name=column["name"], type=str(column["type"]), expression=conditional_quote(column["name"]), - is_temporal=column["type"].python_type.__name__.upper() - in TEMPORAL_TYPES, + is_temporal=is_column_type_temporal(column["type"]), is_aggregation=False, is_physical=True, is_spatial=False, diff --git a/superset/migrations/versions/b8d3a24d9131_new_dataset_models.py b/superset/migrations/versions/b8d3a24d9131_new_dataset_models.py index 533f8a9fdcbe..8728e9adb7b8 100644 --- a/superset/migrations/versions/b8d3a24d9131_new_dataset_models.py +++ b/superset/migrations/versions/b8d3a24d9131_new_dataset_models.py @@ -25,6 +25,7 @@ """ import json +from datetime import date, datetime, time, timedelta from typing import Callable, List, Optional, Set from uuid import uuid4 @@ -34,6 +35,7 @@ from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import backref, relationship, Session from sqlalchemy.schema import UniqueConstraint +from sqlalchemy.sql.type_api import TypeEngine from sqlalchemy_utils import UUIDType from superset import app, db @@ -230,7 +232,14 @@ class NewDataset(Base): external_url = sa.Column(sa.Text, nullable=True) -TEMPORAL_TYPES = {"DATETIME", "DATE", "TIME", "TIMEDELTA"} +TEMPORAL_TYPES = {date, datetime, time, timedelta} + + +def is_column_type_temporal(column_type: TypeEngine) -> bool: + try: + return column_type.python_type in TEMPORAL_TYPES + except NotImplementedError: + return False def load_or_create_tables( @@ -285,8 +294,7 @@ def load_or_create_tables( name=column["name"], type=str(column["type"]), expression=conditional_quote(column["name"]), - is_temporal=column["type"].python_type.__name__.upper() - in TEMPORAL_TYPES, + is_temporal=is_column_type_temporal(column["type"]), is_aggregation=False, is_physical=True, is_spatial=False,