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

Reply via email to