Skip to content

Commit

Permalink
fix(sql_lab): Add custom timestamp type for literal casting for prest…
Browse files Browse the repository at this point in the history
…o timestamps (#13082)

* Add custom timestamp type for literal casting for presto timestamps

* Remove typo in comment

* Use process_bind_params as in sqla docs

* Uncommit local superset config

* Add DATE literal casting

* Fix lint errors and change var name

* Get rid of col_type and whitespace

* Fix linting

* Fix arg type

* Fix isort lint error

* ran black and isort locally..

* accidentally removed EOF

* Dont need eof

* Use newer string formatting style from comments

Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>

* Trigger notification

* Trigger notification

Co-authored-by: Kenny Kwan <kennykwan@salesforce.com>
Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>
  • Loading branch information
3 people committed Apr 20, 2022
1 parent c763baf commit 1b55778
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 4 deletions.
16 changes: 13 additions & 3 deletions superset/db_engine_specs/presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,11 @@
from superset.models.sql_lab import Query
from superset.models.sql_types.presto_sql_types import (
Array,
Date,
Interval,
Map,
Row,
TimeStamp,
TinyInteger,
)
from superset.result_set import destringify
Expand Down Expand Up @@ -1096,10 +1098,18 @@ def where_latest_partition( # pylint: disable=too-many-arguments
if values is None:
return None

column_names = {column.get("name") for column in columns or []}
column_type_by_name = {
column.get("name"): column.get("type") for column in columns or []
}

for col_name, value in zip(col_names, values):
if col_name in column_names:
query = query.where(Column(col_name) == value)
if col_name in column_type_by_name:
if column_type_by_name.get(col_name) == "TIMESTAMP":
query = query.where(Column(col_name, TimeStamp()) == value)
elif column_type_by_name.get(col_name) == "DATE":
query = query.where(Column(col_name, Date()) == value)
else:
query = query.where(Column(col_name) == value)
return query

@classmethod
Expand Down
38 changes: 37 additions & 1 deletion superset/models/sql_types/presto_sql_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,15 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

# pylint: disable=abstract-method
from typing import Any, Dict, List, Optional, Type

from sqlalchemy.sql.sqltypes import Integer
from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.sql.sqltypes import DATE, Integer, TIMESTAMP
from sqlalchemy.sql.type_api import TypeEngine
from sqlalchemy.sql.visitors import Visitable
from sqlalchemy.types import TypeDecorator

# _compiler_dispatch is defined to help with type compilation

Expand Down Expand Up @@ -91,3 +95,35 @@ def python_type(self) -> Optional[Type[Any]]:
@classmethod
def _compiler_dispatch(cls, _visitor: Visitable, **_kw: Any) -> str:
return "ROW"


class TimeStamp(TypeDecorator):
"""
A type to extend functionality of timestamp data type.
"""

impl = TIMESTAMP

@classmethod
def process_bind_param(cls, value: str, dialect: Dialect) -> str:
"""
Used for in-line rendering of TIMESTAMP data type
as Presto does not support automatic casting.
"""
return f"TIMESTAMP '{value}'"


class Date(TypeDecorator):
"""
A type to extend functionality of date data type.
"""

impl = DATE

@classmethod
def process_bind_param(cls, value: str, dialect: Dialect) -> str:
"""
Used for in-line rendering of DATE data type
as Presto does not support automatic casting.
"""
return f"DATE '{value}'"

0 comments on commit 1b55778

Please sign in to comment.