fix(presto preview): re-enable schema previsualization for Trino/Pres…

…to table/schemas (#26782) (cherry picked from commit afdf038)
apache · May 15, 2024 · fc71e98 · fc71e98
1 parent 78d5b47
commit fc71e98
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git 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
@@ -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(