From 0f6038f71de05d5185c59c9060249316a13ee5b1 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Tue, 22 Nov 2022 17:02:08 -0500 Subject: [PATCH 01/16] write class in config --- superset/config.py | 38 +++++++++++++++++++ .../databases/ssh_tunnel/commands/create.py | 6 +++ .../databases/ssh_tunnel/commands/update.py | 8 ++++ 3 files changed, 52 insertions(+) diff --git a/superset/config.py b/superset/config.py index f163997c6ee4b..23afbd1c6a7ec 100644 --- a/superset/config.py +++ b/superset/config.py @@ -471,8 +471,46 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: "DRILL_TO_DETAIL": False, "DATAPANEL_CLOSED_BY_DEFAULT": False, "HORIZONTAL_FILTER_BAR": 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 Tunnel +# ------------------------------ +# Allow users to set the host used when connecting to the SSH Tunnel +# as localhost and any other alias (0.0.0.0) +# ---------------------------------------------------------------------- +# | +# -------------+ | +----------+ +# LOCAL | | | REMOTE | :22 SSH +# CLIENT | <== SSH ========> | SERVER | :8080 web service +# -------------+ | +----------+ +# | +# FIREWALL (only port 22 is open) + +# ---------------------------------------------------------------------- +# class SSHManager: +# def validate(self): +# # validation on CREATE + UPDATE on SSHTunnel Model +# raise NotImplemented() + +# def mutator(self): +# # override any ssh tunnel configuration object +# raise NotImplemented() + +# @property +# def local_bind_address(self): +# # set the local binding address for the local client +# # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder +# # `server` return value +# return "127.0.0.1" + + +SSH_TUNNEL_MANAGER = None + # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars. DEFAULT_FEATURE_FLAGS.update( { diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 1185d20bdf7cb..8ae88a9f2d656 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -20,12 +20,16 @@ from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError +from superset import app from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOCreateFailedError from superset.databases.dao import DatabaseDAO from superset.databases.ssh_tunnel.commands.exceptions import SSHTunnelCreateFailedError from superset.databases.ssh_tunnel.dao import SSHTunnelDAO +config = app.config +ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] + logger = logging.getLogger(__name__) @@ -47,4 +51,6 @@ def run(self) -> Model: def validate(self) -> None: # TODO(hughhh): check to make sure the server port is not localhost # using the config.SSH_TUNNEL_MANAGER + if ssh_tunnel_manager: + ssh_tunnel_manager.validate() return diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index a9d9a6fe6d183..146cc344964f0 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -20,6 +20,7 @@ from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError +from superset import app from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOUpdateFailedError from superset.databases.ssh_tunnel.commands.exceptions import ( @@ -33,6 +34,8 @@ logger = logging.getLogger(__name__) +config = app.config +ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] class UpdateSSHTunnelCommand(BaseCommand): def __init__(self, model_id: int, data: Dict[str, Any]): @@ -53,3 +56,8 @@ def validate(self) -> None: self._model = SSHTunnelDAO.find_by_id(self._model_id) if not self._model: raise SSHTunnelNotFoundError() + + # TODO(hughhh): check to make sure the server port is not localhost + # using the config.SSH_TUNNEL_MANAGER + if ssh_tunnel_manager: + ssh_tunnel_manager.validate() From 2502f635f1446d0c1c7f5227c914b75ebfbfdcca Mon Sep 17 00:00:00 2001 From: hughhhh Date: Tue, 22 Nov 2022 17:15:04 -0500 Subject: [PATCH 02/16] add mutate function --- superset/models/core.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/superset/models/core.py b/superset/models/core.py index fd6ad8be9347e..25adc61942196 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -66,6 +66,8 @@ from superset.utils.memoized import memoized config = app.config + +ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] custom_password_store = config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"] stats_logger = config["STATS_LOGGER"] log_query = config["QUERY_LOGGER"] @@ -381,6 +383,8 @@ def get_sqla_engine_with_context( ssh_tunnel.bind_host = url.host ssh_tunnel.bind_port = url.port ssh_params = ssh_tunnel.parameters() + if ssh_tunnel_manager: + ssh_params = ssh_tunnel_manager.mutate(ssh_params) try: with sshtunnel.open_tunnel(**ssh_params) as server: yield self._get_sqla_engine( From 25226604d1a2dac94dcd9412353bb38e6df24e19 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Tue, 22 Nov 2022 17:25:14 -0500 Subject: [PATCH 03/16] refactor exceptions and add mutator --- superset/models/core.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/superset/models/core.py b/superset/models/core.py index 25adc61942196..37c96feffadd4 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -385,25 +385,18 @@ def get_sqla_engine_with_context( ssh_params = ssh_tunnel.parameters() if ssh_tunnel_manager: ssh_params = ssh_tunnel_manager.mutate(ssh_params) - try: - with sshtunnel.open_tunnel(**ssh_params) as server: - yield self._get_sqla_engine( - schema=schema, - nullpool=nullpool, - source=source, - ssh_tunnel_server=server, - ) - except Exception as ex: - raise ex - - else: - # do look up in table for using database_id - try: + with sshtunnel.open_tunnel(**ssh_params) as server: yield self._get_sqla_engine( - schema=schema, nullpool=nullpool, source=source + schema=schema, + nullpool=nullpool, + source=source, + ssh_tunnel_server=server, ) - except Exception as ex: - raise ex + else: + # do look up in table for using database_id + yield self._get_sqla_engine( + schema=schema, nullpool=nullpool, source=source + ) def _get_sqla_engine( self, From 73b58e5132a99136329a02ac2ecb6573344c95cb Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 23 Nov 2022 14:29:20 -0500 Subject: [PATCH 04/16] add feature flag to enable validation checks --- superset/databases/ssh_tunnel/commands/create.py | 7 +++---- superset/databases/ssh_tunnel/commands/update.py | 6 ++---- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 8ae88a9f2d656..84fea93769986 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -20,7 +20,8 @@ from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError -from superset import app + +from superset import app, is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOCreateFailedError from superset.databases.dao import DatabaseDAO @@ -49,8 +50,6 @@ def run(self) -> Model: return tunnel def validate(self) -> None: - # TODO(hughhh): check to make sure the server port is not localhost - # using the config.SSH_TUNNEL_MANAGER - if ssh_tunnel_manager: + if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: ssh_tunnel_manager.validate() return diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index 146cc344964f0..71c2dd13650f8 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -20,7 +20,7 @@ from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError -from superset import app +from superset import app, is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOUpdateFailedError from superset.databases.ssh_tunnel.commands.exceptions import ( @@ -57,7 +57,5 @@ def validate(self) -> None: if not self._model: raise SSHTunnelNotFoundError() - # TODO(hughhh): check to make sure the server port is not localhost - # using the config.SSH_TUNNEL_MANAGER - if ssh_tunnel_manager: + if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: ssh_tunnel_manager.validate() From 58e37f1097c91ba00976de0276bb1471be278275 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 23 Nov 2022 20:41:50 -0500 Subject: [PATCH 05/16] update class object nots --- superset/config.py | 26 +++++++++---------- .../databases/ssh_tunnel/commands/create.py | 2 +- .../databases/ssh_tunnel/commands/update.py | 2 +- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/superset/config.py b/superset/config.py index 23afbd1c6a7ec..86523502bed6c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -493,22 +493,22 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: # ---------------------------------------------------------------------- # class SSHManager: -# def validate(self): -# # validation on CREATE + UPDATE on SSHTunnel Model -# raise NotImplemented() - -# def mutator(self): -# # override any ssh tunnel configuration object -# raise NotImplemented() - -# @property -# def local_bind_address(self): -# # set the local binding address for the local client -# # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder +# def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: +# # validation on CREATE + UPDATE on SSHTunnel Model +# # to block a request this function most raise an exception +# raise NotImplemented() + +# def mutator(self, ssh_tunnel_params: Dict[str, Any]) -> Dict[str, Any]: +# # override any ssh tunnel configuration object +# raise NotImplemented() + +# @property +# def local_bind_address(self): +# # set the local binding address for the local client +# # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder # # `server` return value # return "127.0.0.1" - SSH_TUNNEL_MANAGER = None # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars. diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 84fea93769986..4c5b666479d41 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -51,5 +51,5 @@ def run(self) -> Model: def validate(self) -> None: if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: - ssh_tunnel_manager.validate() + ssh_tunnel_manager.validate(self._properties) return diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index 71c2dd13650f8..6ff67aca50db2 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -58,4 +58,4 @@ def validate(self) -> None: raise SSHTunnelNotFoundError() if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: - ssh_tunnel_manager.validate() + ssh_tunnel_manager.validate(self._properties) From d720e759e48dcd274aa61cc6d982ec8e71da0f83 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Thu, 24 Nov 2022 21:53:56 -0500 Subject: [PATCH 06/16] add test for ssh manager --- .../databases/ssh_tunnel/commands/create.py | 5 ++- tests/unit_tests/config_test.py | 31 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 4c5b666479d41..5a1885b9a9834 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -35,6 +35,9 @@ class CreateSSHTunnelCommand(BaseCommand): + def __str__(self) -> str: + return super().__str__() + def __init__(self, database_id: int, data: Dict[str, Any]): self._properties = data.copy() self._properties["database_id"] = database_id @@ -52,4 +55,4 @@ def run(self) -> Model: def validate(self) -> None: if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: ssh_tunnel_manager.validate(self._properties) - return + return diff --git a/tests/unit_tests/config_test.py b/tests/unit_tests/config_test.py index 021193a6cd36e..c7029e9c9d971 100644 --- a/tests/unit_tests/config_test.py +++ b/tests/unit_tests/config_test.py @@ -131,7 +131,38 @@ def test_main_dttm_col(mocker: MockerFixture, test_table: "SqlaTable") -> None: test_table.fetch_metadata() assert test_table.main_dttm_col == "event_time" +def test_ssh_manager(mocker: MockerFixture): + from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand + + class TestSSHManager: + def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: + # validation on CREATE + UPDATE on SSHTunnel Model + # to block a request this function most raise an exception + return ssh_tunnel_params + + + ssh_tunnel_properties = { + "server_address": "123.132.123.1", + "bind_host": "localhost", + "bind_port": "5432", + "username": "foo", + "password": "bar", + } + + mock_ssh_manager = mocker.patch( + "superset.databases.ssh_tunnel.commands.create.ssh_tunnel_manager", + return_value=TestSSHManager(), + ) + + mocker.patch( + "superset.databases.ssh_tunnel.commands.create.is_feature_enabled", + return_value=True, + ) + CreateSSHTunnelCommand(1, ssh_tunnel_properties).run() + mock_ssh_manager.validate.assert_called_once() + + def test_main_dttm_col_nonexistent( mocker: MockerFixture, test_table: "SqlaTable", From 097179cf3742fa56a509ab13d1e4d3ec7351a529 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Fri, 25 Nov 2022 19:29:19 -0500 Subject: [PATCH 07/16] add test for update and engine instantiation --- superset/models/core.py | 5 +- tests/unit_tests/config_test.py | 96 ++++++++++++++++++++++++++++++++- 2 files changed, 98 insertions(+), 3 deletions(-) diff --git a/superset/models/core.py b/superset/models/core.py index 37c96feffadd4..8b1d2809cf88e 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -54,7 +54,7 @@ from sqlalchemy.schema import UniqueConstraint from sqlalchemy.sql import expression, Select -from superset import app, db_engine_specs +from superset import app, db_engine_specs, is_feature_enabled from superset.constants import PASSWORD_MASK from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import MetricType, TimeGrain @@ -383,8 +383,9 @@ def get_sqla_engine_with_context( ssh_tunnel.bind_host = url.host ssh_tunnel.bind_port = url.port ssh_params = ssh_tunnel.parameters() - if ssh_tunnel_manager: + if ssh_tunnel_manager and is_feature_enabled("SSH_TUNNELING"): ssh_params = ssh_tunnel_manager.mutate(ssh_params) + with sshtunnel.open_tunnel(**ssh_params) as server: yield self._get_sqla_engine( schema=schema, diff --git a/tests/unit_tests/config_test.py b/tests/unit_tests/config_test.py index c7029e9c9d971..decf1daf1d96a 100644 --- a/tests/unit_tests/config_test.py +++ b/tests/unit_tests/config_test.py @@ -131,7 +131,7 @@ def test_main_dttm_col(mocker: MockerFixture, test_table: "SqlaTable") -> None: test_table.fetch_metadata() assert test_table.main_dttm_col == "event_time" -def test_ssh_manager(mocker: MockerFixture): +def test_ssh_manager_create(mocker: MockerFixture): from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand class TestSSHManager: @@ -161,6 +161,100 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: CreateSSHTunnelCommand(1, ssh_tunnel_properties).run() mock_ssh_manager.validate.assert_called_once() + +def test_ssh_manager_update(mocker: MockerFixture): + from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand + from superset.databases.ssh_tunnel.models import SSHTunnel + + class TestSSHManager: + def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: + # validation on CREATE + UPDATE on SSHTunnel Model + # to block a request this function most raise an exception + return ssh_tunnel_params + + + ssh_tunnel_properties = { + "server_address": "123.132.123.1", + "bind_host": "localhost", + "bind_port": "5432", + "username": "foo", + "password": "bar", + } + + mock_ssh_manager = mocker.patch( + "superset.databases.ssh_tunnel.commands.update.ssh_tunnel_manager", + return_value=TestSSHManager(), + ) + + mocker.patch( + "superset.databases.ssh_tunnel.commands.update.is_feature_enabled", + return_value=True, + ) + + mocker.patch( + "superset.databases.ssh_tunnel.commands.update.SSHTunnelDAO", + return_value=SSHTunnel(), + ) + + UpdateSSHTunnelCommand(1, ssh_tunnel_properties).run() + mock_ssh_manager.validate.assert_called_once() + +def test_ssh_manager_mutate_before_engine(mocker: MockerFixture): + from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand + from superset.databases.ssh_tunnel.models import SSHTunnel + from superset.models.core import Database + + db = Database(database_name="test_database", sqlalchemy_uri="sqlite://locahost:3000") + + class TestSSHManager: + def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: + # validation on CREATE + UPDATE on SSHTunnel Model + # to block a request this function most raise an exception + return ssh_tunnel_params + + def mutate(self, ssh_tunnel_params: Dict[str, Any]) -> None: + return ssh_tunnel_params + + ssh_tunnel_properties = { + "server_address": "123.132.123.1", + "bind_host": "localhost", + "bind_port": "5432", + "username": "foo", + "password": "bar", + } + + ssh_tunnel_params = { + "ssh_address_or_host": "123.132.123.1", + "ssh_port": 8080, + "ssh_username": "hmiles", + "remote_bind_address": ("123.124.2131", 8080), + "local_bind_address": ("127.0.0.1",), + "ssh_password": "bar", + } + + mock_ssh_manager = mocker.patch( + "superset.models.core.ssh_tunnel_manager", + return_value=TestSSHManager().mutate(ssh_tunnel_properties), + ) + mock_ssh_manager.mutate.return_value = ssh_tunnel_params + + mocker.patch( + "superset.models.core.is_feature_enabled", + return_value=True, + ) + + # todo: figure out how to mock tunnel so we dnt try and establish a connection + mocker.patch( + "superset.models.core.sshtunnel.open_tunnel" + ) + + mocker.patch( + "superset.models.core.Database._get_sqla_engine" + ) + + + with db.get_sqla_engine_with_context(ssh_tunnel=SSHTunnel(**ssh_tunnel_properties)) as engine: + mock_ssh_manager.mutate.assert_called_once() def test_main_dttm_col_nonexistent( From 8c6d6e46a67fd1386af4ff146715cfc94bec9ebf Mon Sep 17 00:00:00 2001 From: hughhhh Date: Tue, 29 Nov 2022 11:47:10 -0500 Subject: [PATCH 08/16] fix pre-commit --- .../databases/ssh_tunnel/commands/update.py | 3 +- superset/models/core.py | 4 +-- tests/unit_tests/config_test.py | 28 +++++++++---------- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index 6ff67aca50db2..b1a3da7438a9c 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -37,6 +37,7 @@ config = app.config ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] + class UpdateSSHTunnelCommand(BaseCommand): def __init__(self, model_id: int, data: Dict[str, Any]): self._properties = data.copy() @@ -56,6 +57,6 @@ def validate(self) -> None: self._model = SSHTunnelDAO.find_by_id(self._model_id) if not self._model: raise SSHTunnelNotFoundError() - + if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: ssh_tunnel_manager.validate(self._properties) diff --git a/superset/models/core.py b/superset/models/core.py index 8b1d2809cf88e..48c0649c2833a 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -395,9 +395,7 @@ def get_sqla_engine_with_context( ) else: # do look up in table for using database_id - yield self._get_sqla_engine( - schema=schema, nullpool=nullpool, source=source - ) + yield self._get_sqla_engine(schema=schema, nullpool=nullpool, source=source) def _get_sqla_engine( self, diff --git a/tests/unit_tests/config_test.py b/tests/unit_tests/config_test.py index decf1daf1d96a..68649c3392e7c 100644 --- a/tests/unit_tests/config_test.py +++ b/tests/unit_tests/config_test.py @@ -131,6 +131,7 @@ def test_main_dttm_col(mocker: MockerFixture, test_table: "SqlaTable") -> None: test_table.fetch_metadata() assert test_table.main_dttm_col == "event_time" + def test_ssh_manager_create(mocker: MockerFixture): from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand @@ -139,7 +140,6 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: # validation on CREATE + UPDATE on SSHTunnel Model # to block a request this function most raise an exception return ssh_tunnel_params - ssh_tunnel_properties = { "server_address": "123.132.123.1", @@ -162,6 +162,7 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: CreateSSHTunnelCommand(1, ssh_tunnel_properties).run() mock_ssh_manager.validate.assert_called_once() + def test_ssh_manager_update(mocker: MockerFixture): from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand from superset.databases.ssh_tunnel.models import SSHTunnel @@ -171,7 +172,6 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: # validation on CREATE + UPDATE on SSHTunnel Model # to block a request this function most raise an exception return ssh_tunnel_params - ssh_tunnel_properties = { "server_address": "123.132.123.1", @@ -199,12 +199,15 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: UpdateSSHTunnelCommand(1, ssh_tunnel_properties).run() mock_ssh_manager.validate.assert_called_once() + def test_ssh_manager_mutate_before_engine(mocker: MockerFixture): from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand from superset.databases.ssh_tunnel.models import SSHTunnel from superset.models.core import Database - db = Database(database_name="test_database", sqlalchemy_uri="sqlite://locahost:3000") + db = Database( + database_name="test_database", sqlalchemy_uri="sqlite://locahost:3000" + ) class TestSSHManager: def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: @@ -214,7 +217,7 @@ def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: def mutate(self, ssh_tunnel_params: Dict[str, Any]) -> None: return ssh_tunnel_params - + ssh_tunnel_properties = { "server_address": "123.132.123.1", "bind_host": "localhost", @@ -244,19 +247,16 @@ def mutate(self, ssh_tunnel_params: Dict[str, Any]) -> None: ) # todo: figure out how to mock tunnel so we dnt try and establish a connection - mocker.patch( - "superset.models.core.sshtunnel.open_tunnel" - ) - - mocker.patch( - "superset.models.core.Database._get_sqla_engine" - ) + mocker.patch("superset.models.core.sshtunnel.open_tunnel") + mocker.patch("superset.models.core.Database._get_sqla_engine") - with db.get_sqla_engine_with_context(ssh_tunnel=SSHTunnel(**ssh_tunnel_properties)) as engine: + with db.get_sqla_engine_with_context( + ssh_tunnel=SSHTunnel(**ssh_tunnel_properties) + ) as engine: mock_ssh_manager.mutate.assert_called_once() - - + + def test_main_dttm_col_nonexistent( mocker: MockerFixture, test_table: "SqlaTable", From 96cdbf548748bbd87a10fefcc7dc7468140fadd3 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 30 Nov 2022 13:28:18 -0500 Subject: [PATCH 09/16] pull config from current_app --- superset/databases/ssh_tunnel/commands/create.py | 16 +++++++--------- superset/databases/ssh_tunnel/commands/update.py | 12 ++++-------- 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 5a1885b9a9834..1183bcb1661fd 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -17,30 +17,28 @@ import logging from typing import Any, Dict, List, Optional +from flask import current_app from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError - -from superset import app, is_feature_enabled +from superset import is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOCreateFailedError from superset.databases.dao import DatabaseDAO from superset.databases.ssh_tunnel.commands.exceptions import SSHTunnelCreateFailedError from superset.databases.ssh_tunnel.dao import SSHTunnelDAO -config = app.config -ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] - logger = logging.getLogger(__name__) class CreateSSHTunnelCommand(BaseCommand): def __str__(self) -> str: return super().__str__() - + def __init__(self, database_id: int, data: Dict[str, Any]): self._properties = data.copy() self._properties["database_id"] = database_id + self.ssh_tunnel_manager = current_app.config["SSH_TUNNEL_MANAGER"] def run(self) -> Model: self.validate() @@ -53,6 +51,6 @@ def run(self) -> Model: return tunnel def validate(self) -> None: - if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: - ssh_tunnel_manager.validate(self._properties) - return + if is_feature_enabled("SSH_TUNNELING") and self.ssh_tunnel_manager: + self.ssh_tunnel_manager.validate(self._properties) + return diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index b1a3da7438a9c..c0392ff331a3c 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -17,15 +17,13 @@ import logging from typing import Any, Dict, List, Optional +from flask import current_app from flask_appbuilder.models.sqla import Model -from marshmallow import ValidationError from superset import app, is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOUpdateFailedError from superset.databases.ssh_tunnel.commands.exceptions import ( - SSHTunnelDeleteFailedError, - SSHTunnelInvalidError, SSHTunnelNotFoundError, SSHTunnelUpdateFailedError, ) @@ -34,15 +32,13 @@ logger = logging.getLogger(__name__) -config = app.config -ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] - class UpdateSSHTunnelCommand(BaseCommand): def __init__(self, model_id: int, data: Dict[str, Any]): self._properties = data.copy() self._model_id = model_id self._model: Optional[SSHTunnel] = None + self.ssh_tunnel_manager = current_app.config["SSH_TUNNEL_MANAGER"] def run(self) -> Model: self.validate() @@ -58,5 +54,5 @@ def validate(self) -> None: if not self._model: raise SSHTunnelNotFoundError() - if is_feature_enabled("SSH_TUNNELING") and ssh_tunnel_manager: - ssh_tunnel_manager.validate(self._properties) + if is_feature_enabled("SSH_TUNNELING") and self.ssh_tunnel_manager: + self.ssh_tunnel_manager.validate(self._properties) From 87f3065605fb73e1b7f077cf0ca6c96437947825 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 14 Dec 2022 12:58:43 -0500 Subject: [PATCH 10/16] fix pylint --- superset/config.py | 32 ++++++++----------- .../databases/ssh_tunnel/commands/create.py | 4 --- .../databases/ssh_tunnel/commands/update.py | 1 - 3 files changed, 14 insertions(+), 23 deletions(-) diff --git a/superset/config.py b/superset/config.py index 86523502bed6c..d7ef5aa5b4345 100644 --- a/superset/config.py +++ b/superset/config.py @@ -492,24 +492,20 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: # FIREWALL (only port 22 is open) # ---------------------------------------------------------------------- -# class SSHManager: -# def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: -# # validation on CREATE + UPDATE on SSHTunnel Model -# # to block a request this function most raise an exception -# raise NotImplemented() - -# def mutator(self, ssh_tunnel_params: Dict[str, Any]) -> Dict[str, Any]: -# # override any ssh tunnel configuration object -# raise NotImplemented() - -# @property -# def local_bind_address(self): -# # set the local binding address for the local client -# # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder -# # `server` return value -# return "127.0.0.1" - -SSH_TUNNEL_MANAGER = None +class SSHManager: + def mutator(self, ssh_tunnel_params: Dict[str, Any]) -> Dict[str, Any]: + # override any ssh tunnel configuration object + return ssh_tunnel_params + + @property + def local_bind_address(self): + # set the local binding address for the local client + # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder + # `server` return value + return "127.0.0.1" + + +SSH_TUNNEL_MANAGER = SSHManager() # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars. DEFAULT_FEATURE_FLAGS.update( diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 6a785d40ca160..2fd138565151e 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -21,7 +21,6 @@ from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError -from superset import is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOCreateFailedError from superset.databases.ssh_tunnel.commands.exceptions import ( @@ -36,9 +35,6 @@ class CreateSSHTunnelCommand(BaseCommand): - def __str__(self) -> str: - return super().__str__() - def __init__(self, database_id: int, data: Dict[str, Any]): self._properties = data.copy() self._properties["database_id"] = database_id diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index c5ad9c5a2a7bd..2a913e2cd08b0 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -20,7 +20,6 @@ from flask import current_app from flask_appbuilder.models.sqla import Model -from superset import app, is_feature_enabled from superset.commands.base import BaseCommand from superset.dao.exceptions import DAOUpdateFailedError from superset.databases.ssh_tunnel.commands.exceptions import ( From 1f995036a7fa1889e27c3553a6944b9eb5019bc4 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 14 Dec 2022 13:19:14 -0500 Subject: [PATCH 11/16] remove validation tests --- tests/unit_tests/config_test.py | 125 -------------------------------- 1 file changed, 125 deletions(-) diff --git a/tests/unit_tests/config_test.py b/tests/unit_tests/config_test.py index 68649c3392e7c..021193a6cd36e 100644 --- a/tests/unit_tests/config_test.py +++ b/tests/unit_tests/config_test.py @@ -132,131 +132,6 @@ def test_main_dttm_col(mocker: MockerFixture, test_table: "SqlaTable") -> None: assert test_table.main_dttm_col == "event_time" -def test_ssh_manager_create(mocker: MockerFixture): - from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand - - class TestSSHManager: - def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: - # validation on CREATE + UPDATE on SSHTunnel Model - # to block a request this function most raise an exception - return ssh_tunnel_params - - ssh_tunnel_properties = { - "server_address": "123.132.123.1", - "bind_host": "localhost", - "bind_port": "5432", - "username": "foo", - "password": "bar", - } - - mock_ssh_manager = mocker.patch( - "superset.databases.ssh_tunnel.commands.create.ssh_tunnel_manager", - return_value=TestSSHManager(), - ) - - mocker.patch( - "superset.databases.ssh_tunnel.commands.create.is_feature_enabled", - return_value=True, - ) - - CreateSSHTunnelCommand(1, ssh_tunnel_properties).run() - mock_ssh_manager.validate.assert_called_once() - - -def test_ssh_manager_update(mocker: MockerFixture): - from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand - from superset.databases.ssh_tunnel.models import SSHTunnel - - class TestSSHManager: - def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: - # validation on CREATE + UPDATE on SSHTunnel Model - # to block a request this function most raise an exception - return ssh_tunnel_params - - ssh_tunnel_properties = { - "server_address": "123.132.123.1", - "bind_host": "localhost", - "bind_port": "5432", - "username": "foo", - "password": "bar", - } - - mock_ssh_manager = mocker.patch( - "superset.databases.ssh_tunnel.commands.update.ssh_tunnel_manager", - return_value=TestSSHManager(), - ) - - mocker.patch( - "superset.databases.ssh_tunnel.commands.update.is_feature_enabled", - return_value=True, - ) - - mocker.patch( - "superset.databases.ssh_tunnel.commands.update.SSHTunnelDAO", - return_value=SSHTunnel(), - ) - - UpdateSSHTunnelCommand(1, ssh_tunnel_properties).run() - mock_ssh_manager.validate.assert_called_once() - - -def test_ssh_manager_mutate_before_engine(mocker: MockerFixture): - from superset.databases.ssh_tunnel.commands.update import UpdateSSHTunnelCommand - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - db = Database( - database_name="test_database", sqlalchemy_uri="sqlite://locahost:3000" - ) - - class TestSSHManager: - def validate(self, ssh_tunnel_params: Dict[str, Any]) -> None: - # validation on CREATE + UPDATE on SSHTunnel Model - # to block a request this function most raise an exception - return ssh_tunnel_params - - def mutate(self, ssh_tunnel_params: Dict[str, Any]) -> None: - return ssh_tunnel_params - - ssh_tunnel_properties = { - "server_address": "123.132.123.1", - "bind_host": "localhost", - "bind_port": "5432", - "username": "foo", - "password": "bar", - } - - ssh_tunnel_params = { - "ssh_address_or_host": "123.132.123.1", - "ssh_port": 8080, - "ssh_username": "hmiles", - "remote_bind_address": ("123.124.2131", 8080), - "local_bind_address": ("127.0.0.1",), - "ssh_password": "bar", - } - - mock_ssh_manager = mocker.patch( - "superset.models.core.ssh_tunnel_manager", - return_value=TestSSHManager().mutate(ssh_tunnel_properties), - ) - mock_ssh_manager.mutate.return_value = ssh_tunnel_params - - mocker.patch( - "superset.models.core.is_feature_enabled", - return_value=True, - ) - - # todo: figure out how to mock tunnel so we dnt try and establish a connection - mocker.patch("superset.models.core.sshtunnel.open_tunnel") - - mocker.patch("superset.models.core.Database._get_sqla_engine") - - with db.get_sqla_engine_with_context( - ssh_tunnel=SSHTunnel(**ssh_tunnel_properties) - ) as engine: - mock_ssh_manager.mutate.assert_called_once() - - def test_main_dttm_col_nonexistent( mocker: MockerFixture, test_table: "SqlaTable", From 874c8e3886c1fc9d9791ff39fb1aa9aa53c5e298 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 14 Dec 2022 13:31:49 -0500 Subject: [PATCH 12/16] okkkk --- superset/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/config.py b/superset/config.py index d7ef5aa5b4345..45f0be3b3fcb0 100644 --- a/superset/config.py +++ b/superset/config.py @@ -505,7 +505,7 @@ def local_bind_address(self): return "127.0.0.1" -SSH_TUNNEL_MANAGER = SSHManager() +SSH_TUNNEL_MANAGER = SSHManager # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars. DEFAULT_FEATURE_FLAGS.update( From 71ff1535cecf474cf3963340e2c87fd62ac012ae Mon Sep 17 00:00:00 2001 From: hughhhh Date: Wed, 14 Dec 2022 16:28:25 -0500 Subject: [PATCH 13/16] ignore nonsense --- superset/config.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/superset/config.py b/superset/config.py index 45f0be3b3fcb0..102b272d6f46c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -492,20 +492,17 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: # FIREWALL (only port 22 is open) # ---------------------------------------------------------------------- -class SSHManager: - def mutator(self, ssh_tunnel_params: Dict[str, Any]) -> Dict[str, Any]: +class SSHManager: # pylint: disable=too-few-public-methods + local_bind_address = "127.0.0.1" + + def mutator( # pylint: disable=no-self-use + self, ssh_tunnel_params: Dict[str, Any] + ) -> Dict[str, Any]: # override any ssh tunnel configuration object return ssh_tunnel_params - @property - def local_bind_address(self): - # set the local binding address for the local client - # the port will be dynamically configured by the sshtunnel.SSHTunnelForwarder - # `server` return value - return "127.0.0.1" - -SSH_TUNNEL_MANAGER = SSHManager +SSH_TUNNEL_MANAGER = SSHManager # pylint: disable=invalid-name # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars. DEFAULT_FEATURE_FLAGS.update( @@ -1496,7 +1493,7 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument try: # pylint: disable=import-error,wildcard-import,unused-wildcard-import import superset_config - from superset_config import * # type:ignore + from superset_config import * # type: ignore print(f"Loaded your LOCAL configuration at [{superset_config.__file__}]") except Exception: From 27c5023309af2a62f13aa73d450684926a42a426 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Thu, 15 Dec 2022 01:12:09 -0500 Subject: [PATCH 14/16] ok i need to test this --- superset/config.py | 13 ++++++++++--- superset/databases/ssh_tunnel/commands/create.py | 2 -- superset/databases/ssh_tunnel/commands/update.py | 2 -- superset/databases/ssh_tunnel/models.py | 3 ++- superset/models/core.py | 7 ++++--- 5 files changed, 16 insertions(+), 11 deletions(-) diff --git a/superset/config.py b/superset/config.py index 102b272d6f46c..5117795e63867 100644 --- a/superset/config.py +++ b/superset/config.py @@ -45,6 +45,7 @@ ) import pkg_resources +import sshtunnel from cachelib.base import BaseCache from celery.schedules import crontab from dateutil import tz @@ -496,10 +497,16 @@ class SSHManager: # pylint: disable=too-few-public-methods local_bind_address = "127.0.0.1" def mutator( # pylint: disable=no-self-use - self, ssh_tunnel_params: Dict[str, Any] - ) -> Dict[str, Any]: + self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder + ) -> str: # override any ssh tunnel configuration object - return ssh_tunnel_params + from superset.databases.utils import make_url_safe + + url = make_url_safe(sqlalchemy_url) + return url.set( + host=self.local_bind_address, + port=server.local_bind_port, + ) SSH_TUNNEL_MANAGER = SSHManager # pylint: disable=invalid-name diff --git a/superset/databases/ssh_tunnel/commands/create.py b/superset/databases/ssh_tunnel/commands/create.py index 2fd138565151e..b2e62f340b0bb 100644 --- a/superset/databases/ssh_tunnel/commands/create.py +++ b/superset/databases/ssh_tunnel/commands/create.py @@ -17,7 +17,6 @@ import logging from typing import Any, Dict, List, Optional -from flask import current_app from flask_appbuilder.models.sqla import Model from marshmallow import ValidationError @@ -38,7 +37,6 @@ class CreateSSHTunnelCommand(BaseCommand): def __init__(self, database_id: int, data: Dict[str, Any]): self._properties = data.copy() self._properties["database_id"] = database_id - self.ssh_tunnel_manager = current_app.config["SSH_TUNNEL_MANAGER"] def run(self) -> Model: self.validate() diff --git a/superset/databases/ssh_tunnel/commands/update.py b/superset/databases/ssh_tunnel/commands/update.py index 2a913e2cd08b0..8d2feaf1b0d52 100644 --- a/superset/databases/ssh_tunnel/commands/update.py +++ b/superset/databases/ssh_tunnel/commands/update.py @@ -17,7 +17,6 @@ import logging from typing import Any, Dict, Optional -from flask import current_app from flask_appbuilder.models.sqla import Model from superset.commands.base import BaseCommand @@ -39,7 +38,6 @@ def __init__(self, model_id: int, data: Dict[str, Any]): self._properties = data.copy() self._model_id = model_id self._model: Optional[SSHTunnel] = None - self.ssh_tunnel_manager = current_app.config["SSH_TUNNEL_MANAGER"] def run(self) -> Model: self.validate() diff --git a/superset/databases/ssh_tunnel/models.py b/superset/databases/ssh_tunnel/models.py index 17f4628f8070f..0a1fec9cf7792 100644 --- a/superset/databases/ssh_tunnel/models.py +++ b/superset/databases/ssh_tunnel/models.py @@ -32,6 +32,7 @@ ) app_config = current_app.config +ssh_manager = app_config["SSH_TUNNEL_MANAGER"] class SSHTunnel(Model, AuditMixinNullable, ExtraJSONMixin, ImportExportMixin): @@ -74,7 +75,7 @@ def parameters(self, bind_host: str, bind_port: int) -> Dict[str, Any]: "ssh_port": self.server_port, "ssh_username": self.username, "remote_bind_address": (bind_host, bind_port), - "local_bind_address": (SSH_TUNNELLING_LOCAL_BIND_ADDRESS,), + "local_bind_address": (ssh_manager.local_bind_address,), } if self.password: diff --git a/superset/models/core.py b/superset/models/core.py index adf116b2dbfd2..3a192c9bea77a 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -66,8 +66,7 @@ from superset.utils.memoized import memoized config = app.config - -ssh_tunnel_manager = config["SSH_TUNNEL_MANAGER"] +ssh_manager = config["SSH_TUNNEL_MANAGER"] custom_password_store = config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"] stats_logger = config["STATS_LOGGER"] log_query = config["QUERY_LOGGER"] @@ -450,9 +449,11 @@ def _get_sqla_engine( if ssh_tunnel_server: # update sqlalchemy_url + sqlalchemy_url = ssh_manager.mutate(sqlalchemy_url, ssh_tunnel_server) url = make_url_safe(sqlalchemy_url) + sqlalchemy_url = url.set( - host=SSH_TUNNELLING_LOCAL_BIND_ADDRESS, + host=ssh_tunnel_server.local_bind_address, port=ssh_tunnel_server.local_bind_port, ) From 146c47632d660ff5537d94981c5471139bdb20b7 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Thu, 15 Dec 2022 01:21:51 -0500 Subject: [PATCH 15/16] ok i need to test this --- superset/config.py | 7 ++----- superset/constants.py | 2 -- superset/databases/ssh_tunnel/models.py | 1 - superset/models/core.py | 11 ++--------- 4 files changed, 4 insertions(+), 17 deletions(-) diff --git a/superset/config.py b/superset/config.py index 5117795e63867..486f574b923c6 100644 --- a/superset/config.py +++ b/superset/config.py @@ -57,6 +57,7 @@ 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 @@ -496,12 +497,8 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: class SSHManager: # pylint: disable=too-few-public-methods local_bind_address = "127.0.0.1" - def mutator( # pylint: disable=no-self-use - self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder - ) -> str: + def mutator(self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder) -> str: # override any ssh tunnel configuration object - from superset.databases.utils import make_url_safe - url = make_url_safe(sqlalchemy_url) return url.set( host=self.local_bind_address, diff --git a/superset/constants.py b/superset/constants.py index c0fbb7c2cd8db..7d759acf6741c 100644 --- a/superset/constants.py +++ b/superset/constants.py @@ -34,8 +34,6 @@ NO_TIME_RANGE = "No filter" -SSH_TUNNELLING_LOCAL_BIND_ADDRESS = "127.0.0.1" - class RouteMethod: # pylint: disable=too-few-public-methods """ diff --git a/superset/databases/ssh_tunnel/models.py b/superset/databases/ssh_tunnel/models.py index 0a1fec9cf7792..5f334d8c154bb 100644 --- a/superset/databases/ssh_tunnel/models.py +++ b/superset/databases/ssh_tunnel/models.py @@ -23,7 +23,6 @@ from sqlalchemy.orm import backref, relationship from sqlalchemy_utils import EncryptedType -from superset.constants import SSH_TUNNELLING_LOCAL_BIND_ADDRESS from superset.models.core import Database from superset.models.helpers import ( AuditMixinNullable, diff --git a/superset/models/core.py b/superset/models/core.py index 3a192c9bea77a..19977fa022665 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -55,7 +55,7 @@ from sqlalchemy.sql import expression, Select from superset import app, db_engine_specs -from superset.constants import PASSWORD_MASK, SSH_TUNNELLING_LOCAL_BIND_ADDRESS +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 @@ -448,15 +448,8 @@ def _get_sqla_engine( ) if ssh_tunnel_server: - # update sqlalchemy_url + # update sqlalchemy_url with ssh tunnel manager info sqlalchemy_url = ssh_manager.mutate(sqlalchemy_url, ssh_tunnel_server) - url = make_url_safe(sqlalchemy_url) - - sqlalchemy_url = url.set( - host=ssh_tunnel_server.local_bind_address, - port=ssh_tunnel_server.local_bind_port, - ) - try: return create_engine(sqlalchemy_url, **params) except Exception as ex: From 329e3a5811ebfb513a120bafbb4fc05eff7658c0 Mon Sep 17 00:00:00 2001 From: hughhhh Date: Thu, 15 Dec 2022 14:07:33 -0500 Subject: [PATCH 16/16] verified working update --- superset/config.py | 5 +++-- superset/models/core.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/superset/config.py b/superset/config.py index 486f574b923c6..71b220333977a 100644 --- a/superset/config.py +++ b/superset/config.py @@ -497,11 +497,12 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]: class SSHManager: # pylint: disable=too-few-public-methods local_bind_address = "127.0.0.1" - def mutator(self, sqlalchemy_url: str, server: sshtunnel.SSHTunnelForwarder) -> str: + @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=self.local_bind_address, + host=cls.local_bind_address, port=server.local_bind_port, ) diff --git a/superset/models/core.py b/superset/models/core.py index 19977fa022665..51bd0ce21b6ff 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -449,7 +449,7 @@ def _get_sqla_engine( if ssh_tunnel_server: # update sqlalchemy_url with ssh tunnel manager info - sqlalchemy_url = ssh_manager.mutate(sqlalchemy_url, ssh_tunnel_server) + sqlalchemy_url = ssh_manager.mutator(sqlalchemy_url, ssh_tunnel_server) try: return create_engine(sqlalchemy_url, **params) except Exception as ex: