From 6d4d9ede873c0bb501154b2dac01b3a4b72a1635 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Tue, 3 May 2022 07:19:36 +0300
Subject: [PATCH] fix(sqla): replace custom dttm type with literal_column
 (#19917)

---
 superset/connectors/base/models.py            | 31 ++++++---
 superset/connectors/druid/models.py           |  2 +-
 superset/connectors/sqla/models.py            | 20 ++++--
 superset/db_engine_specs/athena.py            |  4 +-
 superset/db_engine_specs/base.py              |  9 +--
 superset/db_engine_specs/presto.py            |  6 +-
 superset/db_engine_specs/trino.py             |  4 +-
 superset/models/sql_types/base.py             | 65 -------------------
 .../db_engine_specs/presto_tests.py           |  6 +-
 .../db_engine_specs/trino_tests.py            |  4 +-
 tests/integration_tests/model_tests.py        | 24 +------
 .../unit_tests/db_engine_specs/test_athena.py |  9 +--
 .../unit_tests/db_engine_specs/test_presto.py |  8 +--
 .../unit_tests/db_engine_specs/test_trino.py  |  8 +--
 14 files changed, 63 insertions(+), 137 deletions(-)
 delete mode 100644 superset/models/sql_types/base.py

diff --git a/superset/connectors/base/models.py b/superset/connectors/base/models.py
index 3d22857912f10..e72747cc50ff9 100644
--- a/superset/connectors/base/models.py
+++ b/superset/connectors/base/models.py
@@ -14,15 +14,18 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
+
 import json
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, Hashable, List, Optional, Set, Type, Union
+from typing import Any, Dict, Hashable, List, Optional, Set, Type, TYPE_CHECKING, Union
 
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy import and_, Boolean, Column, Integer, String, Text
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.orm import foreign, Query, relationship, RelationshipProperty, Session
+from sqlalchemy.sql import literal_column
 
 from superset import is_feature_enabled, security_manager
 from superset.constants import EMPTY_STRING, NULL_STRING
@@ -33,6 +36,9 @@
 from superset.utils import core as utils
 from superset.utils.core import GenericDataType, MediumText
 
+if TYPE_CHECKING:
+    from superset.db_engine_specs.base import BaseEngineSpec
+
 METRIC_FORM_DATA_PARAMS = [
     "metric",
     "metric_2",
@@ -387,25 +393,34 @@ def data_for_slices(  # pylint: disable=too-many-locals
         return data
 
     @staticmethod
-    def filter_values_handler(
+    def filter_values_handler(  # pylint: disable=too-many-arguments
         values: Optional[FilterValues],
-        target_column_type: utils.GenericDataType,
+        target_generic_type: GenericDataType,
+        target_native_type: Optional[str] = None,
         is_list_target: bool = False,
+        db_engine_spec: Optional[Type[BaseEngineSpec]] = None,
+        db_extra: Optional[Dict[str, Any]] = None,
     ) -> Optional[FilterValues]:
         if values is None:
             return None
 
         def handle_single_value(value: Optional[FilterValue]) -> Optional[FilterValue]:
-            # backward compatibility with previous <select> components
             if (
                 isinstance(value, (float, int))
-                and target_column_type == utils.GenericDataType.TEMPORAL
+                and target_generic_type == utils.GenericDataType.TEMPORAL
+                and target_native_type is not None
+                and db_engine_spec is not None
             ):
-                return datetime.utcfromtimestamp(value / 1000)
+                value = db_engine_spec.convert_dttm(
+                    target_type=target_native_type,
+                    dttm=datetime.utcfromtimestamp(value / 1000),
+                    db_extra=db_extra,
+                )
+                value = literal_column(value)
             if isinstance(value, str):
                 value = value.strip("\t\n")
 
-                if target_column_type == utils.GenericDataType.NUMERIC:
+                if target_generic_type == utils.GenericDataType.NUMERIC:
                     # For backwards compatibility and edge cases
                     # where a column data type might have changed
                     return utils.cast_to_num(value)
@@ -413,7 +428,7 @@ def handle_single_value(value: Optional[FilterValue]) -> Optional[FilterValue]:
                     return None
                 if value == EMPTY_STRING:
                     return ""
-            if target_column_type == utils.GenericDataType.BOOLEAN:
+            if target_generic_type == utils.GenericDataType.BOOLEAN:
                 return utils.cast_to_boolean(value)
             return value
 
diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py
index bc4d0ba3817a4..4025c3f05377b 100644
--- a/superset/connectors/druid/models.py
+++ b/superset/connectors/druid/models.py
@@ -1533,7 +1533,7 @@ def get_filters(
             eq = cls.filter_values_handler(
                 eq,
                 is_list_target=is_list_target,
-                target_column_type=utils.GenericDataType.NUMERIC
+                target_generic_type=utils.GenericDataType.NUMERIC
                 if is_numeric_col
                 else utils.GenericDataType.STRING,
             )
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index e0382c659514c..94190f678b481 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -1459,21 +1459,26 @@ def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-ma
                     )
                 elif col_obj:
                     sqla_col = col_obj.get_sqla_col()
+                col_type = col_obj.type if col_obj else None
                 col_spec = db_engine_spec.get_column_spec(
-                    col_obj.type if col_obj else None
+                    native_type=col_type,
+                    db_extra=self.database.get_extra(),
                 )
                 is_list_target = op in (
                     utils.FilterOperator.IN.value,
                     utils.FilterOperator.NOT_IN.value,
                 )
                 if col_spec:
-                    target_type = col_spec.generic_type
+                    target_generic_type = col_spec.generic_type
                 else:
-                    target_type = GenericDataType.STRING
+                    target_generic_type = GenericDataType.STRING
                 eq = self.filter_values_handler(
                     values=val,
-                    target_column_type=target_type,
+                    target_generic_type=target_generic_type,
+                    target_native_type=col_type,
                     is_list_target=is_list_target,
+                    db_engine_spec=db_engine_spec,
+                    db_extra=self.database.get_extra(),
                 )
                 if is_list_target:
                     assert isinstance(eq, (tuple, list))
@@ -1481,11 +1486,12 @@ def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-ma
                         raise QueryObjectValidationError(
                             _("Filter value list cannot be empty")
                         )
-                    if None in eq:
-                        eq = [x for x in eq if x is not None]
+                    if len(eq) > len(
+                        eq_without_none := [x for x in eq if x is not None]
+                    ):
                         is_null_cond = sqla_col.is_(None)
                         if eq:
-                            cond = or_(is_null_cond, sqla_col.in_(eq))
+                            cond = or_(is_null_cond, sqla_col.in_(eq_without_none))
                         else:
                             cond = is_null_cond
                     else:
diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py
index 19bca8780802e..2394fe088da65 100644
--- a/superset/db_engine_specs/athena.py
+++ b/superset/db_engine_specs/athena.py
@@ -67,10 +67,10 @@ def convert_dttm(
     ) -> Optional[str]:
         tt = target_type.upper()
         if tt == utils.TemporalType.DATE:
-            return f"from_iso8601_date('{dttm.date().isoformat()}')"
+            return f"DATE '{dttm.date().isoformat()}'"
         if tt == utils.TemporalType.TIMESTAMP:
             datetime_formatted = dttm.isoformat(timespec="microseconds")
-            return f"""from_iso8601_timestamp('{datetime_formatted}')"""
+            return f"""TIMESTAMP '{datetime_formatted}'"""
         return None
 
     @classmethod
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 7d133e6245c1d..564367372ef9d 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -61,7 +61,6 @@
 from superset.databases.utils import make_url_safe
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.models.sql_lab import Query
-from superset.models.sql_types.base import literal_dttm_type_factory
 from superset.sql_parse import ParsedQuery, Table
 from superset.superset_typing import ResultSetColumnType
 from superset.utils import core as utils
@@ -269,7 +268,7 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
         ),
         (
             re.compile(r"^date", re.IGNORECASE),
-            types.DateTime(),
+            types.Date(),
             GenericDataType.TEMPORAL,
         ),
         (
@@ -1488,12 +1487,6 @@ def get_column_spec(  # pylint: disable=unused-argument
         )
         if col_types:
             column_type, generic_type = col_types
-            # wrap temporal types in custom type that supports literal binding
-            # using datetimes
-            if generic_type == GenericDataType.TEMPORAL:
-                column_type = literal_dttm_type_factory(
-                    column_type, cls, native_type or "", db_extra=db_extra or {}
-                )
             is_dttm = generic_type == GenericDataType.TEMPORAL
             return ColumnSpec(
                 sqla_type=column_type, generic_type=generic_type, is_dttm=is_dttm
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index bd7424b0de077..dc540ffff2a68 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -508,7 +508,7 @@ def _show_columns(
         ),
         (
             re.compile(r"^date.*", re.IGNORECASE),
-            types.DATETIME(),
+            types.DATE(),
             GenericDataType.TEMPORAL,
         ),
         (
@@ -756,12 +756,12 @@ def convert_dttm(
         """
         tt = target_type.upper()
         if tt == utils.TemporalType.DATE:
-            return f"""from_iso8601_date('{dttm.date().isoformat()}')"""
+            return f"""DATE '{dttm.date().isoformat()}'"""
         if tt in (
             utils.TemporalType.TIMESTAMP,
             utils.TemporalType.TIMESTAMP_WITH_TIME_ZONE,
         ):
-            return f"""from_iso8601_timestamp('{dttm.isoformat(timespec="microseconds")}')"""  # pylint: disable=line-too-long,useless-suppression
+            return f"""TIMESTAMP '{dttm.isoformat(timespec="microseconds")}'"""
         return None
 
     @classmethod
diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
index 8ff4cfde59676..f9d06bb9c6a70 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -71,12 +71,12 @@ def convert_dttm(
         """
         tt = target_type.upper()
         if tt == utils.TemporalType.DATE:
-            return f"from_iso8601_date('{dttm.date().isoformat()}')"
+            return f"DATE '{dttm.date().isoformat()}'"
         if tt in (
             utils.TemporalType.TIMESTAMP,
             utils.TemporalType.TIMESTAMP_WITH_TIME_ZONE,
         ):
-            return f"""from_iso8601_timestamp('{dttm.isoformat(timespec="microseconds")}')"""  # pylint: disable=line-too-long,useless-suppression
+            return f"""TIMESTAMP '{dttm.isoformat(timespec="microseconds")}'"""
         return None
 
     @classmethod
diff --git a/superset/models/sql_types/base.py b/superset/models/sql_types/base.py
deleted file mode 100644
index 597631a372163..0000000000000
--- a/superset/models/sql_types/base.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# 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 datetime import datetime
-from typing import Any, Callable, Dict, Type, TYPE_CHECKING
-
-from flask_babel import gettext as __
-from sqlalchemy import types
-from sqlalchemy.engine.interfaces import Dialect
-
-if TYPE_CHECKING:
-    from superset.db_engine_specs.base import BaseEngineSpec
-
-
-def literal_dttm_type_factory(
-    sqla_type: types.TypeEngine,
-    db_engine_spec: Type["BaseEngineSpec"],
-    col_type: str,
-    db_extra: Dict[str, Any],
-) -> types.TypeEngine:
-    """
-    Create a custom SQLAlchemy type that supports datetime literal binds.
-
-    :param sqla_type: Base type to extend
-    :param db_engine_spec: Database engine spec which supports `convert_dttm` method
-    :param col_type: native column type as defined in table metadata
-    :param db_extra: The database extra object
-    :return: SQLAlchemy type that supports using datetima as literal bind
-    """
-    # pylint: disable=too-few-public-methods
-    class TemporalWrapperType(type(sqla_type)):  # type: ignore
-        # pylint: disable=unused-argument
-        def literal_processor(self, dialect: Dialect) -> Callable[[Any], Any]:
-            def process(value: Any) -> Any:
-                if isinstance(value, datetime):
-                    ts_expression = db_engine_spec.convert_dttm(
-                        col_type, value, db_extra=db_extra
-                    )
-                    if ts_expression is None:
-                        raise NotImplementedError(
-                            __(
-                                "Temporal expression not supported for type: "
-                                "%(col_type)s",
-                                col_type=col_type,
-                            )
-                        )
-                    return ts_expression
-                return super().process(value)
-
-            return process
-
-    return TemporalWrapperType()
diff --git a/tests/integration_tests/db_engine_specs/presto_tests.py b/tests/integration_tests/db_engine_specs/presto_tests.py
index 17c7c2d900e7f..acb9922d45632 100644
--- a/tests/integration_tests/db_engine_specs/presto_tests.py
+++ b/tests/integration_tests/db_engine_specs/presto_tests.py
@@ -517,12 +517,12 @@ def test_convert_dttm(self):
 
         self.assertEqual(
             PrestoEngineSpec.convert_dttm("DATE", dttm),
-            "from_iso8601_date('2019-01-02')",
+            "DATE '2019-01-02'",
         )
 
         self.assertEqual(
             PrestoEngineSpec.convert_dttm("TIMESTAMP", dttm),
-            "from_iso8601_timestamp('2019-01-02T03:04:05.678900')",
+            "TIMESTAMP '2019-01-02T03:04:05.678900'",
         )
 
     def test_query_cost_formatter(self):
@@ -672,12 +672,10 @@ def test_get_sqla_column_type(self):
 
         column_spec = PrestoEngineSpec.get_column_spec("time")
         assert isinstance(column_spec.sqla_type, types.Time)
-        assert type(column_spec.sqla_type).__name__ == "TemporalWrapperType"
         self.assertEqual(column_spec.generic_type, GenericDataType.TEMPORAL)
 
         column_spec = PrestoEngineSpec.get_column_spec("timestamp")
         assert isinstance(column_spec.sqla_type, types.TIMESTAMP)
-        assert type(column_spec.sqla_type).__name__ == "TemporalWrapperType"
         self.assertEqual(column_spec.generic_type, GenericDataType.TEMPORAL)
 
         sqla_type = PrestoEngineSpec.get_sqla_column_type(None)
diff --git a/tests/integration_tests/db_engine_specs/trino_tests.py b/tests/integration_tests/db_engine_specs/trino_tests.py
index 973eb17159ae7..18f00b8988c4d 100644
--- a/tests/integration_tests/db_engine_specs/trino_tests.py
+++ b/tests/integration_tests/db_engine_specs/trino_tests.py
@@ -32,12 +32,12 @@ def test_convert_dttm(self):
 
         self.assertEqual(
             TrinoEngineSpec.convert_dttm("DATE", dttm),
-            "from_iso8601_date('2019-01-02')",
+            "DATE '2019-01-02'",
         )
 
         self.assertEqual(
             TrinoEngineSpec.convert_dttm("TIMESTAMP", dttm),
-            "from_iso8601_timestamp('2019-01-02T03:04:05.678900')",
+            "TIMESTAMP '2019-01-02T03:04:05.678900'",
         )
 
     def test_adjust_database_uri(self):
diff --git a/tests/integration_tests/model_tests.py b/tests/integration_tests/model_tests.py
index 8f90d46a9aa9b..a69ffc931d41d 100644
--- a/tests/integration_tests/model_tests.py
+++ b/tests/integration_tests/model_tests.py
@@ -37,7 +37,6 @@
 from superset.common.db_query_status import QueryStatus
 from superset.models.core import Database
 from superset.models.slice import Slice
-from superset.models.sql_types.base import literal_dttm_type_factory
 from superset.utils.database import get_example_database
 
 from .base_tests import SupersetTestCase
@@ -376,30 +375,22 @@ def test_get_sqla_engine(self, mocked_create_engine):
 class TestSqlaTableModel(SupersetTestCase):
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_get_timestamp_expression(self):
-        col_type = (
-            "VARCHAR"
-            if get_example_database().backend == "presto"
-            else "TemporalWrapperType"
-        )
         tbl = self.get_table(name="birth_names")
         ds_col = tbl.get_column("ds")
         sqla_literal = ds_col.get_timestamp_expression(None)
-        self.assertEqual(str(sqla_literal.compile()), "ds")
-        assert type(sqla_literal.type).__name__ == col_type
+        assert str(sqla_literal.compile()) == "ds"
 
         sqla_literal = ds_col.get_timestamp_expression("P1D")
-        assert type(sqla_literal.type).__name__ == col_type
         compiled = "{}".format(sqla_literal.compile())
         if tbl.database.backend == "mysql":
-            self.assertEqual(compiled, "DATE(ds)")
+            assert compiled == "DATE(ds)"
 
         prev_ds_expr = ds_col.expression
         ds_col.expression = "DATE_ADD(ds, 1)"
         sqla_literal = ds_col.get_timestamp_expression("P1D")
-        assert type(sqla_literal.type).__name__ == col_type
         compiled = "{}".format(sqla_literal.compile())
         if tbl.database.backend == "mysql":
-            self.assertEqual(compiled, "DATE(DATE_ADD(ds, 1))")
+            assert compiled == "DATE(DATE_ADD(ds, 1))"
         ds_col.expression = prev_ds_expr
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
@@ -642,12 +633,3 @@ def test_data_for_slices_with_adhoc_column(self):
 
         # clean up and auto commit
         metadata_db.session.delete(slc)
-
-
-def test_literal_dttm_type_factory():
-    orig_type = DateTime()
-    new_type = literal_dttm_type_factory(
-        orig_type, PostgresEngineSpec, "TIMESTAMP", db_extra={}
-    )
-    assert type(new_type).__name__ == "TemporalWrapperType"
-    assert str(new_type) == str(orig_type)
diff --git a/tests/unit_tests/db_engine_specs/test_athena.py b/tests/unit_tests/db_engine_specs/test_athena.py
index 32a401f595570..a35d8eddfb0ad 100644
--- a/tests/unit_tests/db_engine_specs/test_athena.py
+++ b/tests/unit_tests/db_engine_specs/test_athena.py
@@ -35,13 +35,11 @@ def test_convert_dttm(app_context: AppContext, dttm: datetime) -> None:
 
     from superset.db_engine_specs.athena import AthenaEngineSpec
 
-    assert (
-        AthenaEngineSpec.convert_dttm("DATE", dttm) == "from_iso8601_date('2019-01-02')"
-    )
+    assert AthenaEngineSpec.convert_dttm("DATE", dttm) == "DATE '2019-01-02'"
 
     assert (
         AthenaEngineSpec.convert_dttm("TIMESTAMP", dttm)
-        == "from_iso8601_timestamp('2019-01-02T03:04:05.678900')"
+        == "TIMESTAMP '2019-01-02T03:04:05.678900'"
     )
 
 
@@ -80,8 +78,7 @@ def test_get_text_clause_with_colon(app_context: AppContext) -> None:
     from superset.db_engine_specs.athena import AthenaEngineSpec
 
     query = (
-        "SELECT foo FROM tbl WHERE "
-        "abc >= from_iso8601_timestamp('2021-11-26T00\:00\:00.000000')"
+        "SELECT foo FROM tbl WHERE " "abc >= TIMESTAMP '2021-11-26T00\:00\:00.000000'"
     )
     text_clause = AthenaEngineSpec.get_text_clause(query)
     assert text_clause.text == query
diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py
index eea6b6ec3c362..5bb327ff4ebc3 100644
--- a/tests/unit_tests/db_engine_specs/test_presto.py
+++ b/tests/unit_tests/db_engine_specs/test_presto.py
@@ -26,21 +26,21 @@
     "target_type,dttm,result",
     [
         ("VARCHAR", datetime(2022, 1, 1), None),
-        ("DATE", datetime(2022, 1, 1), "from_iso8601_date('2022-01-01')"),
+        ("DATE", datetime(2022, 1, 1), "DATE '2022-01-01'"),
         (
             "TIMESTAMP",
             datetime(2022, 1, 1, 1, 23, 45, 600000),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000'",
         ),
         (
             "TIMESTAMP WITH TIME ZONE",
             datetime(2022, 1, 1, 1, 23, 45, 600000),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000'",
         ),
         (
             "TIMESTAMP WITH TIME ZONE",
             datetime(2022, 1, 1, 1, 23, 45, 600000, tzinfo=pytz.UTC),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000+00:00')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000+00:00'",
         ),
     ],
 )
diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py
index 9962a0f66d0dc..1e800f3ad6db6 100644
--- a/tests/unit_tests/db_engine_specs/test_trino.py
+++ b/tests/unit_tests/db_engine_specs/test_trino.py
@@ -26,21 +26,21 @@
     "target_type,dttm,result",
     [
         ("VARCHAR", datetime(2022, 1, 1), None),
-        ("DATE", datetime(2022, 1, 1), "from_iso8601_date('2022-01-01')"),
+        ("DATE", datetime(2022, 1, 1), "DATE '2022-01-01'"),
         (
             "TIMESTAMP",
             datetime(2022, 1, 1, 1, 23, 45, 600000),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000'",
         ),
         (
             "TIMESTAMP WITH TIME ZONE",
             datetime(2022, 1, 1, 1, 23, 45, 600000),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000'",
         ),
         (
             "TIMESTAMP WITH TIME ZONE",
             datetime(2022, 1, 1, 1, 23, 45, 600000, tzinfo=pytz.UTC),
-            "from_iso8601_timestamp('2022-01-01T01:23:45.600000+00:00')",
+            "TIMESTAMP '2022-01-01T01:23:45.600000+00:00'",
         ),
     ],
 )