diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 1b46bb35bb4..a7ff8622722 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1276,7 +1276,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods schema: Optional[str], database: Database, query: Select, - columns: Optional[List[Dict[str, str]]] = None, + columns: Optional[List[Dict[str, Any]]] = None, ) -> Optional[Select]: """ Add a where clause to a query to reference only the most recent partition diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 9f9f257a115..6d8986c1c7b 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -418,7 +418,7 @@ class HiveEngineSpec(PrestoEngineSpec): schema: Optional[str], database: "Database", query: Select, - columns: Optional[List[Dict[str, str]]] = None, + columns: Optional[List[Dict[str, Any]]] = None, ) -> Optional[Select]: try: col_names, values = cls.latest_partition( diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 42f6ed9af6d..82b05e53e31 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -495,7 +495,7 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta): schema: Optional[str], database: Database, query: Select, - columns: Optional[List[Dict[str, str]]] = None, + columns: Optional[List[Dict[str, Any]]] = None, ) -> Optional[Select]: try: col_names, values = cls.latest_partition( @@ -513,13 +513,15 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta): } for col_name, value in zip(col_names, values): - 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) + col_type = column_type_by_name.get(col_name) + + if isinstance(col_type, types.DATE): + col_type = Date() + elif isinstance(col_type, types.TIMESTAMP): + col_type = TimeStamp() + + query = query.where(Column(col_name, col_type) == value) + return query @classmethod diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py index 2684db05555..df2ed58c376 100644 --- a/tests/unit_tests/db_engine_specs/test_presto.py +++ b/tests/unit_tests/db_engine_specs/test_presto.py @@ -16,10 +16,12 @@ # under the License. from datetime import datetime from typing import Any, Dict, Optional, Type +from unittest import mock import pytest import pytz -from sqlalchemy import types +from pyhive.sqlalchemy_presto import PrestoDialect +from sqlalchemy import sql, text, types from sqlalchemy.engine.url import make_url from superset.utils.core import GenericDataType @@ -106,3 +108,38 @@ def test_get_schema_from_engine_params() -> None: ) is None ) + + +@mock.patch("superset.db_engine_specs.presto.PrestoEngineSpec.latest_partition") +@pytest.mark.parametrize( + ["column_type", "column_value", "expected_value"], + [ + (types.DATE(), "2023-05-01", "DATE '2023-05-01'"), + (types.TIMESTAMP(), "2023-05-01", "TIMESTAMP '2023-05-01'"), + (types.VARCHAR(), "2023-05-01", "'2023-05-01'"), + (types.INT(), 1234, "1234"), + ], +) +def test_where_latest_partition( + mock_latest_partition, column_type, column_value: Any, expected_value: str +) -> None: + """ + Test the ``where_latest_partition`` method + """ + from superset.db_engine_specs.presto import PrestoEngineSpec as spec + + mock_latest_partition.return_value = (["partition_key"], [column_value]) + + query = sql.select(text("* FROM table")) + columns = [{"name": "partition_key", "type": column_type}] + + expected = f"""SELECT * FROM table \nWHERE "partition_key" = {expected_value}""" + result = spec.where_latest_partition( + "table", mock.MagicMock(), mock.MagicMock(), query, columns + ) + assert result is not None + actual = result.compile( + dialect=PrestoDialect(), compile_kwargs={"literal_binds": True} + ) + + assert str(actual) == expected