fix: Ensure Presto database engine spec correctly handles Trino

superset/db_engine_specs/presto.py

@@ -23,7 +23,18 @@
 from contextlib import closing
 from datetime import datetime
 from distutils.version import StrictVersion
-from typing import Any, cast, Dict, List, Optional, Pattern, Tuple, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    cast,
+    Dict,
+    List,
+    Optional,
+    Pattern,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    Union,
+)
 from urllib import parse
 
 import pandas as pd
@@ -35,13 +46,16 @@
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.result import Row as ResultRow
 from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import DatabaseError
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import ColumnClause, Select
 
 from superset import cache_manager, is_feature_enabled
 from superset.common.db_query_status import QueryStatus
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import BaseEngineSpec, ColumnTypeMapping
+from superset.db_engine_specs.exceptions import SupersetDBAPIDatabaseError
+
 from superset.errors import SupersetErrorType
 from superset.exceptions import SupersetTemplateException
 from superset.models.sql_lab import Query
@@ -224,6 +238,15 @@ class PrestoEngineSpec(BaseEngineSpec):  # pylint: disable=too-many-public-metho
         ),
     }
 
+    @classmethod
+    def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+        # pylint: disable=import-outside-toplevel,import-error
+        from pyhive.exc import DatabaseError
+
+        return {
+            DatabaseError: SupersetDBAPIDatabaseError,
+        }
+
     @classmethod
     def get_allow_cost_estimate(cls, extra: Dict[str, Any]) -> bool:
         version = extra.get("version")
@@ -913,21 +936,23 @@ def extra_table_metadata(
         indexes = database.get_indexes(table_name, schema_name)
         if indexes:
             cols = indexes[0].get("column_names", [])
-            full_table_name = table_name
-            if schema_name and "." not in table_name:
-                full_table_name = "{}.{}".format(schema_name, table_name)
-            pql = cls._partition_query(full_table_name, database)
-            col_names, latest_parts = cls.latest_partition(
-                table_name, schema_name, database, show_first=True
-            )
 
-            if not latest_parts:
-                latest_parts = tuple([None] * len(col_names))
-            metadata["partitions"] = {
-                "cols": cols,
-                "latest": dict(zip(col_names, latest_parts)),
-                "partitionQuery": pql,
-            }
+            if cols:
+                full_table_name = table_name
+                if schema_name and "." not in table_name:
+                    full_table_name = "{}.{}".format(schema_name, table_name)
+                pql = cls._partition_query(full_table_name, database)
+                col_names, latest_parts = cls.latest_partition(
+                    table_name, schema_name, database, show_first=True
+                )
+
+                if not latest_parts:
+                    latest_parts = tuple([None] * len(col_names))
+                metadata["partitions"] = {
+                    "cols": cols,
+                    "latest": dict(zip(col_names, latest_parts)),
+                    "partitionQuery": pql,
+                }
 
         # flake8 is not matching `Optional[str]` to `Any` for some reason...
         metadata["view"] = cast(
@@ -947,20 +972,16 @@ def get_create_view(
         :param schema: Schema name
         :param table: Table (view) name
         """
-        # pylint: disable=import-outside-toplevel
-        from pyhive.exc import DatabaseError
 
         engine = cls.get_engine(database, schema)
         with closing(engine.raw_connection()) as conn:
             cursor = conn.cursor()
             sql = f"SHOW CREATE VIEW {schema}.{table}"
             try:
                 cls.execute(cursor, sql)
-
-            except DatabaseError:  # not a VIEW
+                return cls.fetch_data(cursor, 1)[0][0]
+            except SupersetDBAPIDatabaseError:  # not a VIEW
                 return None
-            rows = cls.fetch_data(cursor, 1)
-        return rows[0][0]
 
     @classmethod
     def get_tracking_url(cls, cursor: "Cursor") -> Optional[str]:

superset/db_engine_specs/trino.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING
 
 import simplejson as json
 from flask import current_app
@@ -25,6 +25,7 @@
 
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.exceptions import SupersetDBAPIDatabaseError
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.models.sql_lab import Query
 from superset.utils import core as utils
@@ -45,6 +46,15 @@ class TrinoEngineSpec(PrestoEngineSpec):
     engine_aliases = {"trinonative"}  # Required for backwards compatibility.
     engine_name = "Trino"
 
+    @classmethod
+    def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+        # pylint: disable=import-outside-toplevel,import-error
+        from trino.exceptions import DatabaseError
+
+        return {
+            DatabaseError: SupersetDBAPIDatabaseError,
+        }
+
     @classmethod
     def update_impersonation_config(
         cls,

tests/integration_tests/db_engine_specs/presto_tests.py

@@ -19,9 +19,11 @@
 
 import pandas as pd
 from sqlalchemy import types
+from sqlalchemy.exc import DatabaseError
 from sqlalchemy.sql import select
 
 from superset.db_engine_specs.presto import PrestoEngineSpec
+from superset.db_engine_specs.exceptions import SupersetDBAPIDatabaseError
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.sql_parse import ParsedQuery
 from superset.utils.core import DatasourceName, GenericDataType
@@ -895,9 +897,7 @@ def test_get_create_view_exception(self):
             PrestoEngineSpec.get_create_view(database, schema=schema, table=table)
 
     def test_get_create_view_database_error(self):
-        from pyhive.exc import DatabaseError
-
-        mock_execute = mock.MagicMock(side_effect=DatabaseError())
+        mock_execute = mock.MagicMock(side_effect=SupersetDBAPIDatabaseError())
         database = mock.MagicMock()
         database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
             mock_execute