From 71e21c9be0e7e84dac227110897a1a4a037eb680 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 16:50:43 +0000 Subject: [PATCH 1/6] allow users to set TUNNEL_TIMEOUT from config --- superset/config.py | 3 ++- superset/extensions/ssh.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/superset/config.py b/superset/config.py index d522e10ac1ec3..065866c928cae 100644 --- a/superset/config.py +++ b/superset/config.py @@ -494,7 +494,7 @@ class D3Format(TypedDict, total=False): # Allow users to enable ssh tunneling when creating a DB. # Users must check whether the DB engine supports SSH Tunnels # otherwise enabling this flag won't have any effect on the DB. - "SSH_TUNNELING": False, + "SSH_TUNNELING": True, "AVOID_COLORS_COLLISION": True, } @@ -515,6 +515,7 @@ class D3Format(TypedDict, total=False): # ---------------------------------------------------------------------- SSH_TUNNEL_MANAGER_CLASS = "superset.extensions.ssh.SSHManager" SSH_TUNNEL_LOCAL_BIND_ADDRESS = "127.0.0.1" +SSH_TUNNEL_TIMEOUT = 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 6a852ea7cd692..2d98310c5d6cc 100644 --- a/superset/extensions/ssh.py +++ b/superset/extensions/ssh.py @@ -22,7 +22,7 @@ from flask import Flask from paramiko import RSAKey -from sshtunnel import open_tunnel, SSHTunnelForwarder +from sshtunnel import open_tunnel, SSHTunnelForwarder, TUNNEL_TIMEOUT from superset.databases.utils import make_url_safe @@ -34,6 +34,7 @@ class SSHManager: def __init__(self, app: Flask) -> None: super().__init__() self.local_bind_address = app.config["SSH_TUNNEL_LOCAL_BIND_ADDRESS"] + TUNNEL_TIMEOUT = app.config["SSH_TUNNEL_TIMEOUT"] def build_sqla_url( # pylint: disable=no-self-use self, sqlalchemy_url: str, server: SSHTunnelForwarder From c823939538650e921d07953c1a93525bc161483c Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 12:56:36 -0400 Subject: [PATCH 2/6] Update config.py --- superset/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/config.py b/superset/config.py index 065866c928cae..c15edd5ab0061 100644 --- a/superset/config.py +++ b/superset/config.py @@ -494,7 +494,7 @@ class D3Format(TypedDict, total=False): # Allow users to enable ssh tunneling when creating a DB. # Users must check whether the DB engine supports SSH Tunnels # otherwise enabling this flag won't have any effect on the DB. - "SSH_TUNNELING": True, + "SSH_TUNNELING": False, "AVOID_COLORS_COLLISION": True, } From 3fdd840518a370a1133ae9c716eeddfda25dc94e Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 17:33:23 +0000 Subject: [PATCH 3/6] refactor --- superset/config.py | 2 +- superset/extensions/ssh.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/superset/config.py b/superset/config.py index c15edd5ab0061..a424a09d23d8c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -515,7 +515,7 @@ class D3Format(TypedDict, total=False): # ---------------------------------------------------------------------- SSH_TUNNEL_MANAGER_CLASS = "superset.extensions.ssh.SSHManager" SSH_TUNNEL_LOCAL_BIND_ADDRESS = "127.0.0.1" -SSH_TUNNEL_TIMEOUT = 1.0 +SSH_TUNNEL_TIMEOUT_SEC = 10.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 2d98310c5d6cc..78b0c4116b192 100644 --- a/superset/extensions/ssh.py +++ b/superset/extensions/ssh.py @@ -20,9 +20,9 @@ from io import StringIO from typing import TYPE_CHECKING +import sshtunnel from flask import Flask from paramiko import RSAKey -from sshtunnel import open_tunnel, SSHTunnelForwarder, TUNNEL_TIMEOUT from superset.databases.utils import make_url_safe @@ -34,10 +34,10 @@ class SSHManager: def __init__(self, app: Flask) -> None: super().__init__() self.local_bind_address = app.config["SSH_TUNNEL_LOCAL_BIND_ADDRESS"] - TUNNEL_TIMEOUT = app.config["SSH_TUNNEL_TIMEOUT"] + sshtunnel.TUNNEL_TIMEOUT = app.config["SSH_TUNNEL_TIMEOUT_SEC"] def build_sqla_url( # pylint: disable=no-self-use - self, sqlalchemy_url: str, server: SSHTunnelForwarder + self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder ) -> str: # override any ssh tunnel configuration object url = make_url_safe(sqlalchemy_url) @@ -50,7 +50,7 @@ def create_tunnel( self, ssh_tunnel: "SSHTunnel", sqlalchemy_database_uri: str, - ) -> SSHTunnelForwarder: + ) -> sshtunnel.SSHTunnelForwarder: url = make_url_safe(sqlalchemy_database_uri) params = { "ssh_address_or_host": (ssh_tunnel.server_address, ssh_tunnel.server_port), @@ -69,7 +69,7 @@ def create_tunnel( ) params["ssh_pkey"] = private_key - return open_tunnel(**params) + return sshtunnel.open_tunnel(**params) class SSHManagerFactory: From 538d6903c6008e40fead87ab47e0e9961ad59c87 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 18:42:45 +0000 Subject: [PATCH 4/6] added test --- tests/unit_tests/extensions/ssh_test.py | 33 +++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 tests/unit_tests/extensions/ssh_test.py diff --git a/tests/unit_tests/extensions/ssh_test.py b/tests/unit_tests/extensions/ssh_test.py new file mode 100644 index 0000000000000..4eddb90694ae3 --- /dev/null +++ b/tests/unit_tests/extensions/ssh_test.py @@ -0,0 +1,33 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from typing import Any +from unittest.mock import Mock, patch + +import pytest +import sshtunnel + +from superset.extensions.ssh import SSHManagerFactory + + +def test_ssh_tunnel_timeout_setting() -> None: + app = Mock() + app.config = { + "SSH_TUNNEL_TIMEOUT_SEC": 123.0, + } + factory = SSHManagerFactory() + factory.init_app(app) + assert sshtunnel.TUNNEL_TIMEOUT == 123.0 From af94f53aa4334315ac92b083d498f206b9d1076d Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 19:07:22 +0000 Subject: [PATCH 5/6] ok --- tests/unit_tests/extensions/ssh_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/extensions/ssh_test.py b/tests/unit_tests/extensions/ssh_test.py index 4eddb90694ae3..65f8ccacd41fd 100644 --- a/tests/unit_tests/extensions/ssh_test.py +++ b/tests/unit_tests/extensions/ssh_test.py @@ -26,7 +26,10 @@ def test_ssh_tunnel_timeout_setting() -> None: app = Mock() app.config = { - "SSH_TUNNEL_TIMEOUT_SEC": 123.0, + "SSH_TUNNEL_MAX_RETRIES": 2, + "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "test", + "SSH_TUNNEL_TIMEOUT_SEC": 10.0, + "SSH_TUNNEL_MANAGER_CLASS": "superset.extensions.ssh.SSHManager", } factory = SSHManagerFactory() factory.init_app(app) From d6008fb9b35a9c467c2dfbae55a30781497a17ac Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Wed, 24 May 2023 19:57:09 +0000 Subject: [PATCH 6/6] ok --- tests/unit_tests/extensions/ssh_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/extensions/ssh_test.py b/tests/unit_tests/extensions/ssh_test.py index 65f8ccacd41fd..0e997729d96fe 100644 --- a/tests/unit_tests/extensions/ssh_test.py +++ b/tests/unit_tests/extensions/ssh_test.py @@ -28,7 +28,7 @@ def test_ssh_tunnel_timeout_setting() -> None: app.config = { "SSH_TUNNEL_MAX_RETRIES": 2, "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "test", - "SSH_TUNNEL_TIMEOUT_SEC": 10.0, + "SSH_TUNNEL_TIMEOUT_SEC": 123.0, "SSH_TUNNEL_MANAGER_CLASS": "superset.extensions.ssh.SSHManager", } factory = SSHManagerFactory()