chore: remove dependency on func_timeout because LGPL (apache#31519)

preset-io · Dec 19, 2024 · d8aba2f · d8aba2f
1 parent db11a2a
commit d8aba2f
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 19 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -53,7 +53,6 @@ dependencies = [
     "flask-migrate>=3.1.0, <4.0",
     "flask-session>=0.4.0, <1.0",
     "flask-wtf>=1.1.0, <2.0",
-    "func_timeout",
     "geopy",
     "greenlet>=3.0.3, <=3.1.1",
     "gunicorn>=22.0.0; sys_platform != 'win32'",
@@ -391,6 +390,5 @@ python-geohash = "0"
 # --------------------------------------------------------------
 
 # TODO REMOVE THESE DEPS FROM CODEBASE
-func-timeout = "4"  # AGPL
 paramiko = "3"  # GPL
 pyxlsb = "1"  # GPL
diff --git a/requirements/base.txt b/requirements/base.txt
@@ -138,8 +138,6 @@ flask-wtf==1.2.2
     # via
     #   apache-superset (pyproject.toml)
     #   flask-appbuilder
-func-timeout==4.3.5
-    # via apache-superset (pyproject.toml)
 geographiclib==2.0
     # via geopy
 geopy==2.4.1

diff --git a/requirements/development.txt b/requirements/development.txt
@@ -255,10 +255,6 @@ fonttools==4.55.0
     # via matplotlib
 freezegun==1.5.1
     # via apache-superset
-func-timeout==4.3.5
-    # via
-    #   -c requirements/base.txt
-    #   apache-superset
 future==1.0.0
     # via pyhive
 geographiclib==2.0

diff --git a/superset/commands/database/test_connection.py b/superset/commands/database/test_connection.py
@@ -21,7 +21,6 @@
 
 from flask import current_app as app
 from flask_babel import gettext as _
-from func_timeout import func_timeout, FunctionTimedOut
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import DBAPIError, NoSuchModuleError
 
@@ -48,6 +47,7 @@
 )
 from superset.extensions import event_logger
 from superset.models.core import Database
+from superset.utils import core as utils
 from superset.utils.ssh_tunnel import unmask_password_info
 
 logger = logging.getLogger(__name__)
@@ -142,16 +142,14 @@ def ping(engine: Engine) -> bool:
 
             with database.get_sqla_engine(override_ssh_tunnel=ssh_tunnel) as engine:
                 try:
-                    alive = func_timeout(
-                        app.config["TEST_DATABASE_CONNECTION_TIMEOUT"].total_seconds(),
-                        ping,
-                        args=(engine,),
-                    )
+                    time_delta = app.config["TEST_DATABASE_CONNECTION_TIMEOUT"]
+                    with utils.timeout(int(time_delta.total_seconds())):
+                        alive = ping(engine)
                 except (sqlite3.ProgrammingError, RuntimeError):
-                    # SQLite can't run on a separate thread, so ``func_timeout`` fails
+                    # SQLite can't run on a separate thread, so ``utils.timeout`` fails
                     # RuntimeError catches the equivalent error from duckdb.
                     alive = engine.dialect.do_ping(engine)
-                except FunctionTimedOut as ex:
+                except SupersetTimeoutException as ex:
                     raise SupersetTimeoutException(
                         error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
                         message=(

diff --git a/tests/integration_tests/databases/commands_tests.py b/tests/integration_tests/databases/commands_tests.py
@@ -19,7 +19,6 @@
 
 import pytest
 import yaml
-from func_timeout import FunctionTimedOut
 from sqlalchemy.exc import DBAPIError
 
 from superset import db, event_logger, security_manager  # noqa: F401
@@ -932,7 +931,7 @@ def test_connection_do_ping_exception(
             == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
         )
 
-    @patch("superset.commands.database.test_connection.func_timeout")
+    @patch("superset.utils.core.timeout")
     @patch("superset.commands.database.test_connection.event_logger.log_with_context")
     @patch("superset.utils.core.g")
     def test_connection_do_ping_timeout(
@@ -941,7 +940,11 @@ def test_connection_do_ping_timeout(
         """Test to make sure do_ping exceptions gets captured"""
         database = get_example_database()
         mock_g.user = security_manager.find_user("admin")
-        mock_func_timeout.side_effect = FunctionTimedOut("Time out")
+        mock_func_timeout.side_effect = SupersetTimeoutException(
+            error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR,
+            message="ERROR",
+            level=ErrorLevel.ERROR,
+        )
         db_uri = database.sqlalchemy_uri_decrypted
         json_payload = {"sqlalchemy_uri": db_uri}
         command_without_db_name = TestConnectionDatabaseCommand(json_payload)