refactor: dbapi exception mapping for dbapi's

superset/db_engine_specs/base.py

@@ -32,6 +32,7 @@
     Optional,
     Pattern,
     Tuple,
+    Type,
     TYPE_CHECKING,
     Union,
 )
@@ -177,6 +178,35 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
         ),
     }
 
+    @classmethod
+    def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+        """
+        Each engine can implement and converge it's own specific exceptions into
+        SQLAlchemy DBAPI exceptions
+
+        Note: On python 3.9 this method can be changed to a classmethod property
+        without the need of implementing a metaclass type
+
+        :return: A map of driver specific exception to superset custom exceptions
+        """
+        return {}
+
+    @classmethod
+    def get_dbapi_mapped_exception(cls, exception: Exception) -> Exception:
+        """
+        Get a superset custom DBAPI exception from the driver specific exception.
+
+        Override if the engine needs to perform extra changes to the exception, for
+        example change the exception message or implement custom more complex logic
+
+        :param exception: The driver specific exception
+        :return: Superset custom DBAPI exception
+        """
+        new_exception = cls.get_dbapi_exception_mapping().get(type(exception))
+        if not new_exception:
+            return exception
+        return new_exception(str(exception))
+
     @classmethod
     def is_db_column_type_match(
         cls, db_column_type: Optional[str], target_column_type: utils.GenericDataType
@@ -314,9 +344,12 @@ def fetch_data(
         """
         if cls.arraysize:
             cursor.arraysize = cls.arraysize
-        if cls.limit_method == LimitMethod.FETCH_MANY and limit:
-            return cursor.fetchmany(limit)
-        return cursor.fetchall()
+        try:
+            if cls.limit_method == LimitMethod.FETCH_MANY and limit:
+                return cursor.fetchmany(limit)
+            return cursor.fetchall()
+        except Exception as ex:
+            raise cls.get_dbapi_mapped_exception(ex)
 
     @classmethod
     def expand_data(
@@ -902,7 +935,10 @@ def execute(cls, cursor: Any, query: str, **kwargs: Any) -> None:
         """
         if cls.arraysize:
             cursor.arraysize = cls.arraysize
-        cursor.execute(query)
+        try:
+            cursor.execute(query)
+        except Exception as ex:
+            raise cls.get_dbapi_mapped_exception(ex)
 
     @classmethod
     def make_label_compatible(cls, label: str) -> Union[str, quoted_name]:

superset/db_engine_specs/clickhouse.py

@@ -15,9 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import Optional
+from typing import Dict, Optional, Type
 
 from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.exceptions import SupersetDBAPIDatabaseError
 from superset.utils import core as utils
 
 
@@ -45,6 +46,21 @@ class ClickHouseEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
         "P1Y": "toStartOfYear(toDateTime({col}))",
     }
 
+    @classmethod
+    def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+        from urllib3.exceptions import NewConnectionError
+
+        return {NewConnectionError: SupersetDBAPIDatabaseError}
+
+    @classmethod
+    def get_dbapi_mapped_exception(cls, exception: Exception) -> Exception:
+        new_exception = cls.get_dbapi_exception_mapping().get(type(exception))
+        if new_exception == SupersetDBAPIDatabaseError:
+            return SupersetDBAPIDatabaseError("Connection failed")
+        if not new_exception:
+            return exception
+        return new_exception(str(exception))
+
     @classmethod
     def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]:
         tt = target_type.upper()

superset/db_engine_specs/elasticsearch.py

@@ -15,9 +15,14 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import Dict, Optional
+from typing import Dict, Optional, Type
 
 from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.exceptions import (
+    SupersetDBAPIDatabaseError,
+    SupersetDBAPIOperationalError,
+    SupersetDBAPIProgrammingError,
+)
 from superset.utils import core as utils
 
 
@@ -41,6 +46,16 @@ class ElasticSearchEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-metho
 
     type_code_map: Dict[int, str] = {}  # loaded from get_datatype only if needed
 
+    @classmethod
+    def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+        import es.exceptions as es_exceptions
+
+        return {
+            es_exceptions.DatabaseError: SupersetDBAPIDatabaseError,
+            es_exceptions.OperationalError: SupersetDBAPIOperationalError,
+            es_exceptions.ProgrammingError: SupersetDBAPIProgrammingError,
+        }
+
     @classmethod
     def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]:
         if target_type.upper() == utils.TemporalType.DATETIME:

superset/db_engine_specs/exceptions.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from superset.exceptions import SupersetException
+
+
+class SupersetDBAPIError(SupersetException):
+    pass
+
+
+class SupersetDBAPIDataError(SupersetDBAPIError):
+    pass
+
+
+class SupersetDBAPIDatabaseError(SupersetDBAPIError):
+    pass
+
+
+class SupersetDBAPIDisconnectionError(SupersetDBAPIError):
+    pass
+
+
+class SupersetDBAPIOperationalError(SupersetDBAPIError):
+    pass
+
+
+class SupersetDBAPIProgrammingError(SupersetDBAPIError):
+    pass

tests/db_engine_specs/clickhouse_tests.py

@@ -14,7 +14,12 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from unittest import mock
+
+import pytest
+
 from superset.db_engine_specs.clickhouse import ClickHouseEngineSpec
+from superset.db_engine_specs.exceptions import SupersetDBAPIDatabaseError
 from tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -30,3 +35,13 @@ def test_convert_dttm(self):
             ClickHouseEngineSpec.convert_dttm("DATETIME", dttm),
             "toDateTime('2019-01-02 03:04:05')",
         )
+
+    def test_execute_connection_error(self):
+        from urllib3.exceptions import NewConnectionError
+
+        cursor = mock.Mock()
+        cursor.execute.side_effect = NewConnectionError(
+            "Dummypool", message="Exception with sensitive data"
+        )
+        with pytest.raises(SupersetDBAPIDatabaseError) as ex:
+            ClickHouseEngineSpec.execute(cursor, "SELECT col1 from table1")