From eb05225f0daca50c88ca8e226b9b9198d5859a7f Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Tue, 13 Jun 2023 12:29:40 -0400 Subject: [PATCH] feat(sshtunnel): add configuration for SSH_TIMEOUT (#24369) --- superset/config.py | 4 ++++ superset/extensions/ssh.py | 1 + superset/models/core.py | 6 ++++-- tests/unit_tests/extensions/ssh_test.py | 2 ++ 4 files changed, 11 insertions(+), 2 deletions(-) diff --git a/superset/config.py b/superset/config.py index f34d99e3cd..53d399110d 100644 --- a/superset/config.py +++ b/superset/config.py @@ -498,7 +498,11 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = { # ---------------------------------------------------------------------- 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 index 78b0c4116b..5cf84099f2 100644 --- a/superset/extensions/ssh.py +++ b/superset/extensions/ssh.py @@ -35,6 +35,7 @@ class SSHManager: 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 index c18d12049e..3a52367242 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -31,6 +31,7 @@ from typing import Any, Callable, Optional, TYPE_CHECKING 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 @@ class Database( 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 index 4538d71969..13bf905e6f 100644 --- 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