From 62d07a9916b91a9551b0a7b5c8ecc77be4f78ebe Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Wed, 3 Dec 2025 14:26:35 -0500 Subject: [PATCH] chore: cleanup ssh tunnel (#34388) (cherry picked from commit c458f99dd490f49237a08482f4d6d59fced30c4c) --- superset/commands/database/create.py | 42 ++--- superset/commands/database/export.py | 2 +- .../commands/database/importers/v1/utils.py | 13 +- .../commands/database/ssh_tunnel/create.py | 100 ----------- .../commands/database/ssh_tunnel/delete.py | 52 ------ .../commands/database/ssh_tunnel/update.py | 82 --------- .../commands/database/sync_permissions.py | 25 +-- superset/commands/database/test_connection.py | 85 ++++++---- superset/commands/database/update.py | 38 +---- superset/commands/database/utils.py | 14 +- superset/commands/dataset/export.py | 3 +- superset/daos/database.py | 123 ++++++++++---- superset/databases/api.py | 84 +-------- superset/databases/schemas.py | 18 ++ superset/databases/ssh_tunnel/models.py | 14 +- superset/extensions/ssh.py | 6 +- superset/models/core.py | 40 +---- superset/utils/ssh_tunnel.py | 23 ++- tests/integration_tests/commands_test.py | 4 +- .../integration_tests/databases/api_tests.py | 105 +++++------- .../databases/commands_tests.py | 15 +- .../databases/ssh_tunnel/commands/__init__.py | 16 -- .../ssh_tunnel/commands/commands_tests.py | 58 ------- tests/integration_tests/tags/api_tests.py | 7 +- .../databases/sync_permissions_test.py | 11 +- tests/unit_tests/databases/api_test.py | 159 +----------------- .../databases/commands/utils_test.py | 9 +- tests/unit_tests/databases/dao/dao_tests.py | 7 +- tests/unit_tests/databases/filters_test.py | 5 +- .../databases/ssh_tunnel/commands/__init__.py | 16 -- .../ssh_tunnel/commands/create_test.py | 148 ---------------- .../ssh_tunnel/commands/delete_test.py | 73 -------- .../ssh_tunnel/commands/update_test.py | 155 ----------------- .../databases/ssh_tunnel/dao_tests.py | 37 ---- tests/unit_tests/models/core_test.py | 2 +- 35 files changed, 304 insertions(+), 1287 deletions(-) delete mode 100644 superset/commands/database/ssh_tunnel/create.py delete mode 100644 superset/commands/database/ssh_tunnel/delete.py delete mode 100644 superset/commands/database/ssh_tunnel/update.py delete mode 100644 tests/integration_tests/databases/ssh_tunnel/commands/__init__.py delete mode 100644 tests/integration_tests/databases/ssh_tunnel/commands/commands_tests.py delete mode 100644 tests/unit_tests/databases/ssh_tunnel/commands/__init__.py delete mode 100644 tests/unit_tests/databases/ssh_tunnel/commands/create_test.py delete mode 100644 tests/unit_tests/databases/ssh_tunnel/commands/delete_test.py delete mode 100644 tests/unit_tests/databases/ssh_tunnel/commands/update_test.py delete mode 100644 tests/unit_tests/databases/ssh_tunnel/dao_tests.py diff --git a/superset/commands/database/create.py b/superset/commands/database/create.py index 90c4e3bd323..ce56b39a829 100644 --- a/superset/commands/database/create.py +++ b/superset/commands/database/create.py @@ -22,7 +22,6 @@ from flask import current_app as app 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.commands.database.exceptions import ( DatabaseConnectionFailedError, @@ -31,7 +30,6 @@ from superset.commands.database.exceptions import ( DatabaseInvalidError, DatabaseRequiredFieldValidationError, ) -from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand from superset.commands.database.ssh_tunnel.exceptions import ( SSHTunnelCreateFailedError, SSHTunnelDatabasePortError, @@ -41,7 +39,7 @@ from superset.commands.database.ssh_tunnel.exceptions import ( from superset.commands.database.test_connection import TestConnectionDatabaseCommand from superset.commands.database.utils import add_permissions from superset.daos.database import DatabaseDAO -from superset.databases.ssh_tunnel.models import SSHTunnel +from superset.databases.utils import make_url_safe from superset.exceptions import OAuth2RedirectError, SupersetErrorsException from superset.extensions import event_logger from superset.models.core import Database @@ -59,6 +57,12 @@ class CreateDatabaseCommand(BaseCommand): def run(self) -> Model: self.validate() + if "sqlalchemy_uri" not in self._properties: + raise DatabaseRequiredFieldValidationError("sqlalchemy_uri") + + url = make_url_safe(self._properties["sqlalchemy_uri"]) + engine = url.get_backend_name() + try: # Test connection before starting create transaction TestConnectionDatabaseCommand(self._properties).run() @@ -74,32 +78,21 @@ class CreateDatabaseCommand(BaseCommand): ) as ex: event_logger.log_with_context( action=f"db_creation_failed.{ex.__class__.__name__}", - engine=self._properties.get("sqlalchemy_uri", "").split(":")[0], + engine=engine, ) # So we can show the original message raise except Exception as ex: event_logger.log_with_context( action=f"db_creation_failed.{ex.__class__.__name__}", - engine=self._properties.get("sqlalchemy_uri", "").split(":")[0], + engine=engine, ) raise DatabaseConnectionFailedError() from ex - ssh_tunnel: Optional[SSHTunnel] = None - try: + # create database and associated schema/catalog permissions database = self._create_database() - - if ssh_tunnel_properties := self._properties.get("ssh_tunnel"): - if not is_feature_enabled("SSH_TUNNELING"): - raise SSHTunnelingNotEnabledError() - - ssh_tunnel = CreateSSHTunnelCommand( - database, ssh_tunnel_properties - ).run() - - # add catalog/schema permissions - add_permissions(database, ssh_tunnel) + add_permissions(database) except ( SSHTunnelInvalidError, SSHTunnelCreateFailedError, @@ -108,7 +101,7 @@ class CreateDatabaseCommand(BaseCommand): ) as ex: event_logger.log_with_context( action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel", - engine=self._properties.get("sqlalchemy_uri", "").split(":")[0], + engine=engine, ) # So we can show the original message raise @@ -118,25 +111,23 @@ class CreateDatabaseCommand(BaseCommand): ) as ex: event_logger.log_with_context( action=f"db_creation_failed.{ex.__class__.__name__}", - engine=database.db_engine_spec.__name__, + engine=engine, ) raise DatabaseCreateFailedError() from ex - if ssh_tunnel: - stats_logger.incr("db_creation_success.ssh_tunnel") - return database def validate(self) -> None: exceptions: list[ValidationError] = [] + sqlalchemy_uri: Optional[str] = self._properties.get("sqlalchemy_uri") - database_name: Optional[str] = self._properties.get("database_name") if not sqlalchemy_uri: exceptions.append(DatabaseRequiredFieldValidationError("sqlalchemy_uri")) + + database_name: Optional[str] = self._properties.get("database_name") if not database_name: exceptions.append(DatabaseRequiredFieldValidationError("database_name")) else: - # Check database_name uniqueness if not DatabaseDAO.validate_uniqueness(database_name): exceptions.append(DatabaseExistsValidationError()) @@ -161,4 +152,5 @@ class CreateDatabaseCommand(BaseCommand): database = DatabaseDAO.create(attributes=self._properties) database.set_sqlalchemy_uri(database.sqlalchemy_uri) + return database diff --git a/superset/commands/database/export.py b/superset/commands/database/export.py index 97333a80aaa..56929155cd3 100644 --- a/superset/commands/database/export.py +++ b/superset/commands/database/export.py @@ -88,7 +88,7 @@ class ExportDatabasesCommand(ExportModelsCommand): "schemas_allowed_for_file_upload" ) - if ssh_tunnel := DatabaseDAO.get_ssh_tunnel(model.id): + if ssh_tunnel := model.ssh_tunnel: ssh_tunnel_payload = ssh_tunnel.export_to_dict( recursive=False, include_parent_ref=False, diff --git a/superset/commands/database/importers/v1/utils.py b/superset/commands/database/importers/v1/utils.py index 55b903744a6..2b2fd2ac3fd 100644 --- a/superset/commands/database/importers/v1/utils.py +++ b/superset/commands/database/importers/v1/utils.py @@ -68,11 +68,11 @@ def import_database( # TODO (betodealmeida): move this logic to import_from_dict config["extra"] = json.dumps(config["extra"]) - # Before it gets removed in import_from_dict ssh_tunnel_config = config.pop("ssh_tunnel", None) # set SQLAlchemy URI via `set_sqlalchemy_uri` so that the password gets masked sqlalchemy_uri = config.pop("sqlalchemy_uri") + # TODO (betodealmeida): we should use the `CreateDatabaseCommand` for imports database: Database = Database.import_from_dict(config, recursive=False) database.set_sqlalchemy_uri(sqlalchemy_uri) @@ -81,14 +81,13 @@ def import_database( if ssh_tunnel_config: ssh_tunnel_config["database_id"] = database.id - ssh_tunnel = SSHTunnel.import_from_dict(ssh_tunnel_config, recursive=False) - else: - ssh_tunnel = None - - # TODO (betodealmeida): we should use the `CreateDatabaseCommand` for imports + database.ssh_tunnel = SSHTunnel.import_from_dict( + ssh_tunnel_config, + recursive=False, + ) try: - add_permissions(database, ssh_tunnel) + add_permissions(database) except SupersetDBAPIConnectionError as ex: logger.warning(ex.message) diff --git a/superset/commands/database/ssh_tunnel/create.py b/superset/commands/database/ssh_tunnel/create.py deleted file mode 100644 index 9e9161ea51f..00000000000 --- a/superset/commands/database/ssh_tunnel/create.py +++ /dev/null @@ -1,100 +0,0 @@ -# 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. -import logging -from functools import partial -from typing import Any, Optional - -from flask_appbuilder.models.sqla import Model -from marshmallow import ValidationError - -from superset.commands.base import BaseCommand -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelCreateFailedError, - SSHTunnelDatabasePortError, - SSHTunnelInvalidError, - SSHTunnelRequiredFieldValidationError, -) -from superset.daos.database import SSHTunnelDAO -from superset.databases.utils import make_url_safe -from superset.extensions import event_logger -from superset.models.core import Database -from superset.utils.decorators import on_error, transaction -from superset.utils.ssh_tunnel import get_default_port - -logger = logging.getLogger(__name__) - - -class CreateSSHTunnelCommand(BaseCommand): - _database: Database - - def __init__(self, database: Database, data: dict[str, Any]): - self._properties = data.copy() - self._properties["database"] = database - self._database = database - - @transaction(on_error=partial(on_error, reraise=SSHTunnelCreateFailedError)) - def run(self) -> Model: - """ - Create an SSH tunnel. - - :returns: The SSH tunnel model - :raises SSHTunnelCreateFailedError: If the model creation fails - :raises SSHTunnelInvalidError: If the configuration are invalid - """ - - self.validate() - return SSHTunnelDAO.create(attributes=self._properties) - - def validate(self) -> None: - # TODO(hughhh): check to make sure the server port is not localhost - # using the config.SSH_TUNNEL_MANAGER - - exceptions: list[ValidationError] = [] - server_address: Optional[str] = self._properties.get("server_address") - server_port: Optional[int] = self._properties.get("server_port") - username: Optional[str] = self._properties.get("username") - password: Optional[str] = self._properties.get("password") - private_key: Optional[str] = self._properties.get("private_key") - private_key_password: Optional[str] = self._properties.get( - "private_key_password" - ) - url = make_url_safe(self._database.sqlalchemy_uri) - backend = url.get_backend_name() - port = url.port or get_default_port(backend) - if not port: - raise SSHTunnelDatabasePortError() - if not server_address: - exceptions.append(SSHTunnelRequiredFieldValidationError("server_address")) - if not server_port: - exceptions.append(SSHTunnelRequiredFieldValidationError("server_port")) - if not username: - exceptions.append(SSHTunnelRequiredFieldValidationError("username")) - if not private_key and not password: - exceptions.append(SSHTunnelRequiredFieldValidationError("password")) - if private_key_password and private_key is None: - exceptions.append(SSHTunnelRequiredFieldValidationError("private_key")) - if exceptions: - exception = SSHTunnelInvalidError() - exception.extend(exceptions) - event_logger.log_with_context( - # pylint: disable=consider-using-f-string - action="ssh_tunnel_creation_failed.{}.{}".format( - exception.__class__.__name__, - ".".join(exception.get_list_classnames()), - ) - ) - raise exception diff --git a/superset/commands/database/ssh_tunnel/delete.py b/superset/commands/database/ssh_tunnel/delete.py deleted file mode 100644 index 8c742307aa8..00000000000 --- a/superset/commands/database/ssh_tunnel/delete.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. -import logging -from functools import partial -from typing import Optional - -from superset import is_feature_enabled -from superset.commands.base import BaseCommand -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelDeleteFailedError, - SSHTunnelingNotEnabledError, - SSHTunnelNotFoundError, -) -from superset.daos.database import SSHTunnelDAO -from superset.databases.ssh_tunnel.models import SSHTunnel -from superset.utils.decorators import on_error, transaction - -logger = logging.getLogger(__name__) - - -class DeleteSSHTunnelCommand(BaseCommand): - def __init__(self, model_id: int): - self._model_id = model_id - self._model: Optional[SSHTunnel] = None - - @transaction(on_error=partial(on_error, reraise=SSHTunnelDeleteFailedError)) - def run(self) -> None: - if not is_feature_enabled("SSH_TUNNELING"): - raise SSHTunnelingNotEnabledError() - self.validate() - assert self._model - SSHTunnelDAO.delete([self._model]) - - def validate(self) -> None: - # Validate/populate model exists - self._model = SSHTunnelDAO.find_by_id(self._model_id) - if not self._model: - raise SSHTunnelNotFoundError() diff --git a/superset/commands/database/ssh_tunnel/update.py b/superset/commands/database/ssh_tunnel/update.py deleted file mode 100644 index 763d36e89a0..00000000000 --- a/superset/commands/database/ssh_tunnel/update.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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. -import logging -from functools import partial -from typing import Any, Optional - -from flask_appbuilder.models.sqla import Model - -from superset.commands.base import BaseCommand -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelDatabasePortError, - SSHTunnelInvalidError, - SSHTunnelNotFoundError, - SSHTunnelRequiredFieldValidationError, - SSHTunnelUpdateFailedError, -) -from superset.daos.database import SSHTunnelDAO -from superset.databases.ssh_tunnel.models import SSHTunnel -from superset.databases.utils import make_url_safe -from superset.utils.decorators import on_error, transaction -from superset.utils.ssh_tunnel import get_default_port - -logger = logging.getLogger(__name__) - - -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 - - @transaction(on_error=partial(on_error, reraise=SSHTunnelUpdateFailedError)) - def run(self) -> Optional[Model]: - self.validate() - - if self._model is None: - return None - - # unset password if private key is provided - if self._properties.get("private_key"): - self._properties["password"] = None - - # unset private key and password if password is provided - if self._properties.get("password"): - self._properties["private_key"] = None - self._properties["private_key_password"] = None - - return SSHTunnelDAO.update(self._model, self._properties) - - def validate(self) -> None: - # Validate/populate model exists - self._model = SSHTunnelDAO.find_by_id(self._model_id) - if not self._model: - raise SSHTunnelNotFoundError() - - url = make_url_safe(self._model.database.sqlalchemy_uri) - private_key: Optional[str] = self._properties.get("private_key") - private_key_password: Optional[str] = self._properties.get( - "private_key_password" - ) - if private_key_password and private_key is None: - raise SSHTunnelInvalidError( - exceptions=[SSHTunnelRequiredFieldValidationError("private_key")] - ) - backend = url.get_backend_name() - port = url.port or get_default_port(backend) - if not port: - raise SSHTunnelDatabasePortError() diff --git a/superset/commands/database/sync_permissions.py b/superset/commands/database/sync_permissions.py index a61d3ad80d2..03e0386ba0e 100644 --- a/superset/commands/database/sync_permissions.py +++ b/superset/commands/database/sync_permissions.py @@ -38,7 +38,6 @@ from superset.commands.database.utils import ( ) from superset.daos.database import DatabaseDAO from superset.daos.dataset import DatasetDAO -from superset.databases.ssh_tunnel.models import SSHTunnel from superset.db_engine_specs.base import GenericDBException from superset.exceptions import OAuth2RedirectError from superset.extensions import celery_app, db @@ -63,7 +62,6 @@ class SyncPermissionsCommand(BaseCommand): username: str | None, old_db_connection_name: str | None = None, db_connection: Database | None = None, - ssh_tunnel: SSHTunnel | None = None, ): """ Constructor method. @@ -72,7 +70,6 @@ class SyncPermissionsCommand(BaseCommand): self.username = username self._old_db_connection_name: str | None = old_db_connection_name self._db_connection: Database | None = db_connection - self.db_connection_ssh_tunnel: SSHTunnel | None = ssh_tunnel self.async_mode: bool = app.config["SYNC_DB_PERMISSIONS_IN_ASYNC_MODE"] @@ -99,20 +96,13 @@ class SyncPermissionsCommand(BaseCommand): if not self._db_connection: raise DatabaseNotFoundError() - if not self.db_connection_ssh_tunnel: - self.db_connection_ssh_tunnel = DatabaseDAO.get_ssh_tunnel( - self.db_connection_id - ) - # Need user info to impersonate for OAuth2 connections if not self.username or not security_manager.get_user_by_username( self.username ): raise UserNotFoundInSessionError() - with self.db_connection.get_sqla_engine( - override_ssh_tunnel=self.db_connection_ssh_tunnel - ) as engine: + with self.db_connection.get_sqla_engine() as engine: try: alive = ping(engine) except Exception as err: @@ -209,10 +199,7 @@ class SyncPermissionsCommand(BaseCommand): self.db_connection.db_engine_spec.supports_cross_catalog_queries or self.db_connection.allow_multi_catalog ): - return self.db_connection.get_all_catalog_names( - force=True, - ssh_tunnel=self.db_connection_ssh_tunnel, - ) + return self.db_connection.get_all_catalog_names(force=True) else: return {self.db_connection.get_default_catalog()} except OAuth2RedirectError: @@ -226,11 +213,7 @@ class SyncPermissionsCommand(BaseCommand): Helper method to load schemas. """ try: - return self.db_connection.get_all_schema_names( - force=True, - catalog=catalog, - ssh_tunnel=self.db_connection_ssh_tunnel, - ) + return self.db_connection.get_all_schema_names(force=True, catalog=catalog) except OAuth2RedirectError: # raise OAuth2 exceptions as-is raise @@ -335,14 +318,12 @@ def sync_database_permissions_task( db_connection = DatabaseDAO.find_by_id(database_id) if not db_connection: raise DatabaseNotFoundError() - ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database_id) SyncPermissionsCommand( database_id, username, old_db_connection_name=old_db_connection_name, db_connection=db_connection, - ssh_tunnel=ssh_tunnel, ).sync_database_permissions() logger.info( diff --git a/superset/commands/database/test_connection.py b/superset/commands/database/test_connection.py index 1bbb0013a47..83deb2f9c29 100644 --- a/superset/commands/database/test_connection.py +++ b/superset/commands/database/test_connection.py @@ -32,8 +32,7 @@ from superset.commands.database.ssh_tunnel.exceptions import ( SSHTunnelingNotEnabledError, ) from superset.commands.database.utils import ping -from superset.daos.database import DatabaseDAO, SSHTunnelDAO -from superset.databases.ssh_tunnel.models import SSHTunnel +from superset.daos.database import DatabaseDAO from superset.databases.utils import make_url_safe from superset.errors import ErrorLevel, SupersetErrorType from superset.exceptions import ( @@ -94,7 +93,9 @@ class TestConnectionDatabaseCommand(BaseCommand): ) -> None: # pylint: disable=too-many-statements,too-many-branches self.validate() ex_str = "" - ssh_tunnel = self._properties.get("ssh_tunnel") + + url = make_url_safe(self._uri) + engine = url.get_backend_name() serialized_encrypted_extra = self._properties.get( "masked_encrypted_extra", @@ -108,33 +109,37 @@ class TestConnectionDatabaseCommand(BaseCommand): ) ) + # collect SSH tunnel info + ssh_tunnel_properties = self._properties.get("ssh_tunnel") + if ssh_tunnel_properties and self._model and self._model.ssh_tunnel: + # unmask password while allowing for updated values + ssh_tunnel_properties = unmask_password_info( + ssh_tunnel_properties, + self._model.ssh_tunnel, + ) + + database: Database | None = None try: database = DatabaseDAO.build_db_for_connection_test( server_cert=self._properties.get("server_cert", ""), extra=self._properties.get("extra", "{}"), impersonate_user=self._properties.get("impersonate_user", False), encrypted_extra=serialized_encrypted_extra, + ssh_tunnel=ssh_tunnel_properties, ) database.set_sqlalchemy_uri(self._uri) database.db_engine_spec.mutate_db_for_connection_test(database) - # Generate tunnel if present in the properties - if ssh_tunnel: - # unmask password while allowing for updated values - if ssh_tunnel_id := ssh_tunnel.pop("id", None): - if existing_ssh_tunnel := SSHTunnelDAO.find_by_id(ssh_tunnel_id): - ssh_tunnel = unmask_password_info( - ssh_tunnel, existing_ssh_tunnel - ) - ssh_tunnel = SSHTunnel(**ssh_tunnel) - event_logger.log_with_context( - action=get_log_connection_action("test_connection_attempt", ssh_tunnel), - engine=database.db_engine_spec.__name__, + action=get_log_connection_action( + "test_connection_attempt", + ssh_tunnel_properties, + ), + engine=engine, ) - with database.get_sqla_engine(override_ssh_tunnel=ssh_tunnel) as engine: + with database.get_sqla_engine() as engine: try: alive = ping(engine) except SupersetTimeoutException as ex: @@ -165,29 +170,40 @@ class TestConnectionDatabaseCommand(BaseCommand): # Log succesful connection test with engine event_logger.log_with_context( - action=get_log_connection_action("test_connection_success", ssh_tunnel), - engine=database.db_engine_spec.__name__, + action=get_log_connection_action( + "test_connection_success", + ssh_tunnel_properties, + ), + engine=engine, ) except (NoSuchModuleError, ModuleNotFoundError) as ex: event_logger.log_with_context( action=get_log_connection_action( - "test_connection_error", ssh_tunnel, ex + "test_connection_error", + ssh_tunnel_properties, + ex, ), - engine=database.db_engine_spec.__name__, + engine=engine, ) raise DatabaseTestConnectionDriverError( - message=_("Could not load database driver: {}").format( - database.db_engine_spec.__name__ + message=_( + "Could not load database driver for: %(engine)s", + engine=engine, ), ) from ex except DBAPIError as ex: event_logger.log_with_context( action=get_log_connection_action( - "test_connection_error", ssh_tunnel, ex + "test_connection_error", + ssh_tunnel_properties, + ex, ), - engine=database.db_engine_spec.__name__, + engine=engine, ) + + if not database: + raise # check for custom errors (wrong username, wrong password, etc) errors = database.db_engine_spec.extract_errors( ex, self._context, database_name=database.unique_name @@ -198,30 +214,39 @@ class TestConnectionDatabaseCommand(BaseCommand): except SupersetSecurityException as ex: event_logger.log_with_context( action=get_log_connection_action( - "test_connection_error", ssh_tunnel, ex + "test_connection_error", + ssh_tunnel_properties, + ex, ), - engine=database.db_engine_spec.__name__, + engine=engine, ) raise DatabaseSecurityUnsafeError(message=str(ex)) from ex except (SupersetTimeoutException, SSHTunnelingNotEnabledError) as ex: event_logger.log_with_context( action=get_log_connection_action( - "test_connection_error", ssh_tunnel, ex + "test_connection_error", + ssh_tunnel_properties, + ex, ), - engine=database.db_engine_spec.__name__, + engine=engine, ) # bubble up the exception to return proper status code raise except Exception as ex: + if not database: + raise + if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2( ex ): database.start_oauth2_dance() event_logger.log_with_context( action=get_log_connection_action( - "test_connection_error", ssh_tunnel, ex + "test_connection_error", + ssh_tunnel_properties, + ex, ), - engine=database.db_engine_spec.__name__, + engine=engine, ) errors = database.db_engine_spec.extract_errors(ex, self._context) raise DatabaseTestConnectionUnexpectedError(errors) from ex diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py index 98e3ffac2b3..124bf11c87b 100644 --- a/superset/commands/database/update.py +++ b/superset/commands/database/update.py @@ -23,7 +23,7 @@ from typing import Any from flask_appbuilder.models.sqla import Model -from superset import db, is_feature_enabled +from superset import db from superset.commands.base import BaseCommand from superset.commands.database.exceptions import ( DatabaseExistsValidationError, @@ -32,15 +32,8 @@ from superset.commands.database.exceptions import ( DatabaseUpdateFailedError, MissingOAuth2TokenError, ) -from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand -from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelingNotEnabledError, -) -from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand from superset.commands.database.sync_permissions import SyncPermissionsCommand from superset.daos.database import DatabaseDAO -from superset.databases.ssh_tunnel.models import SSHTunnel from superset.exceptions import OAuth2RedirectError from superset.models.core import Database from superset.utils import json @@ -95,7 +88,7 @@ class UpdateDatabaseCommand(BaseCommand): # build new DB database = DatabaseDAO.update(self._model, self._properties) database.set_sqlalchemy_uri(database.sqlalchemy_uri) - ssh_tunnel = self._handle_ssh_tunnel(database) + new_catalog = database.get_default_catalog() # update assets when the database catalog changes, if the database was not @@ -118,7 +111,6 @@ class UpdateDatabaseCommand(BaseCommand): current_username, old_db_connection_name=original_database_name, db_connection=database, - ssh_tunnel=ssh_tunnel, ).run() except (OAuth2RedirectError, MissingOAuth2TokenError): pass @@ -158,32 +150,6 @@ class UpdateDatabaseCommand(BaseCommand): self._model.purge_oauth2_tokens() break - def _handle_ssh_tunnel(self, database: Database) -> SSHTunnel | None: - """ - Delete, create, or update an SSH tunnel. - """ - if "ssh_tunnel" not in self._properties: - return None - - if not is_feature_enabled("SSH_TUNNELING"): - raise SSHTunnelingNotEnabledError() - - current_ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id) - ssh_tunnel_properties = self._properties["ssh_tunnel"] - - if ssh_tunnel_properties is None: - if current_ssh_tunnel: - DeleteSSHTunnelCommand(current_ssh_tunnel.id).run() - return None - - if current_ssh_tunnel is None: - return CreateSSHTunnelCommand(database, ssh_tunnel_properties).run() - - return UpdateSSHTunnelCommand( - current_ssh_tunnel.id, - ssh_tunnel_properties, - ).run() - def _update_catalog_attribute( self, database_id: int, diff --git a/superset/commands/database/utils.py b/superset/commands/database/utils.py index 186bccc8a24..e850a9e89ee 100644 --- a/superset/commands/database/utils.py +++ b/superset/commands/database/utils.py @@ -30,7 +30,6 @@ from sqlalchemy.engine import Engine from sqlalchemy.orm import Session from superset import security_manager -from superset.databases.ssh_tunnel.models import SSHTunnel from superset.db_engine_specs.base import GenericDBException from superset.models.core import Database from superset.security.manager import SupersetSecurityManager @@ -51,7 +50,7 @@ def ping(engine: Engine) -> bool: return engine.dialect.do_ping(engine) -def add_permissions(database: Database, ssh_tunnel: SSHTunnel | None) -> None: +def add_permissions(database: Database) -> None: """ Add DAR for catalogs and schemas. """ @@ -66,10 +65,7 @@ def add_permissions(database: Database, ssh_tunnel: SSHTunnel | None) -> None: database.db_engine_spec.supports_cross_catalog_queries or database.allow_multi_catalog ): - catalogs = database.get_all_catalog_names( - cache=False, - ssh_tunnel=ssh_tunnel, - ) + catalogs = database.get_all_catalog_names(cache=False) else: catalogs = {database.get_default_catalog()} @@ -86,11 +82,7 @@ def add_permissions(database: Database, ssh_tunnel: SSHTunnel | None) -> None: for catalog in catalogs: try: - for schema in database.get_all_schema_names( - catalog=catalog, - cache=False, - ssh_tunnel=ssh_tunnel, - ): + for schema in database.get_all_schema_names(catalog=catalog, cache=False): security_manager.add_permission_view_menu( "schema_access", security_manager.get_schema_perm( diff --git a/superset/commands/dataset/export.py b/superset/commands/dataset/export.py index 5422f821577..6c866b310cd 100644 --- a/superset/commands/dataset/export.py +++ b/superset/commands/dataset/export.py @@ -24,7 +24,6 @@ import yaml from superset.commands.export.models import ExportModelsCommand from superset.connectors.sqla.models import SqlaTable -from superset.daos.database import DatabaseDAO from superset.commands.dataset.exceptions import DatasetNotFoundError from superset.daos.dataset import DatasetDAO from superset.utils.dict_import_export import EXPORT_VERSION @@ -111,7 +110,7 @@ class ExportDatasetsCommand(ExportModelsCommand): except json.JSONDecodeError: logger.info("Unable to decode `extra` field: %s", payload["extra"]) - if ssh_tunnel := DatabaseDAO.get_ssh_tunnel(model.database.id): + if ssh_tunnel := model.database.ssh_tunnel: ssh_tunnel_payload = ssh_tunnel.export_to_dict( recursive=False, include_parent_ref=False, diff --git a/superset/daos/database.py b/superset/daos/database.py index fa035534ee8..cd1bc3d51b3 100644 --- a/superset/daos/database.py +++ b/superset/daos/database.py @@ -19,6 +19,10 @@ from __future__ import annotations import logging from typing import Any +from sqlalchemy.orm import joinedload + +from superset import is_feature_enabled +from superset.commands.database.ssh_tunnel.exceptions import SSHTunnelingNotEnabledError from superset.connectors.sqla.models import SqlaTable from superset.daos.base import BaseDAO from superset.databases.filters import DatabaseFilter @@ -37,6 +41,51 @@ logger = logging.getLogger(__name__) class DatabaseDAO(BaseDAO[Database]): base_filter = DatabaseFilter + @classmethod + def create( + cls, + item: Database | None = None, + attributes: dict[str, Any] | None = None, + ) -> Database: + """ + Create a new database, with an optional SSH tunnel. + """ + ssh_tunnel_attributes = ( + attributes.pop("ssh_tunnel", None) if attributes else None + ) + + database = super().create(item, attributes) + + if ssh_tunnel_attributes: + if not is_feature_enabled("SSH_TUNNELING"): + raise SSHTunnelingNotEnabledError() + + database.ssh_tunnel = SSHTunnel(**ssh_tunnel_attributes) + + return database + + @classmethod + def find_by_id( + cls, + model_id: str | int, + skip_base_filter: bool = False, + id_column: str | None = None, + ) -> Database | None: + """ + Find a database by id, eagerly loading the SSH tunnel relationship. + """ + query = db.session.query(cls.model_cls).options(joinedload(Database.ssh_tunnel)) + query = cls._apply_base_filter(query, skip_base_filter) + + column_name = id_column or cls.id_column_name + if not hasattr(cls.model_cls, column_name): + raise AttributeError( + "{0} has no column {1}".format(cls.model_cls, column_name) + ) + + column = getattr(cls.model_cls, column_name) + return query.filter(column == model_id).one_or_none() + @classmethod def update( cls, @@ -52,14 +101,43 @@ class DatabaseDAO(BaseDAO[Database]): The masked values should be unmasked before the database is updated. """ # noqa: E501 + attributes = attributes or {} - if item and attributes and "encrypted_extra" in attributes: + if item and "encrypted_extra" in attributes: attributes["encrypted_extra"] = item.db_engine_spec.unmask_encrypted_extra( item.encrypted_extra, attributes["encrypted_extra"], ) - return super().update(item, attributes) + # update SSH tunnel + if "ssh_tunnel" not in attributes: + # keep existing tunnel if it exists + ssh_tunnel = item.ssh_tunnel if item else None + elif attributes["ssh_tunnel"] is None: + # remove existing tunnel + ssh_tunnel = None + else: + # update existing tunnel or create a new one + ssh_tunnel_attributes = attributes.pop("ssh_tunnel") + + # when updating the tunnel, passwords are sent masked to the frontend; if + # they arrive back masked we need to unmask them + if item and item.ssh_tunnel: + ssh_tunnel_attributes = unmask_password_info( + ssh_tunnel_attributes, + item.ssh_tunnel, + ) + + # delete existing SSH tunnel first + item.ssh_tunnel = None + db.session.flush() + + ssh_tunnel = SSHTunnel(**ssh_tunnel_attributes) + + database = super().update(item, attributes) + database.ssh_tunnel = ssh_tunnel + + return database @staticmethod def validate_uniqueness(database_name: str) -> bool: @@ -86,13 +164,18 @@ class DatabaseDAO(BaseDAO[Database]): @staticmethod def build_db_for_connection_test( - server_cert: str, extra: str, impersonate_user: bool, encrypted_extra: str + server_cert: str, + extra: str, + impersonate_user: bool, + encrypted_extra: str, + ssh_tunnel: dict[str, Any] | None = None, ) -> Database: return Database( server_cert=server_cert, extra=extra, impersonate_user=impersonate_user, encrypted_extra=encrypted_extra, + ssh_tunnel=SSHTunnel(**ssh_tunnel) if ssh_tunnel else None, ) @classmethod @@ -156,40 +239,6 @@ class DatabaseDAO(BaseDAO[Database]): .all() ) - @classmethod - def get_ssh_tunnel(cls, database_id: int) -> SSHTunnel | None: - ssh_tunnel = ( - db.session.query(SSHTunnel) - .filter(SSHTunnel.database_id == database_id) - .one_or_none() - ) - - return ssh_tunnel - - -class SSHTunnelDAO(BaseDAO[SSHTunnel]): - @classmethod - def update( - cls, - item: SSHTunnel | None = None, - attributes: dict[str, Any] | None = None, - ) -> SSHTunnel: - """ - Unmask ``password``, ``private_key`` and ``private_key_password`` before updating. - - When a database is edited the user sees a masked version of - the aforementioned fields. - - The masked values should be unmasked before the ssh tunnel is updated. - """ # noqa: E501 - # ID cannot be updated so we remove it if present in the payload - - if item and attributes: - attributes.pop("id", None) - attributes = unmask_password_info(attributes, item) - - return super().update(item, attributes) - class DatabaseUserOAuth2TokensDAO(BaseDAO[DatabaseUserOAuth2Tokens]): """ diff --git a/superset/databases/api.py b/superset/databases/api.py index c94060b2838..b0f29c5247a 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -54,10 +54,8 @@ from superset.commands.database.exceptions import ( from superset.commands.database.export import ExportDatabasesCommand from superset.commands.database.importers.dispatcher import ImportDatabasesCommand from superset.commands.database.oauth2 import OAuth2StoreTokenCommand -from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand from superset.commands.database.ssh_tunnel.exceptions import ( SSHTunnelDatabasePortError, - SSHTunnelDeleteFailedError, SSHTunnelingNotEnabledError, ) from superset.commands.database.sync_permissions import SyncPermissionsCommand @@ -356,8 +354,8 @@ class DatabaseRestApi(BaseSupersetModelRestApi): "result": database_connection_schema.dump(database, many=False), } try: - if ssh_tunnel := DatabaseDAO.get_ssh_tunnel(pk): - response["result"]["ssh_tunnel"] = ssh_tunnel.data + if database and database.ssh_tunnel: + response["result"]["ssh_tunnel"] = database.ssh_tunnel.data return self.response(200, **response) except SupersetException as ex: return self.response(ex.status, message=ex.message) @@ -394,9 +392,10 @@ class DatabaseRestApi(BaseSupersetModelRestApi): """ data = self.get_headless(pk, **kwargs) try: - if ssh_tunnel := DatabaseDAO.get_ssh_tunnel(pk): + database = DatabaseDAO.find_by_id(pk) + if database and database.ssh_tunnel: payload = data.json - payload["result"]["ssh_tunnel"] = ssh_tunnel.data + payload["result"]["ssh_tunnel"] = database.ssh_tunnel.data return payload return data except SupersetException as ex: @@ -447,7 +446,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi): item = self.add_model_schema.load(request.json) # This validates custom Schema with custom validations except ValidationError as error: - return self.response_400(message=error.messages) + return self.response_422(message=error.messages) try: new_model = CreateDatabaseCommand(item).run() item["uuid"] = new_model.uuid @@ -464,7 +463,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi): # Return SSH Tunnel and hide passwords if any if item.get("ssh_tunnel"): - item["ssh_tunnel"] = mask_password_info(new_model.ssh_tunnel) + item["ssh_tunnel"] = mask_password_info(item["ssh_tunnel"]) return self.response(201, id=new_model.id, result=item) except OAuth2RedirectError: @@ -552,7 +551,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi): item["parameters"] = changed_model.parameters # Return SSH Tunnel and hide passwords if any if item.get("ssh_tunnel"): - item["ssh_tunnel"] = mask_password_info(changed_model.ssh_tunnel) + item["ssh_tunnel"] = mask_password_info(item["ssh_tunnel"]) return self.response(200, id=changed_model.id, result=item) except DatabaseNotFoundError: return self.response_404() @@ -1991,73 +1990,6 @@ class DatabaseRestApi(BaseSupersetModelRestApi): command.run() return self.response(200, message="OK") - @expose("//ssh_tunnel/", methods=("DELETE",)) - @protect() - @statsd_metrics - @deprecated(deprecated_in="4.0") - @event_logger.log_this_with_context( - action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" - f".delete_ssh_tunnel", - log_to_statsd=False, - ) - def delete_ssh_tunnel(self, pk: int) -> Response: - """Delete a SSH tunnel. - --- - delete: - summary: Delete a SSH tunnel - parameters: - - in: path - schema: - type: integer - name: pk - responses: - 200: - description: SSH Tunnel deleted - content: - application/json: - schema: - type: object - properties: - message: - type: string - 401: - $ref: '#/components/responses/401' - 403: - $ref: '#/components/responses/403' - 404: - $ref: '#/components/responses/404' - 422: - $ref: '#/components/responses/422' - 500: - $ref: '#/components/responses/500' - """ - - database = DatabaseDAO.find_by_id(pk) - if not database: - return self.response_404() - try: - existing_ssh_tunnel_model = database.ssh_tunnels - if existing_ssh_tunnel_model: - DeleteSSHTunnelCommand(existing_ssh_tunnel_model.id).run() - return self.response(200, message="OK") - return self.response_404() - except SSHTunnelDeleteFailedError as ex: - logger.error( - "Error deleting SSH Tunnel %s: %s", - self.__class__.__name__, - str(ex), - exc_info=True, - ) - return self.response_422(message=str(ex)) - except SSHTunnelingNotEnabledError as ex: - logger.error( - "Error deleting SSH Tunnel %s: %s", - self.__class__.__name__, - str(ex), - exc_info=True, - ) - return self.response_400(message=str(ex)) - @expose("//schemas_access_for_file_upload/") @protect() @safe diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index ea24ba219b6..105496efa4c 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -450,6 +450,24 @@ class DatabaseSSHTunnel(Schema): private_key = fields.String(required=False) private_key_password = fields.String(required=False) + @validates_schema + def validate_authentication(self, data: dict[str, Any], **kwargs: Any) -> None: + errors: dict[str, str] = {} + + private_key: str | None = data.get("private_key") + password: str | None = data.get("password") + private_key_password: str | None = data.get("private_key_password") + + if not private_key and not password: + errors["password"] = "Either password or private_key is required" # noqa: S105 + if private_key_password and private_key is None: + errors["private_key"] = ( + "private_key is required when private_key_password is provided" + ) + + if errors: + raise ValidationError(errors) + class DatabasePostSchema(DatabaseParametersSchemaMixin, Schema): class Meta: # pylint: disable=too-few-public-methods diff --git a/superset/databases/ssh_tunnel/models.py b/superset/databases/ssh_tunnel/models.py index 4a0e624b521..5bd707d2a7b 100644 --- a/superset/databases/ssh_tunnel/models.py +++ b/superset/databases/ssh_tunnel/models.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from typing import Any import sqlalchemy as sa @@ -41,11 +43,19 @@ class SSHTunnel(AuditMixinNullable, ExtraJSONMixin, ImportExportMixin, Model): id = sa.Column(sa.Integer, primary_key=True) database_id = sa.Column( - sa.Integer, sa.ForeignKey("dbs.id"), nullable=False, unique=True + sa.Integer, + sa.ForeignKey("dbs.id"), + nullable=False, + unique=True, ) database: Database = relationship( "Database", - backref=backref("ssh_tunnels", uselist=False, cascade="all, delete-orphan"), + backref=backref( + "ssh_tunnel", + uselist=False, + cascade="all, delete-orphan", + lazy="joined", + ), foreign_keys=[database_id], ) diff --git a/superset/extensions/ssh.py b/superset/extensions/ssh.py index a3c015b9ebc..74fb44cfd75 100644 --- a/superset/extensions/ssh.py +++ b/superset/extensions/ssh.py @@ -23,6 +23,7 @@ import sshtunnel from flask import Flask from paramiko import RSAKey +from superset.commands.database.ssh_tunnel.exceptions import SSHTunnelDatabasePortError from superset.databases.utils import make_url_safe from superset.utils.class_utils import load_class_from_name @@ -56,10 +57,13 @@ class SSHManager: url = make_url_safe(sqlalchemy_database_uri) backend = url.get_backend_name() + port = url.port or get_default_port(backend) + if not port: + raise SSHTunnelDatabasePortError() params = { "ssh_address_or_host": (ssh_tunnel.server_address, ssh_tunnel.server_port), "ssh_username": ssh_tunnel.username, - "remote_bind_address": (url.host, url.port or get_default_port(backend)), + "remote_bind_address": (url.host, port), "local_bind_address": (self.local_bind_address,), "debug_level": logging.getLogger("flask_appbuilder").level, } diff --git a/superset/models/core.py b/superset/models/core.py index f6643d18ff2..4c5b4929cfc 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -94,7 +94,6 @@ metadata = Model.metadata # pylint: disable=no-member logger = logging.getLogger(__name__) if TYPE_CHECKING: - from superset.databases.ssh_tunnel.models import SSHTunnel from superset.models.sql_lab import Query @@ -202,6 +201,7 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable "external_url", "encrypted_extra", "impersonate_user", + "ssh_tunnel", ] export_children = ["tables"] @@ -426,7 +426,6 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable schema: str | None = None, nullpool: bool = True, source: utils.QuerySource | None = None, - override_ssh_tunnel: SSHTunnel | None = None, ) -> Engine: """ Context manager for a SQLAlchemy engine. @@ -436,19 +435,15 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable to potentially establish SSH tunnels before the connection is created, and clean them up once the engine is no longer used. """ - from superset.daos.database import ( # pylint: disable=import-outside-toplevel - DatabaseDAO, - ) sqlalchemy_uri = self.sqlalchemy_uri_decrypted - ssh_tunnel = override_ssh_tunnel or DatabaseDAO.get_ssh_tunnel(self.id) ssh_context_manager = ( ssh_manager_factory.instance.create_tunnel( - ssh_tunnel=ssh_tunnel, + ssh_tunnel=self.ssh_tunnel, sqlalchemy_database_uri=sqlalchemy_uri, ) - if ssh_tunnel + if self.ssh_tunnel else nullcontext() ) @@ -977,37 +972,23 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable self, catalog: str | None = None, schema: str | None = None, - ssh_tunnel: SSHTunnel | None = None, ) -> Inspector: - with self.get_sqla_engine( - catalog=catalog, - schema=schema, - override_ssh_tunnel=ssh_tunnel, - ) as engine: + with self.get_sqla_engine(catalog=catalog, schema=schema) as engine: yield sqla.inspect(engine) @cache_util.memoized_func( key="db:{self.id}:catalog:{catalog}:schema_list", cache=cache_manager.cache, ) - def get_all_schema_names( - self, - *, - catalog: str | None = None, - ssh_tunnel: SSHTunnel | None = None, - ) -> set[str]: + def get_all_schema_names(self, *, catalog: str | None = None) -> set[str]: """ Return the schemas in a given database :param catalog: override default catalog - :param ssh_tunnel: SSH tunnel information needed to establish a connection :return: schema list """ try: - with self.get_inspector( - catalog=catalog, - ssh_tunnel=ssh_tunnel, - ) as inspector: + with self.get_inspector(catalog=catalog) as inspector: return self.db_engine_spec.get_schema_names(inspector) except Exception as ex: if self.is_oauth2_enabled() and self.db_engine_spec.needs_oauth2(ex): @@ -1019,19 +1000,14 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable key="db:{self.id}:catalog_list", cache=cache_manager.cache, ) - def get_all_catalog_names( - self, - *, - ssh_tunnel: SSHTunnel | None = None, - ) -> set[str]: + def get_all_catalog_names(self) -> set[str]: """ Return the catalogs in a given database - :param ssh_tunnel: SSH tunnel information needed to establish a connection :return: catalog list """ try: - with self.get_inspector(ssh_tunnel=ssh_tunnel) as inspector: + with self.get_inspector() as inspector: return self.db_engine_spec.get_catalog_names(self, inspector) except Exception as ex: if self.is_oauth2_enabled() and self.db_engine_spec.needs_oauth2(ex): diff --git a/superset/utils/ssh_tunnel.py b/superset/utils/ssh_tunnel.py index 1471d54f4b1..7ae52b64386 100644 --- a/superset/utils/ssh_tunnel.py +++ b/superset/utils/ssh_tunnel.py @@ -29,24 +29,21 @@ DEFAULT_PORTS: dict[str, int] = { def mask_password_info(ssh_tunnel: dict[str, Any]) -> dict[str, Any]: - if ssh_tunnel.pop("password", None) is not None: - ssh_tunnel["password"] = PASSWORD_MASK - if ssh_tunnel.pop("private_key", None) is not None: - ssh_tunnel["private_key"] = PASSWORD_MASK - if ssh_tunnel.pop("private_key_password", None) is not None: - ssh_tunnel["private_key_password"] = PASSWORD_MASK + for key in {"password", "private_key", "private_key_password"}: + if ssh_tunnel.pop(key, None) is not None: + ssh_tunnel[key] = PASSWORD_MASK + return ssh_tunnel def unmask_password_info( - ssh_tunnel: dict[str, Any], model: SSHTunnel + ssh_tunnel: dict[str, Any], + model: SSHTunnel, ) -> dict[str, Any]: - if ssh_tunnel.get("password") == PASSWORD_MASK: - ssh_tunnel["password"] = model.password - if ssh_tunnel.get("private_key") == PASSWORD_MASK: - ssh_tunnel["private_key"] = model.private_key - if ssh_tunnel.get("private_key_password") == PASSWORD_MASK: - ssh_tunnel["private_key_password"] = model.private_key_password + for key in {"password", "private_key", "private_key_password"}: + if ssh_tunnel.get(key) == PASSWORD_MASK: + ssh_tunnel[key] = getattr(model, key) + return ssh_tunnel diff --git a/tests/integration_tests/commands_test.py b/tests/integration_tests/commands_test.py index d3cc01b4f62..41776fb9162 100644 --- a/tests/integration_tests/commands_test.py +++ b/tests/integration_tests/commands_test.py @@ -169,7 +169,7 @@ class TestImportAssetsCommand(SupersetTestCase): assert dashboard.owners == [self.user] - mock_add_permissions.assert_called_with(database, None) + mock_add_permissions.assert_called_with(database) db.session.delete(dashboard) db.session.delete(chart) @@ -214,7 +214,7 @@ class TestImportAssetsCommand(SupersetTestCase): dataset = chart.table database = dataset.database - mock_add_permissions.assert_called_with(database, None) + mock_add_permissions.assert_called_with(database) db.session.delete(dashboard) db.session.delete(chart) diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index a67d81a12cd..a08307d0460 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -289,23 +289,21 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_create_database_with_ssh_tunnel( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, ): """ Database API: Test create with SSH Tunnel """ - mock_create_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -337,23 +335,21 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_create_database_with_ssh_tunnel_no_port( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, ): """ Database API: Test create with SSH Tunnel """ - mock_create_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -390,23 +386,21 @@ class TestDatabaseApi(SupersetTestCase): db.session.commit() @pytest.mark.skip("buggy") + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_create_database_with_ssh_tunnel_no_port_no_default( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, ): """ Database API: Test that missing port raises SSHTunnelDatabaseError """ - mock_create_is_feature_enabled.return_value = True self.login(username="admin") example_db = get_example_database() if example_db.backend == "sqlite": @@ -435,30 +429,25 @@ class TestDatabaseApi(SupersetTestCase): == "A database port is required when connecting via SSH Tunnel." ) + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.sync_permissions.SyncPermissionsCommand.run", ) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") - @mock.patch("superset.commands.database.update.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_update_database_with_ssh_tunnel( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_update_is_feature_enabled, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, mock_sync_perms_command, ): """ Database API: Test update Database with SSH Tunnel """ - mock_create_is_feature_enabled.return_value = True - mock_update_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -500,30 +489,25 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.sync_permissions.SyncPermissionsCommand.run", ) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") - @mock.patch("superset.commands.database.update.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_update_database_with_ssh_tunnel_no_port( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_update_is_feature_enabled, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, mock_sync_perms_cmmd_run, ): """ Database API: Test update Database with SSH Tunnel """ - mock_create_is_feature_enabled.return_value = True - mock_update_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -568,26 +552,21 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") - @mock.patch("superset.commands.database.update.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_update_database_no_port_no_default( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_update_is_feature_enabled, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, ): """ Database API: Test that missing port raises SSHTunnelDatabaseError """ - mock_create_is_feature_enabled.return_value = True - mock_update_is_feature_enabled.return_value = True self.login(username="admin") example_db = get_example_database() if example_db.backend == "sqlite": @@ -632,33 +611,25 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.sync_permissions.SyncPermissionsCommand.run", ) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") - @mock.patch("superset.commands.database.update.is_feature_enabled") - @mock.patch("superset.commands.database.ssh_tunnel.delete.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_delete_ssh_tunnel( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_delete_is_feature_enabled, - mock_update_is_feature_enabled, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, mock_sync_perms_command, ): """ Database API: Test deleting a SSH tunnel via Database update """ - mock_create_is_feature_enabled.return_value = True - mock_update_is_feature_enabled.return_value = True - mock_delete_is_feature_enabled.return_value = True self.login(username="admin") example_db = get_example_database() if example_db.backend == "sqlite": @@ -720,30 +691,25 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.sync_permissions.SyncPermissionsCommand.run", ) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") - @mock.patch("superset.commands.database.update.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_update_ssh_tunnel_via_database_api( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_update_is_feature_enabled, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, mock_sync_perms_command, ): """ Database API: Test update SSH Tunnel via Database API """ - mock_create_is_feature_enabled.return_value = True - mock_update_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() @@ -802,15 +768,14 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") - @mock.patch("superset.commands.database.create.is_feature_enabled") def test_cascade_delete_ssh_tunnel( self, - mock_create_is_feature_enabled, mock_get_all_schema_names, mock_get_all_catalog_names, mock_test_connection_database_command_run, @@ -818,7 +783,6 @@ class TestDatabaseApi(SupersetTestCase): """ Database API: SSH Tunnel gets deleted if Database gets deleted """ - mock_create_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -837,6 +801,7 @@ class TestDatabaseApi(SupersetTestCase): uri = "api/v1/database/" rv = self.client.post(uri, json=database_data) response = json.loads(rv.data.decode("utf-8")) + print(rv.text) assert rv.status_code == 201 model_ssh_tunnel = ( db.session.query(SSHTunnel) @@ -855,10 +820,10 @@ class TestDatabaseApi(SupersetTestCase): ) assert model_ssh_tunnel is None + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") @mock.patch("superset.extensions.db.session.rollback") @@ -866,14 +831,12 @@ class TestDatabaseApi(SupersetTestCase): self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, mock_rollback, ): """ Database API: Test rollback is called if SSH Tunnel creation fails """ - mock_create_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -897,7 +860,11 @@ class TestDatabaseApi(SupersetTestCase): "sqlalchemy_uri": example_db.sqlalchemy_uri_decrypted, "ssh_tunnel": ssh_tunnel_properties, } - fail_message = {"message": "SSH Tunnel parameters are invalid."} + fail_message = { + "message": { + "ssh_tunnel": {"password": "Either password or private_key is required"} + } + } uri = "api/v1/database/" rv = self.client.post(uri, json=database_data) @@ -912,9 +879,6 @@ class TestDatabaseApi(SupersetTestCase): assert model_ssh_tunnel is None assert response == fail_message - # Check that rollback was called - mock_rollback.assert_called() - # Clean up any database that might have been created created_db = ( db.session.query(Database) @@ -925,23 +889,21 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(created_db) db.session.commit() + @with_feature_flags(SSH_TUNNELING=True) @mock.patch( "superset.commands.database.test_connection.TestConnectionDatabaseCommand.run", ) - @mock.patch("superset.commands.database.create.is_feature_enabled") @mock.patch("superset.models.core.Database.get_all_catalog_names") @mock.patch("superset.models.core.Database.get_all_schema_names") def test_get_database_returns_related_ssh_tunnel( self, mock_get_all_schema_names, mock_get_all_catalog_names, - mock_create_is_feature_enabled, mock_test_connection_database_command_run, ): """ Database API: Test GET Database returns its related SSH Tunnel """ - mock_create_is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) example_db = get_example_database() if example_db.backend == "sqlite": @@ -1091,7 +1053,7 @@ class TestDatabaseApi(SupersetTestCase): ] } } - assert rv.status_code == 400 + assert rv.status_code == 422 def test_create_database_no_configuration_method(self): """ @@ -1146,7 +1108,7 @@ class TestDatabaseApi(SupersetTestCase): rv = self.client.post(uri, json=database_data) response = json.loads(rv.data.decode("utf-8")) expected_response = {"message": {"server_cert": ["Invalid certificate"]}} - assert rv.status_code == 400 + assert rv.status_code == 422 assert response == expected_response def test_create_database_json_validate(self): @@ -1181,7 +1143,7 @@ class TestDatabaseApi(SupersetTestCase): ], } } - assert rv.status_code == 400 + assert rv.status_code == 422 assert response == expected_response def test_create_database_extra_metadata_validate(self): @@ -1217,7 +1179,7 @@ class TestDatabaseApi(SupersetTestCase): ] } } - assert rv.status_code == 400 + assert rv.status_code == 422 assert response == expected_response def test_create_database_unique_validate(self): @@ -1260,7 +1222,7 @@ class TestDatabaseApi(SupersetTestCase): uri = "api/v1/database/" rv = self.client.post(uri, json=database_data) response = json.loads(rv.data.decode("utf-8")) - assert rv.status_code == 400 + assert rv.status_code == 422 assert "Invalid connection string" in response["message"]["sqlalchemy_uri"][0] @with_config({"PREVENT_UNSAFE_DB_CONNECTIONS": True}) @@ -1287,7 +1249,7 @@ class TestDatabaseApi(SupersetTestCase): } } assert response_data == expected_response - assert response.status_code == 400 + assert response.status_code == 422 def test_create_database_conn_fail(self): """ @@ -2343,7 +2305,7 @@ class TestDatabaseApi(SupersetTestCase): expected_response = { "errors": [ { - "message": "Could not load database driver: BaseEngineSpec", + "message": "Could not load database driver for: broken", "error_type": "GENERIC_COMMAND_ERROR", "level": "warning", "extra": { @@ -2372,7 +2334,7 @@ class TestDatabaseApi(SupersetTestCase): expected_response = { "errors": [ { - "message": "Could not load database driver: MssqlEngineSpec", + "message": "Could not load database driver for: mssql", "error_type": "GENERIC_COMMAND_ERROR", "level": "warning", "extra": { @@ -2991,16 +2953,27 @@ class TestDatabaseApi(SupersetTestCase): assert response == { "errors": [ { - "message": "Must provide credentials for the SSH Tunnel", + "message": ( + "Error importing database: databases/database_1.yaml: " + "{'ssh_tunnel': {'password': 'Either password or private_key " + "is required'}}" + ), "error_type": "GENERIC_COMMAND_ERROR", "level": "warning", "extra": { + "databases/database_1.yaml": { + "ssh_tunnel": { + "password": ( + "Either password or private_key is required" + ), + } + }, "issue_codes": [ { "code": 1010, "message": ( - "Issue 1010 - Superset encountered an " - "error while running a command." + "Issue 1010 - Superset encountered an error while " + "running a command." ), } ], diff --git a/tests/integration_tests/databases/commands_tests.py b/tests/integration_tests/databases/commands_tests.py index 4445d4ea6b1..529e4807e70 100644 --- a/tests/integration_tests/databases/commands_tests.py +++ b/tests/integration_tests/databases/commands_tests.py @@ -813,7 +813,10 @@ class TestImportDatabasesCommand(SupersetTestCase): command = ImportDatabasesCommand(contents) with pytest.raises(CommandInvalidError) as excinfo: command.run() - assert str(excinfo.value) == "Must provide credentials for the SSH Tunnel" + assert str(excinfo.value) == ( + "Error importing database: databases/imported_database.yaml: " + "{'ssh_tunnel': {'password': 'Either password or private_key is required'}}" + ) @patch("superset.databases.schemas.is_feature_enabled") @patch("superset.commands.database.importers.v1.utils.add_permissions") @@ -858,10 +861,12 @@ class TestImportDatabasesCommand(SupersetTestCase): assert str(excinfo.value).startswith("Error importing database") assert excinfo.value.normalized_messages() == { "databases/imported_database.yaml": { - "_schema": [ - "Must provide a private key for the ssh tunnel", - "Must provide a private key password for the ssh tunnel", - ] + "ssh_tunnel": { + "password": "Either password or private_key is required", + "private_key": ( + "private_key is required when private_key_password is provided" + ), + } } } diff --git a/tests/integration_tests/databases/ssh_tunnel/commands/__init__.py b/tests/integration_tests/databases/ssh_tunnel/commands/__init__.py deleted file mode 100644 index 13a83393a91..00000000000 --- a/tests/integration_tests/databases/ssh_tunnel/commands/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/tests/integration_tests/databases/ssh_tunnel/commands/commands_tests.py b/tests/integration_tests/databases/ssh_tunnel/commands/commands_tests.py deleted file mode 100644 index 9e1f33f3933..00000000000 --- a/tests/integration_tests/databases/ssh_tunnel/commands/commands_tests.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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 unittest import mock - -import pytest - -from superset import security_manager -from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelNotFoundError, -) -from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand -from tests.integration_tests.base_tests import SupersetTestCase - - -class TestUpdateSSHTunnelCommand(SupersetTestCase): - @mock.patch("superset.utils.core.g") - def test_update_ssh_tunnel_not_found(self, mock_g): - mock_g.user = security_manager.find_user("admin") - # We have not created a SSH Tunnel yet so id = 1 is invalid - command = UpdateSSHTunnelCommand( - 1, - { - "server_address": "127.0.0.1", - "server_port": 5432, - "username": "test_user", - }, - ) - with pytest.raises(SSHTunnelNotFoundError) as excinfo: - command.run() - assert str(excinfo.value) == ("SSH Tunnel not found.") - - -class TestDeleteSSHTunnelCommand(SupersetTestCase): - @mock.patch("superset.utils.core.g") - @mock.patch("superset.commands.database.ssh_tunnel.delete.is_feature_enabled") - def test_delete_ssh_tunnel_not_found(self, mock_g, mock_delete_is_feature_enabled): - mock_g.user = security_manager.find_user("admin") - mock_delete_is_feature_enabled.return_value = True - # We have not created a SSH Tunnel yet so id = 1 is invalid - command = DeleteSSHTunnelCommand(1) - with pytest.raises(SSHTunnelNotFoundError) as excinfo: - command.run() - assert str(excinfo.value) == ("SSH Tunnel not found.") diff --git a/tests/integration_tests/tags/api_tests.py b/tests/integration_tests/tags/api_tests.py index 8c0d90e26fa..ba4ce8aba6c 100644 --- a/tests/integration_tests/tags/api_tests.py +++ b/tests/integration_tests/tags/api_tests.py @@ -850,9 +850,10 @@ class TestTagApi(InsertChartMixin, SupersetTestCase): rv = self.client.post(uri, json=data) # Should succeed without SQL errors (201 for created or 200 for success) - assert rv.status_code in [200, 201], ( - f"Tag creation should succeed, got {rv.status_code}" - ) + assert rv.status_code in [ + 200, + 201, + ], f"Tag creation should succeed, got {rv.status_code}" # Query the database to verify the tag was created correctly created_tag = db.session.query(Tag).filter_by(name=tag_name).first() diff --git a/tests/unit_tests/commands/databases/sync_permissions_test.py b/tests/unit_tests/commands/databases/sync_permissions_test.py index 78dfe3d0c15..55a6e4501db 100644 --- a/tests/unit_tests/commands/databases/sync_permissions_test.py +++ b/tests/unit_tests/commands/databases/sync_permissions_test.py @@ -43,7 +43,6 @@ def test_sync_permissions_command_sync_mode( """ Test ``SyncPermissionsCommand`` in sync mode. """ - mock_ssh = mocker.MagicMock() user_mock = mocker.patch( "superset.commands.database.sync_permissions.security_manager.get_user_by_username" ) @@ -55,7 +54,9 @@ def test_sync_permissions_command_sync_mode( add_pvm_mock = mocker.patch("superset.commands.database.sync_permissions.add_pvm") cmmd = SyncPermissionsCommand( - 1, "admin", db_connection=database_with_catalog, ssh_tunnel=mock_ssh + 1, + "admin", + db_connection=database_with_catalog, ) mock_refresh_schemas = mocker.patch.object(cmmd, "_refresh_schemas") mock_rename_db_perm = mocker.patch.object(cmmd, "_rename_database_in_permissions") @@ -64,7 +65,6 @@ def test_sync_permissions_command_sync_mode( assert cmmd.db_connection == database_with_catalog assert cmmd.old_db_connection_name == "my_db" - assert cmmd.db_connection_ssh_tunnel == mock_ssh user_mock.assert_called_once_with("admin") add_pvm_mock.assert_has_calls( [ @@ -120,7 +120,6 @@ def test_sync_permissions_command_passing_all_values( """ Test ``SyncPermissionsCommand`` when providing all arguments to the constructor. """ - mock_ssh = mocker.MagicMock() mock_database_dao = mocker.patch( "superset.commands.database.sync_permissions.DatabaseDAO" ) @@ -134,16 +133,13 @@ def test_sync_permissions_command_passing_all_values( "admin", old_db_connection_name="old name", db_connection=database_with_catalog, - ssh_tunnel=mock_ssh, ) mocker.patch.object(cmmd, "sync_database_permissions") cmmd.run() assert cmmd.db_connection == database_with_catalog assert cmmd.old_db_connection_name == "old name" - assert cmmd.db_connection_ssh_tunnel == mock_ssh mock_database_dao.find_by_id.assert_not_called() - mock_database_dao.get_ssh_tunnel.assert_not_called() @with_config({"SYNC_DB_PERMISSIONS_IN_ASYNC_MODE": False}) @@ -159,7 +155,6 @@ def test_sync_permissions_command_raise( "superset.commands.database.sync_permissions.DatabaseDAO" ) mock_database_dao.find_by_id.return_value = database_without_catalog - mock_database_dao.get_ssh_tunnel.return_value = mocker.MagicMock() mock_user = mocker.patch( "superset.commands.database.sync_permissions.security_manager.get_user_by_username" ) diff --git a/tests/unit_tests/databases/api_test.py b/tests/unit_tests/databases/api_test.py index c1d94c1b390..0785d762e50 100644 --- a/tests/unit_tests/databases/api_test.py +++ b/tests/unit_tests/databases/api_test.py @@ -308,6 +308,7 @@ def test_database_connection( }, "server_cert": None, "sqlalchemy_uri": "gsheets://", + "ssh_tunnel": None, "uuid": "02feae18-2dd6-4bb4-a9c0-49e9d4f29d58", }, } @@ -486,160 +487,6 @@ def test_non_zip_import(client: Any, full_api_access: None) -> None: } -def test_delete_ssh_tunnel( - mocker: MockerFixture, - app: Any, - session: Session, - client: Any, - full_api_access: None, -) -> None: - """ - Test that we can delete SSH Tunnel - """ - with app.app_context(): - from superset.daos.database import DatabaseDAO - from superset.databases.api import DatabaseRestApi - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - DatabaseRestApi.datamodel._session = session - - # create table for databases - Database.metadata.create_all(session.get_bind()) # pylint: disable=no-member - - # Create our Database - database = Database( - database_name="my_database", - sqlalchemy_uri="gsheets://", - encrypted_extra=json.dumps( - { - "service_account_info": { - "type": "service_account", - "project_id": "black-sanctum-314419", - "private_key_id": "259b0d419a8f840056158763ff54d8b08f7b8173", - "private_key": "SECRET", - "client_email": "google-spreadsheets-demo-servi@black-sanctum-314419.iam.gserviceaccount.com", # noqa: E501 - "client_id": "SSH_TUNNEL_CREDENTIALS_CLIENT", - "auth_uri": "https://accounts.google.com/o/oauth2/auth", - "token_uri": "https://oauth2.googleapis.com/token", - "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/google-spreadsheets-demo-servi%40black-sanctum-314419.iam.gserviceaccount.com", - }, - } - ), - ) - db.session.add(database) - db.session.commit() - - # mock the lookup so that we don't need to include the driver - mocker.patch("sqlalchemy.engine.URL.get_driver_name", return_value="gsheets") - mocker.patch("superset.utils.log.DBEventLogger.log") - mocker.patch( - "superset.commands.database.ssh_tunnel.delete.is_feature_enabled", - return_value=True, - ) - - # Create our SSHTunnel - tunnel = SSHTunnel( - database_id=1, - database=database, - ) - - db.session.add(tunnel) - db.session.commit() - - # Get our recently created SSHTunnel - response_tunnel = DatabaseDAO.get_ssh_tunnel(1) - assert response_tunnel - assert isinstance(response_tunnel, SSHTunnel) - assert 1 == response_tunnel.database_id - - # Delete the recently created SSHTunnel - response_delete_tunnel = client.delete( - f"/api/v1/database/{database.id}/ssh_tunnel/" - ) - assert response_delete_tunnel.json["message"] == "OK" - - response_tunnel = DatabaseDAO.get_ssh_tunnel(1) - assert response_tunnel is None - - -def test_delete_ssh_tunnel_not_found( - mocker: MockerFixture, - app: Any, - session: Session, - client: Any, - full_api_access: None, -) -> None: - """ - Test that we cannot delete a tunnel that does not exist - """ - with app.app_context(): - from superset.daos.database import DatabaseDAO - from superset.databases.api import DatabaseRestApi - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - DatabaseRestApi.datamodel._session = session - - # create table for databases - Database.metadata.create_all(session.get_bind()) # pylint: disable=no-member - - # Create our Database - database = Database( - database_name="my_database", - sqlalchemy_uri="gsheets://", - encrypted_extra=json.dumps( - { - "service_account_info": { - "type": "service_account", - "project_id": "black-sanctum-314419", - "private_key_id": "259b0d419a8f840056158763ff54d8b08f7b8173", - "private_key": "SECRET", - "client_email": "google-spreadsheets-demo-servi@black-sanctum-314419.iam.gserviceaccount.com", # noqa: E501 - "client_id": "SSH_TUNNEL_CREDENTIALS_CLIENT", - "auth_uri": "https://accounts.google.com/o/oauth2/auth", - "token_uri": "https://oauth2.googleapis.com/token", - "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/google-spreadsheets-demo-servi%40black-sanctum-314419.iam.gserviceaccount.com", - }, - } - ), - ) - db.session.add(database) - db.session.commit() - - # mock the lookup so that we don't need to include the driver - mocker.patch("sqlalchemy.engine.URL.get_driver_name", return_value="gsheets") - mocker.patch("superset.utils.log.DBEventLogger.log") - mocker.patch( - "superset.commands.database.ssh_tunnel.delete.is_feature_enabled", - return_value=True, - ) - - # Create our SSHTunnel - tunnel = SSHTunnel( - database_id=1, - database=database, - ) - - db.session.add(tunnel) - db.session.commit() - - # Delete the recently created SSHTunnel - response_delete_tunnel = client.delete("/api/v1/database/2/ssh_tunnel/") - assert response_delete_tunnel.json["message"] == "Not found" - - # Get our recently created SSHTunnel - response_tunnel = DatabaseDAO.get_ssh_tunnel(1) - assert response_tunnel - assert isinstance(response_tunnel, SSHTunnel) - assert 1 == response_tunnel.database_id - - response_tunnel = DatabaseDAO.get_ssh_tunnel(2) - assert response_tunnel is None - - def test_apply_dynamic_database_filter( mocker: MockerFixture, app: Any, @@ -698,10 +545,6 @@ def test_apply_dynamic_database_filter( # mock the lookup so that we don't need to include the driver mocker.patch("sqlalchemy.engine.URL.get_driver_name", return_value="gsheets") mocker.patch("superset.utils.log.DBEventLogger.log") - mocker.patch( - "superset.commands.database.ssh_tunnel.delete.is_feature_enabled", - return_value=False, - ) def _base_filter(query): from superset.models.core import Database diff --git a/tests/unit_tests/databases/commands/utils_test.py b/tests/unit_tests/databases/commands/utils_test.py index 0e5393639fd..266c4e5d864 100644 --- a/tests/unit_tests/databases/commands/utils_test.py +++ b/tests/unit_tests/databases/commands/utils_test.py @@ -29,13 +29,12 @@ def test_add_permissions(mocker: MockerFixture) -> None: database.db_engine_spec.supports_catalog = True database.get_all_catalog_names.return_value = ["catalog1", "catalog2"] database.get_all_schema_names.side_effect = [["schema1"], ["schema2"]] - ssh_tunnel = mocker.MagicMock() add_permission_view_menu = mocker.patch( "superset.commands.database.importers.v1.utils.security_manager." "add_permission_view_menu" ) - add_permissions(database, ssh_tunnel) + add_permissions(database) add_permission_view_menu.assert_has_calls( [ @@ -60,13 +59,12 @@ def test_add_permissions_get_default_catalog(mocker: MockerFixture): database.get_all_catalog_names.return_value = ["catalog1", "catalog2"] database.get_default_catalog.return_value = "catalog1" database.get_all_schema_names.side_effect = [["schema1"], ["schema2"]] - ssh_tunnel = mocker.MagicMock() add_permission_view_menu = mocker.patch( "superset.commands.database.importers.v1.utils.security_manager." "add_permission_view_menu" ) - add_permissions(database, ssh_tunnel) + add_permissions(database) add_permission_view_menu.assert_has_calls( [ @@ -88,13 +86,12 @@ def test_add_permissions_handle_failures(mocker: MockerFixture) -> None: database.db_engine_spec.supports_catalog = True database.get_all_catalog_names.return_value = ["catalog1", "catalog2", "catalog3"] database.get_all_schema_names.side_effect = [["schema1"], Exception, ["schema3"]] - ssh_tunnel = mocker.MagicMock() add_permission_view_menu = mocker.patch( "superset.commands.database.importers.v1.utils.security_manager." "add_permission_view_menu" ) - add_permissions(database, ssh_tunnel) + add_permissions(database) add_permission_view_menu.assert_has_calls( [ diff --git a/tests/unit_tests/databases/dao/dao_tests.py b/tests/unit_tests/databases/dao/dao_tests.py index a826d01be8e..3fbae7e2662 100644 --- a/tests/unit_tests/databases/dao/dao_tests.py +++ b/tests/unit_tests/databases/dao/dao_tests.py @@ -54,7 +54,9 @@ def test_database_get_ssh_tunnel(session_with_data: Session) -> None: from superset.daos.database import DatabaseDAO from superset.databases.ssh_tunnel.models import SSHTunnel - result = DatabaseDAO.get_ssh_tunnel(1) + database = DatabaseDAO.find_by_id(1, skip_base_filter=True) + assert database is not None + result = database.ssh_tunnel assert result assert isinstance(result, SSHTunnel) @@ -64,6 +66,7 @@ def test_database_get_ssh_tunnel(session_with_data: Session) -> None: def test_database_get_ssh_tunnel_not_found(session_with_data: Session) -> None: from superset.daos.database import DatabaseDAO - result = DatabaseDAO.get_ssh_tunnel(2) + database = DatabaseDAO.find_by_id(2, skip_base_filter=True) + result = database.ssh_tunnel if database else None assert result is None diff --git a/tests/unit_tests/databases/filters_test.py b/tests/unit_tests/databases/filters_test.py index 253798d9cdf..066c31a63d8 100644 --- a/tests/unit_tests/databases/filters_test.py +++ b/tests/unit_tests/databases/filters_test.py @@ -117,10 +117,7 @@ def test_database_filter(mocker: MockerFixture) -> None: engine, compile_kwargs={"literal_binds": True}, ) - space = " " # pre-commit removes trailing spaces... assert ( str(compiled_query) - == f"""SELECT dbs.uuid, dbs.created_on, dbs.changed_on, dbs.id, dbs.verbose_name, dbs.database_name, dbs.sqlalchemy_uri, dbs.password, dbs.cache_timeout, dbs.select_as_create_table_as, dbs.expose_in_sqllab, dbs.configuration_method, dbs.allow_run_async, dbs.allow_file_upload, dbs.allow_ctas, dbs.allow_cvas, dbs.allow_dml, dbs.force_ctas_schema, dbs.extra, dbs.encrypted_extra, dbs.impersonate_user, dbs.server_cert, dbs.is_managed_externally, dbs.external_url, dbs.created_by_fk, dbs.changed_by_fk{space} -FROM dbs{space} -WHERE '[' || dbs.database_name || '].(id:' || CAST(dbs.id AS VARCHAR) || ')' IN ('[my_db].(id:42)', '[my_other_db].(id:43)') OR dbs.database_name IN ('my_db', 'my_other_db', 'third_db')""" # noqa: S608, E501 + == "SELECT dbs.uuid, dbs.created_on, dbs.changed_on, dbs.id, dbs.verbose_name, dbs.database_name, dbs.sqlalchemy_uri, dbs.password, dbs.cache_timeout, dbs.select_as_create_table_as, dbs.expose_in_sqllab, dbs.configuration_method, dbs.allow_run_async, dbs.allow_file_upload, dbs.allow_ctas, dbs.allow_cvas, dbs.allow_dml, dbs.force_ctas_schema, dbs.extra, dbs.encrypted_extra, dbs.impersonate_user, dbs.server_cert, dbs.is_managed_externally, dbs.external_url, dbs.created_by_fk, dbs.changed_by_fk, ssh_tunnels_1.uuid AS uuid_1, ssh_tunnels_1.created_on AS created_on_1, ssh_tunnels_1.changed_on AS changed_on_1, ssh_tunnels_1.extra_json, ssh_tunnels_1.id AS id_1, ssh_tunnels_1.database_id, ssh_tunnels_1.server_address, ssh_tunnels_1.server_port, ssh_tunnels_1.username, ssh_tunnels_1.password AS password_1, ssh_tunnels_1.private_key, ssh_tunnels_1.private_key_password, ssh_tunnels_1.created_by_fk AS created_by_fk_1, ssh_tunnels_1.changed_by_fk AS changed_by_fk_1 \nFROM dbs LEFT OUTER JOIN ssh_tunnels AS ssh_tunnels_1 ON dbs.id = ssh_tunnels_1.database_id \nWHERE '[' || dbs.database_name || '].(id:' || CAST(dbs.id AS VARCHAR) || ')' IN ('[my_db].(id:42)', '[my_other_db].(id:43)') OR dbs.database_name IN ('my_db', 'my_other_db', 'third_db')" # noqa: E501 ) diff --git a/tests/unit_tests/databases/ssh_tunnel/commands/__init__.py b/tests/unit_tests/databases/ssh_tunnel/commands/__init__.py deleted file mode 100644 index 13a83393a91..00000000000 --- a/tests/unit_tests/databases/ssh_tunnel/commands/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/tests/unit_tests/databases/ssh_tunnel/commands/create_test.py b/tests/unit_tests/databases/ssh_tunnel/commands/create_test.py deleted file mode 100644 index 168c0fc3d53..00000000000 --- a/tests/unit_tests/databases/ssh_tunnel/commands/create_test.py +++ /dev/null @@ -1,148 +0,0 @@ -# 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. - - -import pytest -from sqlalchemy.orm.session import Session - -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelDatabasePortError, - SSHTunnelInvalidError, -) - - -def test_create_ssh_tunnel_command(session: Session) -> None: - from superset import db - from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - engine = db.session.get_bind() - Database.metadata.create_all(engine) # pylint: disable=no-member - - database = Database( - database_name="my_database", - sqlalchemy_uri="postgresql://u:p@localhost:5432/db", - ) - - properties = { - "database": database, - "server_address": "123.132.123.1", - "server_port": "3005", - "username": "foo", - "password": "bar", - } - - result = CreateSSHTunnelCommand(database, properties).run() - - assert result is not None - assert isinstance(result, SSHTunnel) - - -def test_create_ssh_tunnel_command_invalid_params(session: Session) -> None: - from superset import db - from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand - from superset.models.core import Database - - engine = db.session.get_bind() - Database.metadata.create_all(engine) # pylint: disable=no-member - - database = Database( - database_name="my_database", - sqlalchemy_uri="postgresql://u:p@localhost:5432/db", - ) - - # If we are trying to create a tunnel with a private_key_password - # then a private_key is mandatory - properties = { - "database": database, - "server_address": "123.132.123.1", - "server_port": "3005", - "username": "foo", - "private_key_password": "bar", - } - - command = CreateSSHTunnelCommand(database, properties) - - with pytest.raises(SSHTunnelInvalidError) as excinfo: - command.run() - assert str(excinfo.value) == ("SSH Tunnel parameters are invalid.") - - -def test_create_ssh_tunnel_command_no_port(session: Session) -> None: - """ - Test that SSH Tunnel can be created without explicit port but with a default one. - """ - from superset import db - from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - engine = db.session.get_bind() - Database.metadata.create_all(engine) # pylint: disable=no-member - - database = Database( - database_name="my_database", - sqlalchemy_uri="postgresql://u:p@localhost/db", - ) - - properties = { - "database": database, - "server_address": "123.132.123.1", - "server_port": "3005", - "username": "foo", - "password": "bar", - } - - result = CreateSSHTunnelCommand(database, properties).run() - - assert result is not None - assert isinstance(result, SSHTunnel) - - -def test_create_ssh_tunnel_command_no_port_no_default(session: Session) -> None: - """ - Test that error is raised when creating SSH Tunnel without explicit/default ports. - """ - from superset import db - from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand - from superset.models.core import Database - - engine = db.session.get_bind() - Database.metadata.create_all(engine) # pylint: disable=no-member - - database = Database( - id=1, - database_name="my_database", - sqlalchemy_uri="weird+db://u:p@localhost/db", - ) - - properties = { - "database": database, - "server_address": "123.132.123.1", - "server_port": "3005", - "username": "foo", - "password": "bar", - } - - command = CreateSSHTunnelCommand(database, properties) - - with pytest.raises(SSHTunnelDatabasePortError) as excinfo: - command.run() - assert str(excinfo.value) == ( - "A database port is required when connecting via SSH Tunnel." - ) diff --git a/tests/unit_tests/databases/ssh_tunnel/commands/delete_test.py b/tests/unit_tests/databases/ssh_tunnel/commands/delete_test.py deleted file mode 100644 index ba1c874276e..00000000000 --- a/tests/unit_tests/databases/ssh_tunnel/commands/delete_test.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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 collections.abc import Iterator - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy.orm.session import Session - - -@pytest.fixture -def session_with_data(session: Session) -> Iterator[Session]: - from superset.connectors.sqla.models import SqlaTable - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - engine = session.get_bind() - SqlaTable.metadata.create_all(engine) # pylint: disable=no-member - - database = Database(database_name="my_database", sqlalchemy_uri="sqlite://") - sqla_table = SqlaTable( - table_name="my_sqla_table", - columns=[], - metrics=[], - database=database, - ) - ssh_tunnel = SSHTunnel( - database_id=database.id, - database=database, - ) - - session.add(database) - session.add(sqla_table) - session.add(ssh_tunnel) - session.flush() - yield session - session.rollback() - - -def test_delete_ssh_tunnel_command( - mocker: MockerFixture, session_with_data: Session -) -> None: - from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand - from superset.daos.database import DatabaseDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert 1 == result.database_id - mocker.patch( - "superset.commands.database.ssh_tunnel.delete.is_feature_enabled", - return_value=True, - ) - DeleteSSHTunnelCommand(1).run() - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result is None diff --git a/tests/unit_tests/databases/ssh_tunnel/commands/update_test.py b/tests/unit_tests/databases/ssh_tunnel/commands/update_test.py deleted file mode 100644 index f223014ee34..00000000000 --- a/tests/unit_tests/databases/ssh_tunnel/commands/update_test.py +++ /dev/null @@ -1,155 +0,0 @@ -# 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 collections.abc import Iterator - -import pytest -from sqlalchemy.orm.session import Session - -from superset.commands.database.ssh_tunnel.exceptions import ( - SSHTunnelDatabasePortError, - SSHTunnelInvalidError, -) - - -@pytest.fixture -def session_with_data(request, session: Session) -> Iterator[Session]: - from superset.connectors.sqla.models import SqlaTable - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - engine = session.get_bind() - SqlaTable.metadata.create_all(engine) # pylint: disable=no-member - - sqlalchemy_uri = getattr(request, "param", "postgresql://u:p@localhost:5432/db") - database = Database(database_name="my_database", sqlalchemy_uri=sqlalchemy_uri) - sqla_table = SqlaTable( - table_name="my_sqla_table", - columns=[], - metrics=[], - database=database, - ) - ssh_tunnel = SSHTunnel( - database_id=database.id, database=database, server_address="Test" - ) - - session.add(database) - session.add(sqla_table) - session.add(ssh_tunnel) - session.flush() - yield session - session.rollback() - - -def test_update_shh_tunnel_command(session_with_data: Session) -> None: - from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand - from superset.daos.database import DatabaseDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert 1 == result.database_id - assert "Test" == result.server_address - - update_payload = {"server_address": "Test2"} - UpdateSSHTunnelCommand(1, update_payload).run() - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert "Test2" == result.server_address - - -def test_update_shh_tunnel_invalid_params(session_with_data: Session) -> None: - from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand - from superset.daos.database import DatabaseDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert 1 == result.database_id - assert "Test" == result.server_address - - # If we are trying to update a tunnel with a private_key_password - # then a private_key is mandatory - update_payload = {"private_key_password": "pass"} - command = UpdateSSHTunnelCommand(1, update_payload) - - with pytest.raises(SSHTunnelInvalidError) as excinfo: - command.run() - assert str(excinfo.value) == ("SSH Tunnel parameters are invalid.") - - -@pytest.mark.parametrize( - "session_with_data", ["postgresql://u:p@localhost/testdb"], indirect=True -) -def test_update_shh_tunnel_no_port(session_with_data: Session) -> None: - """ - Test that SSH Tunnel can be updated without explicit port but with a default one. - """ - from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand - from superset.daos.database import DatabaseDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert 1 == result.database_id - assert "Test" == result.server_address - - update_payload = {"server_address": "Test2"} - UpdateSSHTunnelCommand(1, update_payload).run() - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert "Test2" == result.server_address - - -@pytest.mark.parametrize( - "session_with_data", ["weird+db://u:p@localhost/testdb"], indirect=True -) -def test_update_shh_tunnel_no_port_no_default(session_with_data: Session) -> None: - """ - Test that error is raised when updating SSH Tunnel without explicit/default ports. - """ - from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand - from superset.daos.database import DatabaseDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - - result = DatabaseDAO.get_ssh_tunnel(1) - - assert result - assert isinstance(result, SSHTunnel) - assert 1 == result.database_id - assert "Test" == result.server_address - - update_payload = {"server_address": "Test update"} - command = UpdateSSHTunnelCommand(1, update_payload) - - with pytest.raises(SSHTunnelDatabasePortError) as excinfo: - command.run() - assert str(excinfo.value) == ( - "A database port is required when connecting via SSH Tunnel." - ) diff --git a/tests/unit_tests/databases/ssh_tunnel/dao_tests.py b/tests/unit_tests/databases/ssh_tunnel/dao_tests.py deleted file mode 100644 index a24a94ec36d..00000000000 --- a/tests/unit_tests/databases/ssh_tunnel/dao_tests.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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. - - -def test_create_ssh_tunnel(): - from superset.daos.database import SSHTunnelDAO - from superset.databases.ssh_tunnel.models import SSHTunnel - from superset.models.core import Database - - database = Database(id=1, database_name="my_database", sqlalchemy_uri="sqlite://") - - result = SSHTunnelDAO.create( - attributes={ - "database_id": database.id, - "server_address": "123.132.123.1", - "server_port": "3005", - "username": "foo", - "password": "bar", - }, - ) - - assert result is not None - assert isinstance(result, SSHTunnel) diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index a2c3adc1c21..26e4cd2736f 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -333,7 +333,7 @@ def test_get_all_catalog_names(mocker: MockerFixture) -> None: inspector.bind.execute.return_value = [("examples",), ("other",)] assert database.get_all_catalog_names(force=True) == {"examples", "other"} - get_inspector.assert_called_with(ssh_tunnel=None) + get_inspector.assert_called_with() def test_get_all_schema_names_needs_oauth2(mocker: MockerFixture) -> None: