fix(sql_lab): Add custom timestamp type for literal casting for presto 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>
This commit is contained in:
Kenny Kwan 2022-04-20 15:01:24 -07:00 committed by GitHub
parent c763baf09e
commit 1b55778427
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 50 additions and 4 deletions

View File

@ -47,9 +47,11 @@ from superset.exceptions import SupersetTemplateException
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
@ -1096,10 +1098,18 @@ class PrestoEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-metho
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

View File

@ -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
@ -91,3 +95,35 @@ class Row(TypeEngine):
@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}'"