[incubator-superset] branch master updated: chore: Update testconn logic (#11315)

[dp...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dpgaspar pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 8ee6493  chore: Update testconn logic (#11315)
8ee6493 is described below

commit 8ee649323375d074b5d22b2e7e1f51e0183fe476
Author: John Bodley <45...@users.noreply.github.com>
AuthorDate: Tue Oct 27 02:15:05 2020 -0700

    chore: Update testconn logic (#11315)
    
    * chore: Update testconn logic
    
    * [databases] Address @dpgaspar's comments
    
    Co-authored-by: John Bodley <jo...@airbnb.com>
---
 superset/databases/api.py                      |  3 ++-
 superset/databases/commands/test_connection.py |  7 ++++---
 superset/views/core.py                         | 15 ++++++---------
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/superset/databases/api.py b/superset/databases/api.py
index ca1eb3b..5dd1225 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 99e6f98..34a60e0 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 da6c0b1..84c9cb1 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