From fc71e986e2d2843ce0aab78101684568c4eeaad2 Mon Sep 17 00:00:00 2001 From: Balthazar Rouberol Date: Tue, 14 May 2024 00:16:12 +0200 Subject: [PATCH] fix(presto preview): re-enable schema previsualization for Trino/Presto table/schemas (#26782) (cherry picked from commit afdf0389a309345e934a6f47a8e13bb28c38d54a) --- superset/db_engine_specs/presto.py | 5 ++++- tests/unit_tests/db_engine_specs/test_presto.py | 8 ++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 44f8f9668a224..561c91cbb65c7 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -507,7 +507,10 @@ def where_latest_partition( # pylint: disable=too-many-arguments } for col_name, value in zip(col_names, values): - col_type = column_type_by_name.get(col_name) + col_type = None + if col_type_name := column_type_by_name.get(col_name): + if col_type_class := getattr(types, col_type_name, None): + col_type = col_type_class() if isinstance(col_type, types.DATE): col_type = Date() diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py index 8d57d4ed1a8c3..04eade7fd1681 100644 --- a/tests/unit_tests/db_engine_specs/test_presto.py +++ b/tests/unit_tests/db_engine_specs/test_presto.py @@ -115,10 +115,10 @@ def test_get_schema_from_engine_params() -> None: @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"), + ("DATE", "2023-05-01", "DATE '2023-05-01'"), + ("TIMESTAMP", "2023-05-01", "TIMESTAMP '2023-05-01'"), + ("VARCHAR", "2023-05-01", "'2023-05-01'"), + ("INT", 1234, "1234"), ], ) def test_where_latest_partition(