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_lab import Query
from superset.models.sql_types.presto_sql_types import ( from superset.models.sql_types.presto_sql_types import (
Array, Array,
Date,
Interval, Interval,
Map, Map,
Row, Row,
TimeStamp,
TinyInteger, TinyInteger,
) )
from superset.result_set import destringify from superset.result_set import destringify
@ -1096,9 +1098,17 @@ class PrestoEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-metho
if values is None: if values is None:
return 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): for col_name, value in zip(col_names, values):
if col_name in column_names: 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) query = query.where(Column(col_name) == value)
return query return query

View File

@ -14,11 +14,15 @@
# KIND, either express or implied. See the License for the # KIND, either express or implied. See the License for the
# specific language governing permissions and limitations # specific language governing permissions and limitations
# under the License. # under the License.
# pylint: disable=abstract-method
from typing import Any, Dict, List, Optional, Type 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.type_api import TypeEngine
from sqlalchemy.sql.visitors import Visitable from sqlalchemy.sql.visitors import Visitable
from sqlalchemy.types import TypeDecorator
# _compiler_dispatch is defined to help with type compilation # _compiler_dispatch is defined to help with type compilation
@ -91,3 +95,35 @@ class Row(TypeEngine):
@classmethod @classmethod
def _compiler_dispatch(cls, _visitor: Visitable, **_kw: Any) -> str: def _compiler_dispatch(cls, _visitor: Visitable, **_kw: Any) -> str:
return "ROW" 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}'"