From e0217e9b84a9b87a40f46d65b9de9134ebacd0da Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Mon, 1 Feb 2021 14:12:27 +0000 Subject: [PATCH 1/4] refactor: dbapi exception mapping for dbapi's --- superset/db_engine_specs/base.py | 44 ++++++++++++++++++++--- superset/db_engine_specs/clickhouse.py | 18 +++++++++- superset/db_engine_specs/elasticsearch.py | 17 ++++++++- superset/db_engine_specs/exceptions.py | 41 +++++++++++++++++++++ tests/db_engine_specs/clickhouse_tests.py | 16 +++++++++ 5 files changed, 130 insertions(+), 6 deletions(-) create mode 100644 superset/db_engine_specs/exceptions.py diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index aa52abf9e3e35..6c62126aacd23 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 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]: 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..4f84ff9edb11d 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 + + 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..b16b81c0b7ca7 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,14 @@ 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") + assert str(ex) == "xpyo" From 91d71cb267878663c5c94e16db04a43a9b1f0de8 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Mon, 1 Feb 2021 14:43:49 +0000 Subject: [PATCH 2/4] fix test --- tests/db_engine_specs/clickhouse_tests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/db_engine_specs/clickhouse_tests.py b/tests/db_engine_specs/clickhouse_tests.py index b16b81c0b7ca7..c6f506cdf4ffc 100644 --- a/tests/db_engine_specs/clickhouse_tests.py +++ b/tests/db_engine_specs/clickhouse_tests.py @@ -45,4 +45,3 @@ def test_execute_connection_error(self): ) with pytest.raises(SupersetDBAPIDatabaseError) as ex: ClickHouseEngineSpec.execute(cursor, "SELECT col1 from table1") - assert str(ex) == "xpyo" From a13035264fa528b919bff0a65411bdb8970d620d Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Mon, 1 Feb 2021 16:30:01 +0000 Subject: [PATCH 3/4] fix lint --- superset/db_engine_specs/elasticsearch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index 4f84ff9edb11d..2d52d3c4e8aad 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/superset/db_engine_specs/elasticsearch.py @@ -48,7 +48,7 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho @classmethod def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]: - import es.exceptions as es_exceptions + import es.exceptions as es_exceptions # pylint: disable=import-error return { es_exceptions.DatabaseError: SupersetDBAPIDatabaseError, From c99864b95af000211016bb55790af41a8cc45eaa Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 2 Feb 2021 10:57:39 +0000 Subject: [PATCH 4/4] fix grammar on comment --- superset/db_engine_specs/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 6c62126aacd23..1e560e6ed5c44 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -181,8 +181,8 @@ 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 + 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