diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index 11c221953a5..3b3a8e41f67 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -73,7 +73,11 @@ class TestConnectionDatabaseCommand(BaseCommand): username = self._actor.username if self._actor is not None else None engine = database.get_sqla_engine(user_name=username) with closing(engine.raw_connection()) as conn: - if not engine.dialect.do_ping(conn): + try: + alive = engine.dialect.do_ping(conn) + except Exception: # pylint: disable=broad-except + alive = False + if not alive: raise DBAPIError(None, None, None) # Log succesful connection test with engine diff --git a/tests/databases/commands_tests.py b/tests/databases/commands_tests.py index 0ff25194b62..da475d11ab1 100644 --- a/tests/databases/commands_tests.py +++ b/tests/databases/commands_tests.py @@ -37,7 +37,7 @@ from superset.databases.commands.export import ExportDatabasesCommand from superset.databases.commands.importers.v1 import ImportDatabasesCommand from superset.databases.commands.test_connection import TestConnectionDatabaseCommand from superset.databases.schemas import DatabaseTestConnectionSchema -from superset.errors import SupersetError +from superset.errors import SupersetError, SupersetErrorType from superset.exceptions import SupersetSecurityException from superset.models.core import Database from superset.utils.core import backend, get_example_database @@ -547,6 +547,31 @@ class TestTestConnectionDatabaseCommand(SupersetTestCase): ) mock_event_logger.assert_called() + @mock.patch("superset.databases.dao.Database.get_sqla_engine") + @mock.patch( + "superset.databases.commands.test_connection.event_logger.log_with_context" + ) + def test_connection_do_ping_exception( + self, mock_event_logger, mock_get_sqla_engine + ): + """Test to make sure do_ping exceptions gets captured""" + database = get_example_database() + mock_get_sqla_engine.return_value.dialect.do_ping.side_effect = Exception( + "An error has occurred!" + ) + db_uri = database.sqlalchemy_uri_decrypted + json_payload = {"sqlalchemy_uri": db_uri} + command_without_db_name = TestConnectionDatabaseCommand( + security_manager.find_user("admin"), json_payload + ) + + with pytest.raises(DatabaseTestConnectionFailedError) as excinfo: + command_without_db_name.run() + assert ( + excinfo.value.errors[0].error_type + == SupersetErrorType.GENERIC_DB_ENGINE_ERROR + ) + @mock.patch("superset.databases.dao.Database.get_sqla_engine") @mock.patch( "superset.databases.commands.test_connection.event_logger.log_with_context"