This is an automated email from the ASF dual-hosted git repository. hugh pushed a commit to branch test-ssh-tunnel-1 in repository https://gitbox.apache.org/repos/asf/superset.git
commit f00bd69dc9d6d5d99af79d8fb6b809f4067e77e2 Author: Antonio Rivero <[email protected]> AuthorDate: Tue Nov 22 15:54:21 2022 -0300 SSH Tunnel: - Initial DELETE and CREATE ssh tunnel APIs - Add happy path tests for each API NOTE: - Updating a DB might also result in creating a Tunnel, that flow will get tested in the UPDATE API changes. --- superset/constants.py | 1 + superset/databases/api.py | 62 ++++++++++++++++++++++ superset/databases/commands/create.py | 13 +++++ superset/databases/schemas.py | 41 ++++++++------- tests/integration_tests/databases/api_tests.py | 31 +++++++++++ tests/unit_tests/databases/api_test.py | 72 ++++++++++++++++++++++++++ 6 files changed, 200 insertions(+), 20 deletions(-) diff --git a/superset/constants.py b/superset/constants.py index 7d759acf67..b775926b21 100644 --- a/superset/constants.py +++ b/superset/constants.py @@ -136,6 +136,7 @@ MODEL_API_RW_METHOD_PERMISSION_MAP = { "validate_sql": "read", "get_data": "read", "samples": "read", + "delete_ssh_tunnel": "write", } EXTRA_FORM_DATA_APPEND_KEYS = { diff --git a/superset/databases/api.py b/superset/databases/api.py index aced8e7c6f..9cab40ede4 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -72,6 +72,12 @@ from superset.databases.schemas import ( ValidateSQLRequest, ValidateSQLResponse, ) +from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand +from superset.databases.ssh_tunnel.commands.delete import DeleteSSHTunnelCommand +from superset.databases.ssh_tunnel.commands.exceptions import ( + SSHTunnelDeleteFailedError, + SSHTunnelNotFoundError, +) from superset.databases.utils import get_table_metadata from superset.db_engine_specs import get_available_engine_specs from superset.errors import ErrorLevel, SupersetError, SupersetErrorType @@ -107,6 +113,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi): "available", "validate_parameters", "validate_sql", + "delete_ssh_tunnel", } resource_name = "database" class_permission_name = "Database" @@ -1204,3 +1211,58 @@ class DatabaseRestApi(BaseSupersetModelRestApi): command = ValidateDatabaseParametersCommand(payload) command.run() return self.response(200, message="OK") + + @expose("/<int:pk>/ssh_tunnel/", methods=["DELETE"]) + @protect() + @safe + @statsd_metrics + @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: + """Deletes a SSH Tunnel + --- + delete: + description: >- + Deletes 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' + """ + try: + DeleteSSHTunnelCommand(pk).run() + return self.response(200, message="OK") + except SSHTunnelNotFoundError: + 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)) diff --git a/superset/databases/commands/create.py b/superset/databases/commands/create.py index 4dc8e8eda4..421e060a4b 100644 --- a/superset/databases/commands/create.py +++ b/superset/databases/commands/create.py @@ -31,6 +31,9 @@ from superset.databases.commands.exceptions import ( ) from superset.databases.commands.test_connection import TestConnectionDatabaseCommand from superset.databases.dao import DatabaseDAO +from superset.databases.ssh_tunnel.commands.exceptions import SSHTunnelInvalidError +from superset.databases.ssh_tunnel.dao import SSHTunnelDAO +from superset.databases.ssh_tunnel.models import SSHTunnel from superset.exceptions import SupersetErrorsException from superset.extensions import db, event_logger, security_manager @@ -77,6 +80,16 @@ class CreateDatabaseCommand(BaseCommand): security_manager.add_permission_view_menu( "schema_access", security_manager.get_schema_perm(database, schema) ) + + if self._properties.get("ssh_tunnel"): + SSHTunnelDAO.create( + { + **self._properties.get("ssh_tunnel"), + "database_id": database.id, + }, + commit=False, + ) + db.session.commit() except DAOCreateFailedError as ex: db.session.rollback() diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index b874d147a2..9c2de0b18e 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -365,6 +365,26 @@ class DatabaseValidateParametersSchema(Schema): ) +class DatabaseSSHTunnel(Schema): + id = fields.Integer() + database_id = fields.Integer() + + server_address = fields.String() + server_port = fields.Integer() + username = fields.String() + + # Basic Authentication + password = fields.String(required=False) + + # password protected private key authentication + private_key = fields.String(required=False) + private_key_password = fields.String(required=False) + + # remote binding port + bind_host = fields.String() + bind_port = fields.Integer() + + class DatabasePostSchema(Schema, DatabaseParametersSchemaMixin): class Meta: # pylint: disable=too-few-public-methods unknown = EXCLUDE @@ -409,6 +429,7 @@ class DatabasePostSchema(Schema, DatabaseParametersSchemaMixin): is_managed_externally = fields.Boolean(allow_none=True, default=False) external_url = fields.String(allow_none=True) uuid = fields.String(required=False) + ssh_tunnel = fields.Nested(DatabaseSSHTunnel, allow_none=True) class DatabasePutSchema(Schema, DatabaseParametersSchemaMixin): @@ -456,26 +477,6 @@ class DatabasePutSchema(Schema, DatabaseParametersSchemaMixin): external_url = fields.String(allow_none=True) -class DatabaseSSHTunnel(Schema): - id = fields.Integer() - database_id = fields.Integer() - - server_address = fields.String() - server_port = fields.Integer() - username = fields.String() - - # Basic Authentication - password = fields.String(required=False) - - # password protected private key authentication - private_key = fields.String(required=False) - private_key_password = fields.String(required=False) - - # remote binding port - bind_host = fields.String() - bind_port = fields.Integer() - - class DatabaseTestConnectionSchema(Schema, DatabaseParametersSchemaMixin): rename_encrypted_extra = pre_load(rename_encrypted_extra) diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index 9698c7e42a..7821900c87 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -35,6 +35,7 @@ from sqlalchemy.sql import func from superset import db, security_manager from superset.connectors.sqla.models import SqlaTable +from superset.databases.utils import make_url_safe from superset.db_engine_specs.mysql import MySQLEngineSpec from superset.db_engine_specs.postgres import PostgresEngineSpec from superset.db_engine_specs.redshift import RedshiftEngineSpec @@ -279,6 +280,36 @@ class TestDatabaseApi(SupersetTestCase): db.session.delete(model) db.session.commit() + def test_create_database_with_ssh_tunnel(self): + """ + Database API: Test create with SSH Tunnel + """ + self.login(username="admin") + example_db = get_example_database() + if example_db.backend == "sqlite": + return + ssh_tunnel_properties = { + "server_address": "123.132.123.1", + "bind_host": "localhost", + "bind_port": "5432", + "username": "foo", + "password": "bar", + } + database_data = { + "database_name": "test-create-db-with-ssh-tunnel", + "sqlalchemy_uri": example_db.sqlalchemy_uri_decrypted, + "ssh_tunnel": ssh_tunnel_properties, + } + + uri = "api/v1/database/" + rv = self.client.post(uri, json=database_data) + response = json.loads(rv.data.decode("utf-8")) + self.assertEqual(rv.status_code, 201) + # Cleanup + model = db.session.query(Database).get(response.get("id")) + db.session.delete(model) + db.session.commit() + def test_create_database_invalid_configuration_method(self): """ Database API: Test create with an invalid configuration method. diff --git a/tests/unit_tests/databases/api_test.py b/tests/unit_tests/databases/api_test.py index d6f8897c4a..73e79f3a82 100644 --- a/tests/unit_tests/databases/api_test.py +++ b/tests/unit_tests/databases/api_test.py @@ -191,3 +191,75 @@ def test_non_zip_import(client: Any, full_api_access: None) -> None: } ] } + + +def test_delete_ssh_tunnel( + mocker: MockFixture, + app: Any, + session: Session, + client: Any, + full_api_access: None, +) -> None: + """ + Test that we can delete SSH Tunnel + """ + with app.app_context(): + from superset.databases.api import DatabaseRestApi + from superset.databases.dao import DatabaseDAO + 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-se...@black-sanctum-314419.iam.gserviceaccount.com", + "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", + }, + } + ), + ) + session.add(database) + 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") + + # Create our SSHTunnel + tunnel = SSHTunnel( + database_id=1, + database=database, + ) + + session.add(tunnel) + 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("/api/v1/database/1/ssh_tunnel/") + assert response_delete_tunnel.json["message"] == "OK" + + response_tunnel = DatabaseDAO.get_ssh_tunnel(1) + assert response_tunnel is None
