diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index c62a9d58b434..c6faf6db6c9b 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 c4bcf759c043..41a4f4e0f38d 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'", + )