feat(sshtunnel): add configuration for SSH_TIMEOUT (apache#24369)

weareferal · Jun 13, 2023 · eb05225 · eb05225
1 parent 1328c56
commit eb05225
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 deletions.
diff --git a/superset/config.py b/superset/config.py
@@ -498,7 +498,11 @@ class D3Format(TypedDict, total=False):
 # ----------------------------------------------------------------------
 SSH_TUNNEL_MANAGER_CLASS = "superset.extensions.ssh.SSHManager"
 SSH_TUNNEL_LOCAL_BIND_ADDRESS = "127.0.0.1"
+#: Timeout (seconds) for tunnel connection (open_channel timeout)
 SSH_TUNNEL_TIMEOUT_SEC = 10.0
+#: Timeout (seconds) for transport socket (``socket.settimeout``)
+SSH_TUNNEL_PACKET_TIMEOUT_SEC = 1.0
+
 
 # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars.
 DEFAULT_FEATURE_FLAGS.update(

diff --git a/superset/extensions/ssh.py b/superset/extensions/ssh.py
@@ -35,6 +35,7 @@ def __init__(self, app: Flask) -> None:
         super().__init__()
         self.local_bind_address = app.config["SSH_TUNNEL_LOCAL_BIND_ADDRESS"]
         sshtunnel.TUNNEL_TIMEOUT = app.config["SSH_TUNNEL_TIMEOUT_SEC"]
+        sshtunnel.SSH_TIMEOUT = app.config["SSH_TUNNEL_PACKET_TIMEOUT_SEC"]
 
     def build_sqla_url(  # pylint: disable=no-self-use
         self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder

diff --git a/superset/models/core.py b/superset/models/core.py
@@ -31,6 +31,7 @@
 import numpy
 import pandas as pd
 import sqlalchemy as sqla
+import sshtunnel
 from flask import g, request
 from flask_appbuilder import Model
 from sqlalchemy import (
@@ -406,9 +407,10 @@ def get_sqla_engine_with_context(
         with engine_context as server_context:
             if ssh_tunnel and server_context:
                 logger.info(
-                    "[SSH] Successfully create tunnel at %s: %s",
+                    "[SSH] Successfully created tunnel w/ %s tunnel_timeout + %s ssh_timeout at %s",
+                    sshtunnel.TUNNEL_TIMEOUT,
+                    sshtunnel.SSH_TIMEOUT,
                     server_context.local_bind_address,
-                    server_context.local_bind_port,
                 )
                 sqlalchemy_uri = ssh_manager_factory.instance.build_sqla_url(
                     sqlalchemy_uri, server_context

diff --git a/tests/unit_tests/extensions/ssh_test.py b/tests/unit_tests/extensions/ssh_test.py
@@ -28,8 +28,10 @@ def test_ssh_tunnel_timeout_setting() -> None:
         "SSH_TUNNEL_MAX_RETRIES": 2,
         "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "test",
         "SSH_TUNNEL_TIMEOUT_SEC": 123.0,
+        "SSH_TUNNEL_PACKET_TIMEOUT_SEC": 321.0,
         "SSH_TUNNEL_MANAGER_CLASS": "superset.extensions.ssh.SSHManager",
     }
     factory = SSHManagerFactory()
     factory.init_app(app)
     assert sshtunnel.TUNNEL_TIMEOUT == 123.0
+    assert sshtunnel.SSH_TIMEOUT == 321.0