chore: Update testconn logic (#11315)

* chore: Update testconn logic

* [databases] Address @dpgaspar's comments

Co-authored-by: John Bodley <john.bodley@airbnb.com>
This commit is contained in:
John Bodley 2020-10-27 02:15:05 -07:00 committed by GitHub
parent ecdff724ea
commit 8ee6493233
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 12 additions and 13 deletions

View File

@ -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,

View File

@ -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()

View File

@ -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