From ba7c6519757ab3fa7e322420c66513d448e051c3 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 11 Jan 2022 07:30:34 -0800
Subject: [PATCH 1/5] feat: shorter timeout on test_connection

---
 setup.py                                       |  1 +
 .../database/DatabaseModal/SqlAlchemyForm.tsx  |  3 +++
 .../CRUD/data/database/DatabaseModal/index.tsx | 15 ++++++++++++++-
 superset/config.py                             |  5 +++++
 superset/databases/commands/test_connection.py | 18 +++++++++++++++++-
 5 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 3217b23fa5a2..db127790ef96 100644
--- a/setup.py
+++ b/setup.py
@@ -81,6 +81,7 @@ def get_git_sha() -> str:
         "flask-talisman",
         "flask-migrate",
         "flask-wtf",
+        "func_timeout",
         "geopy",
         "graphlib-backport",
         "gunicorn>=20.1.0",
diff --git a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/SqlAlchemyForm.tsx b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/SqlAlchemyForm.tsx
index d9221499295c..a48cea33e413 100644
--- a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/SqlAlchemyForm.tsx
+++ b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/SqlAlchemyForm.tsx
@@ -29,12 +29,14 @@ const SqlAlchemyTab = ({
   testConnection,
   conf,
   isEditMode = false,
+  testInProgress = false,
 }: {
   db: DatabaseObject | null;
   onInputChange: EventHandler<ChangeEvent<HTMLInputElement>>;
   testConnection: EventHandler<MouseEvent<HTMLElement>>;
   conf: { SQLALCHEMY_DOCS_URL: string; SQLALCHEMY_DISPLAY_TEXT: string };
   isEditMode?: boolean;
+  testInProgress?: boolean;
 }) => (
   <>
     <StyledInputContainer>
@@ -88,6 +90,7 @@ const SqlAlchemyTab = ({
     </StyledInputContainer>
     <Button
       onClick={testConnection}
+      disabled={testInProgress}
       cta
       buttonStyle="link"
       css={(theme: SupersetTheme) => wideButton(theme)}
diff --git a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.tsx b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.tsx
index 86b221808a40..ae4cff8bb5b3 100644
--- a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.tsx
+++ b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/index.tsx
@@ -440,6 +440,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
   const [dbName, setDbName] = useState('');
   const [editNewDb, setEditNewDb] = useState<boolean>(false);
   const [isLoading, setLoading] = useState<boolean>(false);
+  const [testInProgress, setTestInProgress] = useState<boolean>(false);
   const conf = useCommonConf();
   const dbImages = getDatabaseImages();
   const connectionAlert = getConnectionAlert();
@@ -494,7 +495,18 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
       encrypted_extra: db?.encrypted_extra || '',
       server_cert: db?.server_cert || undefined,
     };
-    testDatabaseConnection(connection, addDangerToast, addSuccessToast);
+    setTestInProgress(true);
+    testDatabaseConnection(
+      connection,
+      (errorMsg: string) => {
+        setTestInProgress(false);
+        addDangerToast(errorMsg);
+      },
+      (errorMsg: string) => {
+        setTestInProgress(false);
+        addSuccessToast(errorMsg);
+      },
+    );
   };
 
   const onClose = () => {
@@ -1047,6 +1059,7 @@ const DatabaseModal: FunctionComponent<DatabaseModalProps> = ({
                 conf={conf}
                 testConnection={testConnection}
                 isEditMode={isEditMode}
+                testInProgress={testInProgress}
               />
               {isDynamic(db?.backend || db?.engine) && !isEditMode && (
                 <div css={(theme: SupersetTheme) => infoTooltip(theme)}>
diff --git a/superset/config.py b/superset/config.py
index 8bee76bee0ea..d97fbc52b50f 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1164,6 +1164,11 @@ def SQL_QUERY_MUTATOR(  # pylint: disable=invalid-name,unused-argument
     "SQLite",
     # etc.
 ]
+# When adding a new database we try to connect to it. Depending on which parameters are
+# incorrect this could take a couple minutes, until the SQLAlchemy driver pinging the
+# database times out. Instead of relying on the driver timeout we can specify a shorter
+# one here.
+TEST_DATABASE_CONNECTION_TIMEOUT = timedelta(seconds=30)
 
 # Do you want Talisman enabled?
 TALISMAN_ENABLED = False
diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py
index 2d44bed940a0..fb490337f814 100644
--- a/superset/databases/commands/test_connection.py
+++ b/superset/databases/commands/test_connection.py
@@ -18,8 +18,10 @@
 from contextlib import closing
 from typing import Any, Dict, Optional
 
+from flask import current_app as app
 from flask_appbuilder.security.sqla.models import User
 from flask_babel import gettext as _
+from func_timeout import func_timeout, FunctionTimedOut
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.exc import DBAPIError, NoSuchModuleError
 
@@ -78,7 +80,21 @@ def run(self) -> None:
             )
             with closing(engine.raw_connection()) as conn:
                 try:
-                    alive = engine.dialect.do_ping(conn)
+                    alive = func_timeout(
+                        int(
+                            app.config[
+                                "TEST_DATABASE_CONNECTION_TIMEOUT"
+                            ].total_seconds()
+                        ),
+                        engine.dialect.do_ping,
+                        args=(conn,),
+                    )
+                except FunctionTimedOut as ex:
+                    raise Exception(
+                        "Please check your connection details and database settings, "
+                        "and ensure that your database is accepting connections, then "
+                        "try connecting again."
+                    ) from ex
                 except Exception:  # pylint: disable=broad-except
                     alive = False
                 if not alive:

From 4577933adb1f0d54b5d90de281180330e8fd1287 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 11 Jan 2022 09:27:10 -0800
Subject: [PATCH 2/5] pip-compile-multi --no-upgrade

---
 requirements/base.txt | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/requirements/base.txt b/requirements/base.txt
index ced4edfb2a0b..06e3d0a76fb0 100644
--- a/requirements/base.txt
+++ b/requirements/base.txt
@@ -103,6 +103,8 @@ flask-wtf==0.14.3
     # via
     #   apache-superset
     #   flask-appbuilder
+func-timeout==4.3.5
+    # via apache-superset
 geographiclib==1.52
     # via geopy
 geopy==2.2.0

From b370cb2996061ee8587528dd1bf895bbfc4b9867 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 11 Jan 2022 10:15:25 -0800
Subject: [PATCH 3/5] Fix for SQLite

---
 superset/databases/commands/test_connection.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py
index fb490337f814..54f48c41018f 100644
--- a/superset/databases/commands/test_connection.py
+++ b/superset/databases/commands/test_connection.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+import sqlite3
 from contextlib import closing
 from typing import Any, Dict, Optional
 
@@ -89,6 +90,9 @@ def run(self) -> None:
                         engine.dialect.do_ping,
                         args=(conn,),
                     )
+                except sqlite3.ProgrammingError:
+                    # SQLite can't run on a separate thread, so ``func_timeout`` fails
+                    alive = engine.dialect.do_ping(conn)
                 except FunctionTimedOut as ex:
                     raise Exception(
                         "Please check your connection details and database settings, "

From b603df8315fc3fd217ddb34182ff712b3b280818 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 11 Jan 2022 15:50:21 -0800
Subject: [PATCH 4/5] Return 408

---
 .../databases/commands/test_connection.py     | 23 +++++++++++++++----
 superset/errors.py                            |  2 ++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/superset/databases/commands/test_connection.py b/superset/databases/commands/test_connection.py
index 54f48c41018f..d3076017bc40 100644
--- a/superset/databases/commands/test_connection.py
+++ b/superset/databases/commands/test_connection.py
@@ -34,7 +34,8 @@
     DatabaseTestConnectionUnexpectedError,
 )
 from superset.databases.dao import DatabaseDAO
-from superset.exceptions import SupersetSecurityException
+from superset.errors import ErrorLevel, SupersetErrorType
+from superset.exceptions import SupersetSecurityException, SupersetTimeoutException
 from superset.extensions import event_logger
 from superset.models.core import Database
 
@@ -94,10 +95,15 @@ def run(self) -> None:
                     # SQLite can't run on a separate thread, so ``func_timeout`` fails
                     alive = engine.dialect.do_ping(conn)
                 except FunctionTimedOut as ex:
-                    raise Exception(
-                        "Please check your connection details and database settings, "
-                        "and ensure that your database is accepting connections, then "
-                        "try connecting again."
+                    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
@@ -134,6 +140,13 @@ def run(self) -> None:
                 engine=database.db_engine_spec.__name__,
             )
             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__,
+            )
+            # bubble up the exception to return a 408
+            raise ex
         except Exception as ex:
             event_logger.log_with_context(
                 action=f"test_connection_error.{ex.__class__.__name__}",
diff --git a/superset/errors.py b/superset/errors.py
index 8dc4f90308f7..9b3414ecb57e 100644
--- a/superset/errors.py
+++ b/superset/errors.py
@@ -51,6 +51,7 @@ class SupersetErrorType(str, Enum):
     CONNECTION_MISSING_PARAMETERS_ERROR = "CONNECTION_MISSING_PARAMETERS_ERROR"
     OBJECT_DOES_NOT_EXIST_ERROR = "OBJECT_DOES_NOT_EXIST_ERROR"
     SYNTAX_ERROR = "SYNTAX_ERROR"
+    CONNECTION_DATABASE_TIMEOUT = "CONNECTION_DATABASE_TIMEOUT"
 
     # Viz errors
     VIZ_GET_DF_ERROR = "VIZ_GET_DF_ERROR"
@@ -172,6 +173,7 @@ class SupersetErrorType(str, Enum):
     SupersetErrorType.CONNECTION_INVALID_PORT_ERROR: [1034],
     SupersetErrorType.ASYNC_WORKERS_ERROR: [1035],
     SupersetErrorType.DATABASE_NOT_FOUND_ERROR: [1011, 1036],
+    SupersetErrorType.CONNECTION_DATABASE_TIMEOUT: [1001, 1009],
 }
 
 

From 2e43a865c8b64d9c379941e2148176a9fb4adfd9 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 11 Jan 2022 17:50:21 -0800
Subject: [PATCH 5/5] Add test

---
 .../databases/commands_tests.py               | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/tests/integration_tests/databases/commands_tests.py b/tests/integration_tests/databases/commands_tests.py
index 5860b23fd91b..3d682f700cc4 100644
--- a/tests/integration_tests/databases/commands_tests.py
+++ b/tests/integration_tests/databases/commands_tests.py
@@ -19,6 +19,7 @@
 
 import pytest
 import yaml
+from func_timeout import FunctionTimedOut
 from sqlalchemy.exc import DBAPIError
 
 from superset import db, event_logger, security_manager
@@ -38,7 +39,11 @@
 from superset.databases.commands.validate import ValidateDatabaseParametersCommand
 from superset.databases.schemas import DatabaseTestConnectionSchema
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import SupersetErrorsException, SupersetSecurityException
+from superset.exceptions import (
+    SupersetErrorsException,
+    SupersetSecurityException,
+    SupersetTimeoutException,
+)
 from superset.models.core import Database
 from superset.utils.core import backend, get_example_database
 from tests.integration_tests.base_tests import SupersetTestCase
@@ -612,6 +617,28 @@ def test_connection_do_ping_exception(
             == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
         )
 
+    @mock.patch("superset.databases.commands.test_connection.func_timeout")
+    @mock.patch(
+        "superset.databases.commands.test_connection.event_logger.log_with_context"
+    )
+    def test_connection_do_ping_timeout(self, mock_event_logger, mock_func_timeout):
+        """Test to make sure do_ping exceptions gets captured"""
+        database = get_example_database()
+        mock_func_timeout.side_effect = FunctionTimedOut("Time out")
+        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(SupersetTimeoutException) as excinfo:
+            command_without_db_name.run()
+        assert excinfo.value.status == 408
+        assert (
+            excinfo.value.error.error_type
+            == SupersetErrorType.CONNECTION_DATABASE_TIMEOUT
+        )
+
     @mock.patch("superset.databases.dao.Database.get_sqla_engine")
     @mock.patch(
         "superset.databases.commands.test_connection.event_logger.log_with_context"