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 <[email protected]> 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()
