feat(ssh-tunnelling): Setup SSH Tunneling Commands for Database Connections

requirements/base.txt

@@ -19,6 +19,8 @@ babel==2.9.1
     # via flask-babel
 backoff==1.11.1
     # via apache-superset
+bcrypt==4.0.1
+    # via paramiko
 billiard==3.6.4.0
     # via celery
 bleach==3.3.1
@@ -30,7 +32,9 @@ cachelib==0.4.1
 celery==5.2.2
     # via apache-superset
 cffi==1.14.6
-    # via cryptography
+    # via
+    #   cryptography
+    #   pynacl
 click==8.0.4
     # via
     #   apache-superset
@@ -57,7 +61,9 @@ cron-descriptor==1.2.24
 croniter==1.0.15
     # via apache-superset
 cryptography==3.4.7
-    # via apache-superset
+    # via
+    #   apache-superset
+    #   paramiko
 deprecation==2.1.0
     # via apache-superset
 dnspython==2.1.0
@@ -167,6 +173,8 @@ packaging==21.3
     #   deprecation
 pandas==1.4.4
     # via apache-superset
+paramiko==2.11.0
+    # via sshtunnel
 parsedatetime==2.6
     # via apache-superset
 pgsanity==0.2.9
@@ -188,6 +196,8 @@ pyjwt==2.4.0
     #   flask-jwt-extended
 pymeeus==0.5.11
     # via convertdate
+pynacl==1.5.0
+    # via paramiko
 pyparsing==3.0.6
     # via
     #   apache-superset
@@ -231,6 +241,7 @@ six==1.16.0
     #   flask-talisman
     #   isodate
     #   jsonschema
+    #   paramiko
     #   polyline
     #   prison
     #   pyrsistent
@@ -252,6 +263,8 @@ sqlalchemy-utils==0.38.3
     #   flask-appbuilder
 sqlparse==0.4.3
     # via apache-superset
+sshtunnel==0.4.0
+    # via apache-superset
 tabulate==0.8.9
     # via apache-superset
 typing-extensions==4.4.0

setup.py

@@ -113,6 +113,7 @@ def get_git_sha() -> str:
         "PyJWT>=2.4.0, <3.0",
         "redis",
         "selenium>=3.141.0",
+        "sshtunnel>=0.4.0, <0.5",
         "simplejson>=3.15.0",
         "slack_sdk>=3.1.1, <4",
         "sqlalchemy>=1.4, <2",

superset/constants.py

@@ -34,6 +34,8 @@
 
 NO_TIME_RANGE = "No filter"
 
