diff --git a/setup.py b/setup.py index d082733fa38d4..19a08951df0bd 100644 --- a/setup.py +++ b/setup.py @@ -147,6 +147,7 @@ def get_git_sha() -> str: "trino": ["sqlalchemy-trino>=0.2"], "prophet": ["prophet>=1.0.1, <1.1", "pystan<3.0"], "redshift": ["sqlalchemy-redshift>=0.8.1, < 0.9"], + "rockset": ["rockset>=0.7.68, <0.8"], "snowflake": ["snowflake-sqlalchemy>=1.2.3, <1.3"], "teradata": ["sqlalchemy-teradata==0.9.0.dev0"], "thumbnails": ["Pillow>=7.0.0, <8.0.0"], diff --git a/superset/db_engine_specs/rockset.py b/superset/db_engine_specs/rockset.py new file mode 100644 index 0000000000000..27e626776d264 --- /dev/null +++ b/superset/db_engine_specs/rockset.py @@ -0,0 +1,69 @@ +# 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, TYPE_CHECKING + +from superset.db_engine_specs.base import BaseEngineSpec +from superset.utils import core as utils + +if TYPE_CHECKING: + from superset.connectors.sqla.models import TableColumn + + +class RocksetEngineSpec(BaseEngineSpec): + + engine = "rockset" + engine_name = "Rockset" + + _time_grain_expressions = { + None: "{col}", + "PT1S": "DATE_TRUNC('second', {col})", + "PT1M": "DATE_TRUNC('minute', {col})", + "PT1H": "DATE_TRUNC('hour', {col})", + "P1D": "DATE_TRUNC('day', {col})", + "P1W": "DATE_TRUNC('week', {col})", + "P1M": "DATE_TRUNC('month', {col})", + "P0.25Y": "DATE_TRUNC('quarter', {col})", + "P1Y": "DATE_TRUNC('year', {col})", + } + + @classmethod + def epoch_to_dttm(cls) -> str: + return "{col} * 1000" + + @classmethod + def epoch_ms_to_dttm(cls) -> str: + return "{col}" + + @classmethod + def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: + print(cls) + tt = target_type.upper() + if tt == utils.TemporalType.DATE: + return f"DATE '{dttm.date().isoformat()}'" + if tt == utils.TemporalType.DATETIME: + dttm_formatted = dttm.isoformat(sep=" ", timespec="microseconds") + return f"""DATETIME '{dttm_formatted}'""" + if tt == utils.TemporalType.TIMESTAMP: + dttm_formatted = dttm.isoformat(timespec="microseconds") + return f"""TIMESTAMP '{dttm_formatted}'""" + return None + + @classmethod + def alter_new_orm_column(cls, orm_col: "TableColumn") -> None: + if orm_col.type == "TIMESTAMP": + orm_col.python_date_format = "epoch_ms"