diff --git a/superset/databases/api.py b/superset/databases/api.py index ca1eb3bc7de..5dd1225a007 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -27,6 +27,7 @@ from flask_babel import gettext as _ from marshmallow import ValidationError from sqlalchemy.engine.url import make_url from sqlalchemy.exc import ( + DBAPIError, NoSuchModuleError, NoSuchTableError, OperationalError, @@ -589,7 +590,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi): ) except DatabaseSecurityUnsafeError as ex: return self.response_422(message=ex) - except OperationalError: + except DBAPIError: logger.warning("Connection failed") return self.response( 500, diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py index 99e6f987217..34a60e0f256 100644 --- a/superset/databases/commands/test_connection.py +++ b/superset/databases/commands/test_connection.py @@ -19,7 +19,7 @@ from contextlib import closing from typing import Any, Dict, Optional from flask_appbuilder.security.sqla.models import User -from sqlalchemy import select +from sqlalchemy.exc import DBAPIError from superset.commands.base import BaseCommand from superset.databases.commands.exceptions import DatabaseSecurityUnsafeError @@ -54,8 +54,9 @@ class TestConnectionDatabaseCommand(BaseCommand): database.db_engine_spec.mutate_db_for_connection_test(database) username = self._actor.username if self._actor is not None else None engine = database.get_sqla_engine(user_name=username) - with closing(engine.connect()) as conn: - conn.scalar(select([1])) + with closing(engine.raw_connection()) as conn: + if not engine.dialect.do_ping(conn): + raise DBAPIError(None, None, None) except DBSecurityException as ex: logger.warning(ex) raise DatabaseSecurityUnsafeError() diff --git a/superset/views/core.py b/superset/views/core.py index da6c0b1df63..84c9cb16e46 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -34,12 +34,7 @@ from flask_babel import gettext as __, lazy_gettext as _ from jinja2.exceptions import TemplateError from sqlalchemy import and_, or_ from sqlalchemy.engine.url import make_url -from sqlalchemy.exc import ( - ArgumentError, - NoSuchModuleError, - OperationalError, - SQLAlchemyError, -) +from sqlalchemy.exc import ArgumentError, DBAPIError, NoSuchModuleError, SQLAlchemyError from sqlalchemy.orm.session import Session from werkzeug.urls import Href @@ -1152,8 +1147,10 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods engine = database.get_sqla_engine(user_name=username) with closing(engine.raw_connection()) as conn: - engine.dialect.do_ping(conn) - return json_success('"OK"') + if engine.dialect.do_ping(conn): + return json_success('"OK"') + + raise DBAPIError(None, None, None) except CertificateException as ex: logger.info("Certificate exception") return json_error_response(ex.message) @@ -1175,7 +1172,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods "'DRIVER://USER:PASSWORD@DB-HOST/DATABASE-NAME'" ) ) - except OperationalError: + except DBAPIError: logger.warning("Connection failed") return json_error_response( _("Connection failed, please check your connection settings"), 400