refactor TEMPORAL_TYPES

apache · Apr 7, 2022 · b95fb1a · b95fb1a
1 parent 90ba057
commit b95fb1a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git 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(
@@ -174,7 +175,7 @@ def validate_adhoc_subquery(
 
 def is_column_type_temporal(column_type: TypeEngine) -> bool:
     try:
-        return column_type.python_type.__name__.upper() in TEMPORAL_TYPES
+        return column_type.python_type in TEMPORAL_TYPES
     except NotImplementedError:
         return False
 

diff --git 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
 
@@ -231,12 +232,12 @@ 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.__name__.upper() in TEMPORAL_TYPES
+        return column_type.python_type in TEMPORAL_TYPES
     except NotImplementedError:
         return False