From cef7a1e781a49eede1b1e2a8ad3bca1ff9f04194 Mon Sep 17 00:00:00 2001
From: Michiel De Smet <michiel.desmet@starburstdata.com>
Date: Fri, 7 Oct 2022 16:31:05 +0200
Subject: [PATCH] Fix Trino timestamp conversion

---
 superset/db_engine_specs/trino.py             | 25 ++++++++++++++++
 .../db_engine_specs/trino_tests.py            | 30 +++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
index c62a9d58b4343..c6faf6db6c9b9 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -17,6 +17,8 @@
 from __future__ import annotations
 
 import logging
+import re
+from datetime import datetime
 from typing import Any, Dict, Optional, Type, TYPE_CHECKING
 
 import simplejson as json
@@ -47,6 +49,29 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
     engine = "trino"
     engine_name = "Trino"
 
+    @classmethod
+    def convert_dttm(
+        cls, target_type: str, dttm: datetime, db_extra: Optional[Dict[str, Any]] = None
+    ) -> Optional[str]:
+        """
+        Convert a Python `datetime` object to a SQL expression.
+        :param target_type: The target type of expression
+        :param dttm: The datetime object
+        :param db_extra: The database extra object
+        :return: The SQL expression
+        Superset only defines time zone naive `datetime` objects, though this method
+        handles both time zone naive and aware conversions.
+        """
+        tt = target_type.upper()
+        if tt == utils.TemporalType.DATE:
+            return f"DATE '{dttm.date().isoformat()}'"
+        if re.sub(r"\(\d\)", "", tt) in (
+            utils.TemporalType.TIMESTAMP,
+            utils.TemporalType.TIMESTAMP_WITH_TIME_ZONE,
+        ):
+            return f"""TIMESTAMP '{dttm.isoformat(timespec="microseconds", sep=" ")}'"""
+        return None
+
     @classmethod
     def extra_table_metadata(
         cls,
diff --git a/tests/integration_tests/db_engine_specs/trino_tests.py b/tests/integration_tests/db_engine_specs/trino_tests.py
index c4bcf759c043e..41a4f4e0f38da 100644
--- a/tests/integration_tests/db_engine_specs/trino_tests.py
+++ b/tests/integration_tests/db_engine_specs/trino_tests.py
@@ -19,10 +19,12 @@
 from unittest.mock import Mock, patch
 
 import pytest
+from sqlalchemy import types
 
 import superset.config
 from superset.constants import USER_AGENT
 from superset.db_engine_specs.trino import TrinoEngineSpec
+from superset.utils.core import GenericDataType
 from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -166,3 +168,31 @@ def test_auth_custom_auth_denied(self):
             f"For security reason, custom authentication '{auth_method}' "
             f"must be listed in 'ALLOWED_EXTRA_AUTHENTICATIONS' config"
         )
+
+    def test_convert_dttm(self):
+        dttm = self.get_dttm()
+
+        self.assertEqual(
+            TrinoEngineSpec.convert_dttm("TIMESTAMP", dttm),
+            "TIMESTAMP '2019-01-02 03:04:05.678900'",
+        )
+
+        self.assertEqual(
+            TrinoEngineSpec.convert_dttm("TIMESTAMP(3)", dttm),
+            "TIMESTAMP '2019-01-02 03:04:05.678900'",
+        )
+
+        self.assertEqual(
+            TrinoEngineSpec.convert_dttm("TIMESTAMP WITH TIME ZONE", dttm),
+            "TIMESTAMP '2019-01-02 03:04:05.678900'",
+        )
+
+        self.assertEqual(
+            TrinoEngineSpec.convert_dttm("TIMESTAMP(3) WITH TIME ZONE", dttm),
+            "TIMESTAMP '2019-01-02 03:04:05.678900'",
+        )
+
+        self.assertEqual(
+            TrinoEngineSpec.convert_dttm("DATE", dttm),
+            "DATE '2019-01-02'",
+        )