fix: catch exception when create connection (#16692)
* fix: catch exception when create connection
* fix lint
* added UT
diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
index 7349844..ccd28e7 100644
--- a/superset/db_engine_specs/bigquery.py
+++ b/superset/db_engine_specs/bigquery.py
@@ -17,7 +17,7 @@
import re
import urllib
from datetime import datetime
-from typing import Any, Dict, List, Optional, Pattern, Tuple, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, Pattern, Tuple, Type, TYPE_CHECKING
import pandas as pd
from apispec import APISpec
@@ -32,6 +32,7 @@
from superset.databases.schemas import encrypted_field_properties, EncryptedField
from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.exceptions import SupersetDBAPIDisconnectionError
from superset.errors import SupersetError, SupersetErrorType
from superset.sql_parse import Table
from superset.utils import core as utils
@@ -389,6 +390,13 @@
raise ValidationError("Invalid service credentials")
@classmethod
+ def get_dbapi_exception_mapping(cls) -> Dict[Type[Exception], Type[Exception]]:
+ # pylint: disable=import-error,import-outside-toplevel
+ from google.auth.exceptions import DefaultCredentialsError
+
+ return {DefaultCredentialsError: SupersetDBAPIDisconnectionError}
+
+ @classmethod
def validate_parameters(
cls, parameters: BigQueryParametersType # pylint: disable=unused-argument
) -> List[SupersetError]:
diff --git a/superset/models/core.py b/superset/models/core.py
index 6144345..1ea60e6 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -371,7 +371,10 @@
sqlalchemy_url, params, effective_username, security_manager, source
)
- return create_engine(sqlalchemy_url, **params)
+ try:
+ return create_engine(sqlalchemy_url, **params)
+ except Exception as ex:
+ raise self.db_engine_spec.get_dbapi_mapped_exception(ex)
def get_reserved_words(self) -> Set[str]:
return self.get_dialect().preparer.reserved_words
diff --git a/tests/integration_tests/model_tests.py b/tests/integration_tests/model_tests.py
index 567fdfe..c8499ce 100644
--- a/tests/integration_tests/model_tests.py
+++ b/tests/integration_tests/model_tests.py
@@ -18,6 +18,8 @@
import textwrap
import unittest
from unittest import mock
+
+from superset.exceptions import SupersetException
from tests.integration_tests.fixtures.birth_names_dashboard import (
load_birth_names_dashboard_with_slices,
)
@@ -337,6 +339,18 @@
df = main_db.get_df("USE superset; SELECT ';';", None)
self.assertEqual(df.iat[0, 0], ";")
+ @mock.patch("superset.models.core.create_engine")
+ def test_get_sqla_engine(self, mocked_create_engine):
+ model = Database(
+ database_name="test_database", sqlalchemy_uri="mysql://root@localhost",
+ )
+ model.db_engine_spec.get_dbapi_exception_mapping = mock.Mock(
+ return_value={Exception: SupersetException}
+ )
+ mocked_create_engine.side_effect = Exception()
+ with self.assertRaises(SupersetException):
+ model.get_sqla_engine()
+
class TestSqlaTableModel(SupersetTestCase):
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