+SSH_TUNNELLING_LOCAL_BIND_ADDRESS = "127.0.0.1"
+
 
 class RouteMethod:  # pylint: disable=too-few-public-methods
     """

superset/databases/commands/create.py

@@ -31,6 +31,7 @@
 )
 from superset.databases.commands.test_connection import TestConnectionDatabaseCommand
 from superset.databases.dao import DatabaseDAO
+from superset.databases.ssh_tunnel.dao import SSHTunnelDAO
 from superset.exceptions import SupersetErrorsException
 from superset.extensions import db, event_logger, security_manager
 
@@ -70,13 +71,25 @@ def run(self) -> Model:
         try:
             database = DatabaseDAO.create(self._properties, commit=False)
             database.set_sqlalchemy_uri(database.sqlalchemy_uri)
+            db.session.flush()
+
+            ssh_tunnel = None
+            if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
+                ssh_tunnel = SSHTunnelDAO.create(
+                    {
+                        **ssh_tunnel_properties,
+                        "database_id": database.id,
+                    },
+                    commit=False,
+                )
 
             # adding a new database we always want to force refresh schema list
-            schemas = database.get_all_schema_names(cache=False)
+            schemas = database.get_all_schema_names(cache=False, ssh_tunnel=ssh_tunnel)
             for schema in schemas:
                 security_manager.add_permission_view_menu(
                     "schema_access", security_manager.get_schema_perm(database, schema)
                 )
+
             db.session.commit()
         except DAOCreateFailedError as ex:
             db.session.rollback()

superset/databases/commands/test_connection.py

@@ -32,6 +32,7 @@
     DatabaseTestConnectionUnexpectedError,
 )
 from superset.databases.dao import DatabaseDAO
+from superset.databases.ssh_tunnel.models import SSHTunnel
 from superset.databases.utils import make_url_safe
 from superset.errors import ErrorLevel, SupersetErrorType
 from superset.exceptions import (
@@ -90,6 +91,10 @@ def run(self) -> None:  # pylint: disable=too-many-statements
             database.set_sqlalchemy_uri(uri)
             database.db_engine_spec.mutate_db_for_connection_test(database)
 
+            # Generate tunnel if present in the properties
+            if ssh_tunnel := self._properties.get("ssh_tunnel"):
+                ssh_tunnel = SSHTunnel(**ssh_tunnel)
+
             event_logger.log_with_context(
                 action="test_connection_attempt",
                 engine=database.db_engine_spec.__name__,
@@ -99,7 +104,9 @@ def ping(engine: Engine) -> bool:
                 with closing(engine.raw_connection()) as conn:
                     return engine.dialect.do_ping(conn)
 
-            with database.get_sqla_engine_with_context() as engine:
+            with database.get_sqla_engine_with_context(
+                override_ssh_tunnel=ssh_tunnel
+            ) as engine:
                 try:
                     alive = func_timeout(
                         app.config["TEST_DATABASE_CONNECTION_TIMEOUT"].total_seconds(),

superset/databases/dao.py

@@ -19,6 +19,7 @@
 
 from superset.dao.base import BaseDAO
 from superset.databases.filters import DatabaseFilter
+from superset.databases.ssh_tunnel.models import SSHTunnel
 from superset.extensions import db
 from superset.models.core import Database
 from superset.models.dashboard import Dashboard
@@ -124,3 +125,13 @@ def get_related_objects(cls, database_id: int) -> Dict[str, Any]:
         return dict(
             charts=charts, dashboards=dashboards, sqllab_tab_states=sqllab_tab_states
         )
+
+    @classmethod
+    def get_ssh_tunnel(cls, database_id: int) -> Optional[SSHTunnel]:
+        ssh_tunnel = (
+            db.session.query(SSHTunnel)
+            .filter(SSHTunnel.database_id == database_id)
+            .one_or_none()
+        )
+
+        return ssh_tunnel

superset/databases/schemas.py

@@ -365,6 +365,19 @@ class Meta:  # pylint: disable=too-few-public-methods
     )
 
 
+class DatabaseSSHTunnel(Schema):
+    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)
+
+
 class DatabasePostSchema(Schema, DatabaseParametersSchemaMixin):
     class Meta:  # pylint: disable=too-few-public-methods
         unknown = EXCLUDE
@@ -409,6 +422,7 @@ class Meta:  # pylint: disable=too-few-public-methods
     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):
@@ -482,6 +496,8 @@ class DatabaseTestConnectionSchema(Schema, DatabaseParametersSchemaMixin):
         validate=[Length(1, 1024), sqlalchemy_uri_validator],
     )
 
+    ssh_tunnel = fields.Nested(DatabaseSSHTunnel, allow_none=True)
+
 
 class TableMetadataOptionsResponseSchema(Schema):
     deferrable = fields.Bool()

superset/databases/ssh_tunnel/__init__.py

@@ -0,0 +1,16 @@
+# 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.

superset/databases/ssh_tunnel/commands/__init__.py

@@ -0,0 +1,16 @@
+# 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.

superset/databases/ssh_tunnel/commands/create.py

@@ -0,0 +1,82 @@
+# 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 typing import Any, Dict, List, Optional
+
+from flask_appbuilder.models.sqla import Model
+from marshmallow import ValidationError
+
+from superset.commands.base import BaseCommand
+from superset.dao.exceptions import DAOCreateFailedError
+from superset.databases.ssh_tunnel.commands.exceptions import (
+    SSHTunnelCreateFailedError,
+    SSHTunnelInvalidError,
+    SSHTunnelRequiredFieldValidationError,
+)
+from superset.databases.ssh_tunnel.dao import SSHTunnelDAO
+from superset.extensions import event_logger
+
+logger = logging.getLogger(__name__)
+
+
+class CreateSSHTunnelCommand(BaseCommand):
+    def __init__(self, database_id: int, data: Dict[str, Any]):
+        self._properties = data.copy()
+        self._properties["database_id"] = database_id
+
+    def run(self) -> Model:
+        self.validate()
+
+        try:
+            tunnel = SSHTunnelDAO.create(self._properties, commit=False)
+        except DAOCreateFailedError as ex:
+            raise SSHTunnelCreateFailedError() from ex
+
+        return tunnel
+
+    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] = []
+        database_id: Optional[int] = self._properties.get("database_id")
+        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")
+        private_key: Optional[str] = self._properties.get("private_key")
+        private_key_password: Optional[str] = self._properties.get(
+            "private_key_password"
+        )
+        if not database_id:
+            exceptions.append(SSHTunnelRequiredFieldValidationError("database_id"))
+        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 private_key_password and private_key is None:
+            exceptions.append(SSHTunnelRequiredFieldValidationError("private_key"))
+        if exceptions:
+            exception = SSHTunnelInvalidError()
+            exception.add_list(exceptions)
+            event_logger.log_with_context(
+                action="ssh_tunnel_creation_failed.{}.{}".format(
+                    exception.__class__.__name__,
+                    ".".join(exception.get_list_classnames()),
+                )
+            )
+            raise exception

superset/databases/ssh_tunnel/commands/delete.py

@@ -0,0 +1,51 @@
+# 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 typing import Optional
+
+from flask_appbuilder.models.sqla import Model
+
+from superset.commands.base import BaseCommand
+from superset.dao.exceptions import DAODeleteFailedError
+from superset.databases.ssh_tunnel.commands.exceptions import (
+    SSHTunnelDeleteFailedError,
+    SSHTunnelNotFoundError,
+)
+from superset.databases.ssh_tunnel.dao import SSHTunnelDAO
+from superset.databases.ssh_tunnel.models import SSHTunnel
+
+logger = logging.getLogger(__name__)
+
+
+class DeleteSSHTunnelCommand(BaseCommand):
+    def __init__(self, model_id: int):
+        self._model_id = model_id
+        self._model: Optional[SSHTunnel] = None
+
+    def run(self) -> Model:
+        self.validate()
+        try:
+            ssh_tunnel = SSHTunnelDAO.delete(self._model)
+        except DAODeleteFailedError as ex:
+            raise SSHTunnelDeleteFailedError() from ex
+        return ssh_tunnel
+
+    def validate(self) -> None:
+        # Validate/populate model exists
+        self._model = SSHTunnelDAO.find_by_id(self._model_id)
+        if not self._model:
+            raise SSHTunnelNotFoundError()