Skip to content

Commit

Permalink
fix: catch exception when create connection (#16692)
Browse files Browse the repository at this point in the history
* fix: catch exception when create connection

* fix lint

* added UT
  • Loading branch information
zhaoyongjie authored Sep 16, 2021
1 parent 092ef5b commit 21f98dd
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 2 deletions.
10 changes: 9 additions & 1 deletion superset/db_engine_specs/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion superset/models/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions tests/integration_tests/model_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 21f98dd

Please sign in to comment.