From 21f98ddc2118f30430529437bf9774f6f107b966 Mon Sep 17 00:00:00 2001
From: Yongjie Zhao <yongjie.zhao@gmail.com>
Date: Thu, 16 Sep 2021 08:55:57 +0100
Subject: [PATCH] fix: catch exception when create connection (#16692)

* fix: catch exception when create connection

* fix lint

* added UT
---
 superset/db_engine_specs/bigquery.py   | 10 +++++++++-
 superset/models/core.py                |  5 ++++-
 tests/integration_tests/model_tests.py | 14 ++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
index 734984499fcea..ccd28e76a83fd 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 6144345bf5154..1ea60e6361fd7 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 567fdfe719b53..c8499ce59207a 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")