diff --git a/setup.py b/setup.py index 198a2578bcc73..e39284f147e2f 100644 --- a/setup.py +++ b/setup.py @@ -140,6 +140,7 @@ def get_git_sha(): "pinot": ["pinotdb>=0.3.3, <0.4"], "postgres": ["psycopg2-binary==2.8.5"], "presto": ["pyhive[presto]>=0.4.0"], + "trino": ["sqlalchemy-trino>=0.2"], "prophet": ["fbprophet>=0.6, <0.7"], "redshift": ["sqlalchemy-redshift>=0.8.1, < 0.9"], "snowflake": ["snowflake-sqlalchemy>=1.2.3, <1.3"], diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py new file mode 100644 index 0000000000000..5913be0b39855 --- /dev/null +++ b/superset/db_engine_specs/trino.py @@ -0,0 +1,58 @@ +# 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 + +from superset.db_engine_specs import BaseEngineSpec +from superset.utils import core as utils + + +class TrinoEngineSpec(BaseEngineSpec): + engine = "trino" + engine_name = "Trino" + + # pylint: disable=line-too-long + _time_grain_expressions = { + None: "{col}", + "PT1S": "date_trunc('second', CAST({col} AS TIMESTAMP))", + "PT1M": "date_trunc('minute', CAST({col} AS TIMESTAMP))", + "PT1H": "date_trunc('hour', CAST({col} AS TIMESTAMP))", + "P1D": "date_trunc('day', CAST({col} AS TIMESTAMP))", + "P1W": "date_trunc('week', CAST({col} AS TIMESTAMP))", + "P1M": "date_trunc('month', CAST({col} AS TIMESTAMP))", + "P0.25Y": "date_trunc('quarter', CAST({col} AS TIMESTAMP))", + "P1Y": "date_trunc('year', CAST({col} AS TIMESTAMP))", + # "1969-12-28T00:00:00Z/P1W", # Week starting Sunday + # "1969-12-29T00:00:00Z/P1W", # Week starting Monday + # "P1W/1970-01-03T00:00:00Z", # Week ending Saturday + # "P1W/1970-01-04T00:00:00Z", # Week ending Sunday + } + + @classmethod + def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: + 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 None + + @classmethod + def epoch_to_dttm(cls) -> str: + return "from_unixtime({col})" diff --git a/tests/db_engine_specs/trino_tests.py b/tests/db_engine_specs/trino_tests.py new file mode 100644 index 0000000000000..2f827721fa07b --- /dev/null +++ b/tests/db_engine_specs/trino_tests.py @@ -0,0 +1,34 @@ +# 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 superset.db_engine_specs.trino import TrinoEngineSpec +from tests.db_engine_specs.base_tests import TestDbEngineSpec + + +class TestPrestoDbEngineSpec(TestDbEngineSpec): + def test_convert_dttm(self): + dttm = self.get_dttm() + + self.assertEqual( + TrinoEngineSpec.convert_dttm("DATE", dttm), + "from_iso8601_date('2019-01-02')", + ) + + self.assertEqual( + TrinoEngineSpec.convert_dttm("TIMESTAMP", dttm), + "from_iso8601_timestamp('2019-01-02T03:04:05.678900')", + )