diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 734984499fce..ccd28e76a83f 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 @@ -388,6 +389,13 @@ def get_parameters_from_uri( 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 diff --git a/superset/models/core.py b/superset/models/core.py index 6144345bf515..1ea60e6361fd 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -371,7 +371,10 @@ def get_sqla_engine( 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 567fdfe719b5..c8499ce59207 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 @@ def test_multi_statement(self): 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")