From 82a6811e7ef3b00e5abbead15377481298764891 Mon Sep 17 00:00:00 2001 From: John Bodley <4567245+john-bodley@users.noreply.github.com> Date: Tue, 22 Mar 2022 06:33:38 +1300 Subject: [PATCH] fix(presto/trino): Add TIME/TIMESTAMP WITH TIME ZONE (#19263) Co-authored-by: John Bodley --- superset/db_engine_specs/base.py | 2 +- superset/db_engine_specs/presto.py | 16 +++++- superset/db_engine_specs/trino.py | 22 ++++++-- superset/utils/core.py | 2 + .../unit_tests/db_engine_specs/test_presto.py | 53 +++++++++++++++++++ .../unit_tests/db_engine_specs/test_trino.py | 53 +++++++++++++++++++ 6 files changed, 141 insertions(+), 7 deletions(-) create mode 100644 tests/unit_tests/db_engine_specs/test_presto.py create mode 100644 tests/unit_tests/db_engine_specs/test_trino.py diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index e867f200d91a8..5c73e2f666e26 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -812,7 +812,7 @@ def convert_dttm( # pylint: disable=unused-argument cls, target_type: str, dttm: datetime, db_extra: Optional[Dict[str, Any]] = None ) -> Optional[str]: """ - Convert Python datetime object to a SQL expression + Convert a Python `datetime` object to a SQL expression. :param target_type: The target type of expression :param dttm: The datetime object diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 376151587cdee..62e2f349f1e4e 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -725,10 +725,24 @@ def adjust_database_uri( 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"""from_iso8601_date('{dttm.date().isoformat()}')""" - if tt == utils.TemporalType.TIMESTAMP: + 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 None diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 4e5f153ad2ab2..d902a917a6a11 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -57,13 +57,25 @@ class TrinoEngineSpec(BaseEngineSpec): 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: - value = dttm.date().isoformat() - return f"from_iso8601_date('{value}')" - if tt == utils.TemporalType.TIMESTAMP: - value = dttm.isoformat(timespec="microseconds") - return f"from_iso8601_timestamp('{value}')" + return f"from_iso8601_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 None @classmethod diff --git a/superset/utils/core.py b/superset/utils/core.py index d2527a32c72cf..b6ae3272cfed5 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -325,7 +325,9 @@ class TemporalType(str, Enum): SMALLDATETIME = "SMALLDATETIME" TEXT = "TEXT" TIME = "TIME" + TIME_WITH_TIME_ZONE = "TIME WITH TIME ZONE" TIMESTAMP = "TIMESTAMP" + TIMESTAMP_WITH_TIME_ZONE = "TIMESTAMP WITH TIME ZONE" class ColumnTypeSource(Enum): diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py new file mode 100644 index 0000000000000..370af3f48d604 --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_presto.py @@ -0,0 +1,53 @@ +# 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 Optional + +import pytest +import pytz +from flask.ctx import AppContext + + +@pytest.mark.parametrize( + "target_type,dttm,result", + [ + ("VARCHAR", datetime(2022, 1, 1), None), + ("DATE", datetime(2022, 1, 1), "from_iso8601_date('2022-01-01')"), + ( + "TIMESTAMP", + datetime(2022, 1, 1, 1, 23, 45, 600000), + "from_iso8601_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 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')", + ), + ], +) +def test_convert_dttm( + app_context: AppContext, target_type: str, dttm: datetime, result: Optional[str], +) -> None: + from superset.db_engine_specs.presto import PrestoEngineSpec + + for case in (str.lower, str.upper): + assert PrestoEngineSpec.convert_dttm(case(target_type), dttm) == result diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py new file mode 100644 index 0000000000000..ff00c4ff4595d --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -0,0 +1,53 @@ +# 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 Optional + +import pytest +import pytz +from flask.ctx import AppContext + + +@pytest.mark.parametrize( + "target_type,dttm,result", + [ + ("VARCHAR", datetime(2022, 1, 1), None), + ("DATE", datetime(2022, 1, 1), "from_iso8601_date('2022-01-01')"), + ( + "TIMESTAMP", + datetime(2022, 1, 1, 1, 23, 45, 600000), + "from_iso8601_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 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')", + ), + ], +) +def test_convert_dttm( + app_context: AppContext, target_type: str, dttm: datetime, result: Optional[str], +) -> None: + from superset.db_engine_specs.trino import TrinoEngineSpec + + for case in (str.lower, str.upper): + assert TrinoEngineSpec.convert_dttm(case(target_type), dttm) == result