[superset] 01/01: move to extensions

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hugh pushed a commit to branch sshmanager-2
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 295b1486e95b7cda30e76e931f3e9cbd83bc70fa
Author: hughhhh <hu...@gmail.com>
AuthorDate: Thu Dec 15 15:52:48 2022 -0500

    move to extensions
---
 superset/config.py              | 16 +---------------
 superset/extensions/__init__.py | 28 ++++++++++++++++++++++++++++
 superset/models/core.py         | 10 +++++++---
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index 71b2203339..de3770a191 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -57,7 +57,6 @@ from superset.advanced_data_type.plugins.internet_address import internet_addres
 from superset.advanced_data_type.plugins.internet_port import internet_port
 from superset.advanced_data_type.types import AdvancedDataType
 from superset.constants import CHANGE_ME_SECRET_KEY
-from superset.databases.utils import make_url_safe
 from superset.jinja_context import BaseTemplateProcessor
 from superset.reports.types import ReportScheduleExecutor
 from superset.stats_logger import DummyStatsLogger
@@ -494,20 +493,7 @@ DEFAULT_FEATURE_FLAGS: Dict[str, bool] = {
 #                          FIREWALL (only port 22 is open)
 
 # ----------------------------------------------------------------------
-class SSHManager:  # pylint: disable=too-few-public-methods
-    local_bind_address = "127.0.0.1"
-
-    @classmethod
-    def mutator(cls, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder) -> str:
-        # override any ssh tunnel configuration object
-        url = make_url_safe(sqlalchemy_url)
-        return url.set(
-            host=cls.local_bind_address,
-            port=server.local_bind_port,
-        )
-
-
-SSH_TUNNEL_MANAGER = SSHManager  # pylint: disable=invalid-name
+SSH_TUNNEL_MANAGER = None
 
 # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars.
 DEFAULT_FEATURE_FLAGS.update(
diff --git a/superset/extensions/__init__.py b/superset/extensions/__init__.py
index 1f5882f749..b15bc32b1f 100644
--- a/superset/extensions/__init__.py
+++ b/superset/extensions/__init__.py
@@ -20,6 +20,7 @@ from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
 import celery
+import sshtunnel
 from cachelib.base import BaseCache
 from flask import Flask
 from flask_appbuilder import AppBuilder, SQLA
@@ -28,6 +29,7 @@ from flask_talisman import Talisman
 from flask_wtf.csrf import CSRFProtect
 from werkzeug.local import LocalProxy
 
+from superset.databases.utils import make_url_safe
 from superset.utils.async_query_manager import AsyncQueryManager
 from superset.utils.cache_manager import CacheManager
 from superset.utils.encrypt import EncryptedFieldFactory
@@ -109,6 +111,31 @@ class ProfilingExtension:  # pylint: disable=too-few-public-methods
         app.wsgi_app = SupersetProfiler(app.wsgi_app, self.interval)  # type: ignore
 
 
+class SSHManager:  # pylint: disable=too-few-public-methods
+    local_bind_address = "127.0.0.1"
+
+    @classmethod
+    def mutator(cls, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder) -> str:
+        # override any ssh tunnel configuration object
+        url = make_url_safe(sqlalchemy_url)
+        return url.set(
+            host=cls.local_bind_address,
+            port=server.local_bind_port,
+        )
+
+    @classmethod
+    def create_tunnel(
+        cls,
+        ssh_address_or_host: str,
+        ssh_port: int,
+        ssh_username: str,
+        remote_bind_address: str,
+        local_bind_address: str,
+        **kwargs: Dict[str, Any],
+    ) -> sshtunnel.SSHTunnelForwarder:
+        return sshtunnel.open_tunnel(**kwargs)
+
+
 APP_DIR = os.path.join(os.path.dirname(__file__), os.path.pardir)
 appbuilder = AppBuilder(update_perms=False)
 async_query_manager = AsyncQueryManager()
@@ -127,3 +154,4 @@ profiling = ProfilingExtension()
 results_backend_manager = ResultsBackendManager()
 security_manager = LocalProxy(lambda: appbuilder.sm)
 talisman = Talisman()
+ssh_manager = SSHManager()
diff --git a/superset/models/core.py b/superset/models/core.py
index ee0bc2c465..58b95f84a4 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -58,7 +58,12 @@ from superset import app, db_engine_specs
 from superset.constants import PASSWORD_MASK
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import MetricType, TimeGrain
-from superset.extensions import cache_manager, encrypted_field_factory, security_manager
+from superset.extensions import (
+    cache_manager,
+    encrypted_field_factory,
+    security_manager,
+    ssh_manager,
+)
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
 from superset.result_set import SupersetResultSet
 from superset.utils import cache as cache_util, core as utils
@@ -66,7 +71,6 @@ from superset.utils.core import get_username
 from superset.utils.memoized import memoized
 
 config = app.config
-ssh_manager = config["SSH_TUNNEL_MANAGER"]
 custom_password_store = config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"]
 stats_logger = config["STATS_LOGGER"]
 log_query = config["QUERY_LOGGER"]
@@ -388,7 +392,7 @@ class Database(
             ssh_params = ssh_tunnel.kwarg_parameters(
                 bind_host=url.host, bind_port=url.port
             )
-            engine_context = sshtunnel.open_tunnel(**ssh_params)
+            engine_context = ssh_manager.create_tunnel(**ssh_params)
         else:
             engine_context = nullcontext()