From 5111011de9de614e68c3c373dc9e938a9df3791f Mon Sep 17 00:00:00 2001
From: Phillip Kelley-Dotson <pkelleydotson@yahoo.com>
Date: Tue, 17 May 2022 10:04:52 -0700
Subject: [PATCH] fix: dbmodal test connection error timeout (#20068)

* fix: check for connect on ping

* clean up merge

* fix merge

* precommit
---
 .../databases/commands/test_connection.py     | 67 ++++++++++---------
 1 file changed, 37 insertions(+), 30 deletions(-)

diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py
index 706697412859..2e217ab01a80 100644
--- a/superset/databases/commands/test_connection.py
+++ b/superset/databases/commands/test_connection.py
@@ -23,6 +23,7 @@
 from flask_appbuilder.security.sqla.models import User
 from flask_babel import gettext as _
 from func_timeout import func_timeout, FunctionTimedOut
+from sqlalchemy.engine import Engine
 from sqlalchemy.exc import DBAPIError, NoSuchModuleError
 
 from superset.commands.base import BaseCommand
@@ -82,36 +83,41 @@ def run(self) -> None:
                     action="test_connection_attempt",
                     engine=database.db_engine_spec.__name__,
                 )
-                with closing(engine.raw_connection()) as conn:
-                    try:
-                        alive = func_timeout(
-                            int(
-                                app.config[
-                                    "TEST_DATABASE_CONNECTION_TIMEOUT"
-                                ].total_seconds()
-                            ),
-                            engine.dialect.do_ping,
-                            args=(conn,),
-                        )
-                    except (sqlite3.ProgrammingError, RuntimeError):
-                        # SQLite can't run on a separate thread, so ``func_timeout`` fails
-                        # RuntimeError catches the equivalent error from duckdb.
-                        alive = engine.dialect.do_ping(conn)
-                    except FunctionTimedOut as ex:
-                        raise SupersetTimeoutException(
-                            error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
-                            message=(
-                                "Please check your connection details and database "
-                                "settings, and ensure that your database is accepting "
-                                "connections, then try connecting again."
-                            ),
-                            level=ErrorLevel.ERROR,
-                            extra={"sqlalchemy_uri": database.sqlalchemy_uri},
-                        ) from ex
-                    except Exception:  # pylint: disable=broad-except
-                        alive = False
-                    if not alive:
-                        raise DBAPIError(None, None, None)
+
+                def ping(engine: Engine) -> bool:
+                    with closing(engine.raw_connection()) as conn:
+                        return engine.dialect.do_ping(conn)
+
+                try:
+                    alive = func_timeout(
+                        int(
+                            app.config[
+                                "TEST_DATABASE_CONNECTION_TIMEOUT"
+                            ].total_seconds()
+                        ),
+                        ping,
+                        args=(engine,),
+                    )
+
+                except (sqlite3.ProgrammingError, RuntimeError):
+                    # SQLite can't run on a separate thread, so ``func_timeout`` fails
+                    # RuntimeError catches the equivalent error from duckdb.
+                    alive = engine.dialect.do_ping(engine)
+                except FunctionTimedOut as ex:
+                    raise SupersetTimeoutException(
+                        error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
+                        message=(
+                            "Please check your connection details and database settings, "
+                            "and ensure that your database is accepting connections, "
+                            "then try connecting again."
+                        ),
+                        level=ErrorLevel.ERROR,
+                        extra={"sqlalchemy_uri": database.sqlalchemy_uri},
+                    ) from ex
+                except Exception:  # pylint: disable=broad-except
+                    alive = False
+                if not alive:
+                    raise DBAPIError(None, None, None)
 
             # Log succesful connection test with engine
             event_logger.log_with_context(
@@ -144,6 +150,7 @@ def run(self) -> None:
             )
             raise DatabaseSecurityUnsafeError(message=str(ex)) from ex
         except SupersetTimeoutException as ex:
+
             event_logger.log_with_context(
                 action=f"test_connection_error.{ex.__class__.__name__}",
                 engine=database.db_engine_spec.__name__,