From 1b55778427cdb5e4b40074536a3ae2a597f30a69 Mon Sep 17 00:00:00 2001
From: Kenny Kwan <kenny.kwan9@gmail.com>
Date: Wed, 20 Apr 2022 15:01:24 -0700
Subject: [PATCH] fix(sql_lab): Add custom timestamp type for literal casting
 for presto timestamps (#13082)

* Add custom timestamp type for literal casting for presto timestamps

* Remove typo in comment

* Use process_bind_params as in sqla docs

* Uncommit local superset config

* Add DATE literal casting

* Fix lint errors and change var name

* Get rid of col_type and whitespace

* Fix linting

* Fix arg type

* Fix isort lint error

* ran black and isort locally..

* accidentally removed EOF

* Dont need eof

* Use newer string formatting style from comments

Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>

* Trigger notification

* Trigger notification

Co-authored-by: Kenny Kwan <kennykwan@salesforce.com>
Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>
---
 superset/db_engine_specs/presto.py            | 16 ++++++--
 superset/models/sql_types/presto_sql_types.py | 38 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index 8675607848..645dd32d26 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -47,9 +47,11 @@ from superset.exceptions import SupersetTemplateException
 from superset.models.sql_lab import Query
 from superset.models.sql_types.presto_sql_types import (
     Array,
+    Date,
     Interval,
     Map,
     Row,
+    TimeStamp,
     TinyInteger,
 )
 from superset.result_set import destringify
@@ -1096,10 +1098,18 @@ class PrestoEngineSpec(BaseEngineSpec):  # pylint: disable=too-many-public-metho
         if values is None:
             return None
 
-        column_names = {column.get("name") for column in columns or []}
+        column_type_by_name = {
+            column.get("name"): column.get("type") for column in columns or []
+        }
+
         for col_name, value in zip(col_names, values):
-            if col_name in column_names:
-                query = query.where(Column(col_name) == value)
+            if col_name in column_type_by_name:
+                if column_type_by_name.get(col_name) == "TIMESTAMP":
+                    query = query.where(Column(col_name, TimeStamp()) == value)
+                elif column_type_by_name.get(col_name) == "DATE":
+                    query = query.where(Column(col_name, Date()) == value)
+                else:
+                    query = query.where(Column(col_name) == value)
         return query
 
     @classmethod
diff --git a/superset/models/sql_types/presto_sql_types.py b/superset/models/sql_types/presto_sql_types.py
index a314639ca6..5f36266cca 100644
--- a/superset/models/sql_types/presto_sql_types.py
+++ b/superset/models/sql_types/presto_sql_types.py
@@ -14,11 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+# pylint: disable=abstract-method
 from typing import Any, Dict, List, Optional, Type
 
-from sqlalchemy.sql.sqltypes import Integer
+from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.sql.sqltypes import DATE, Integer, TIMESTAMP
 from sqlalchemy.sql.type_api import TypeEngine
 from sqlalchemy.sql.visitors import Visitable
+from sqlalchemy.types import TypeDecorator
 
 # _compiler_dispatch is defined to help with type compilation
 
@@ -91,3 +95,35 @@ class Row(TypeEngine):
     @classmethod
     def _compiler_dispatch(cls, _visitor: Visitable, **_kw: Any) -> str:
         return "ROW"
+
+
+class TimeStamp(TypeDecorator):
+    """
+    A type to extend functionality of timestamp data type.
+    """
+
+    impl = TIMESTAMP
+
+    @classmethod
+    def process_bind_param(cls, value: str, dialect: Dialect) -> str:
+        """
+        Used for in-line rendering of TIMESTAMP data type
+        as Presto does not support automatic casting.
+        """
+        return f"TIMESTAMP '{value}'"
+
+
+class Date(TypeDecorator):
+    """
+    A type to extend functionality of date data type.
+    """
+
+    impl = DATE
+
+    @classmethod
+    def process_bind_param(cls, value: str, dialect: Dialect) -> str:
+        """
+        Used for in-line rendering of DATE data type
+        as Presto does not support automatic casting.
+        """
+        return f"DATE '{value}'"