diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index aa52abf9e3e35..1e560e6ed5c44 100644 --- a/superset/db_engine_specs/base.py +++ b/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 its own specific exceptions into + Superset 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]: diff --git a/superset/db_engine_specs/clickhouse.py b/superset/db_engine_specs/clickhouse.py index 658aef01f2199..4db5684119042 100644 --- a/superset/db_engine_specs/clickhouse.py +++ b/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() diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index 4f6067710f671..2d52d3c4e8aad 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/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 # pylint: disable=import-error + + 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: diff --git a/superset/db_engine_specs/exceptions.py b/superset/db_engine_specs/exceptions.py new file mode 100644 index 0000000000000..6b4fb5549dd20 --- /dev/null +++ b/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 diff --git a/tests/db_engine_specs/clickhouse_tests.py b/tests/db_engine_specs/clickhouse_tests.py index b6416c8d0c8ed..c6f506cdf4ffc 100644 --- a/tests/db_engine_specs/clickhouse_tests.py +++ b/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")