feat: shorter timeout on test_connection (apache#18001)

* feat: shorter timeout on test_connection

* pip-compile-multi --no-upgrade

* Fix for SQLite

* Return 408

* Add test

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

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

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)}

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)}>

superset/config.py

@@ -1171,6 +1171,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

superset/databases/commands/test_connection.py

@@ -15,11 +15,14 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+import sqlite3
 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
 
@@ -31,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
 
@@ -78,7 +82,29 @@ def run(self) -> None:
             )
             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:
+                    # SQLite can't run on a separate thread, so ``func_timeout`` fails
                     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:
@@ -114,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__}",

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],
 }
 
 

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"