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

…ctions (#21912) Co-authored-by: Antonio Rivero Martinez <[email protected]> Co-authored-by: Elizabeth Thompson <[email protected]>
apache · Jan 3, 2023 · ebaad10 · ebaad10
1 parent a7a4561
commit ebaad10
Show file tree

Hide file tree

Showing 40 changed files with 1,905 additions and 47 deletions.
diff --git a/requirements/base.txt b/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
@@ -57,7 +59,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 +171,8 @@ packaging==21.3
     #   deprecation
 pandas==1.5.2
     # via apache-superset
+paramiko==2.11.0
+    # via sshtunnel
 parsedatetime==2.6
     # via apache-superset
 pgsanity==0.2.9
@@ -188,6 +194,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 +239,7 @@ six==1.16.0
     #   flask-talisman
     #   isodate
     #   jsonschema
+    #   paramiko
     #   polyline
     #   prison
     #   pyrsistent
@@ -252,6 +261,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

diff --git a/setup.py b/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",

diff --git a/superset/config.py b/superset/config.py
@@ -476,8 +476,30 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]:
     "DRILL_TO_DETAIL": False,
     "DATAPANEL_CLOSED_BY_DEFAULT": False,
     "HORIZONTAL_FILTER_BAR": False,
+    # Allow users to enable ssh tunneling when creating a DB.
+    # Users must check whether the DB engine supports SSH Tunnels
+    # otherwise enabling this flag won't have any effect on the DB.
+    "SSH_TUNNELING": False,
 }
 
+# ------------------------------
+# SSH Tunnel
+# ------------------------------
+# Allow users to set the host used when connecting to the SSH Tunnel
+# as localhost and any other alias (0.0.0.0)
+# ----------------------------------------------------------------------
+#                             |
+# -------------+              |    +----------+
+#     LOCAL    |              |    |  REMOTE  | :22 SSH
+#     CLIENT   | <== SSH ========> |  SERVER  | :8080 web service
+# -------------+              |    +----------+
+#                             |
+#                          FIREWALL (only port 22 is open)
+
+# ----------------------------------------------------------------------
+SSH_TUNNEL_MANAGER_CLASS = "superset.extensions.ssh.SSHManager"
+SSH_TUNNEL_LOCAL_BIND_ADDRESS = "127.0.0.1"
+
 # Feature flags may also be set via 'SUPERSET_FEATURE_' prefixed environment vars.
 DEFAULT_FEATURE_FLAGS.update(
     {
@@ -1506,7 +1528,7 @@ def EMAIL_HEADER_MUTATOR(  # pylint: disable=invalid-name,unused-argument
     try:
         # pylint: disable=import-error,wildcard-import,unused-wildcard-import
         import superset_config
-        from superset_config import *  # type:ignore
+        from superset_config import *  # type: ignore
 
         print(f"Loaded your LOCAL configuration at [{superset_config.__file__}]")
     except Exception:

diff --git a/superset/constants.py b/superset/constants.py
@@ -139,6 +139,7 @@ class RouteMethod:  # pylint: disable=too-few-public-methods
     "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
@@ -72,6 +72,11 @@
     ValidateSQLRequest,
     ValidateSQLResponse,
 )
+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
@@ -80,6 +85,7 @@
 from superset.models.core import Database
 from superset.superset_typing import FlaskResponse
 from superset.utils.core import error_msg_from_exception, parse_js_uri_path_item
+from superset.utils.ssh_tunnel import mask_password_info
 from superset.views.base import json_errors_response
 from superset.views.base_api import (
     BaseSupersetModelRestApi,
@@ -107,6 +113,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
         "available",
         "validate_parameters",
         "validate_sql",
+        "delete_ssh_tunnel",
     }
     resource_name = "database"
     class_permission_name = "Database"
@@ -219,6 +226,47 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
         ValidateSQLResponse,
     )
 
+    @expose("/<int:pk>", methods=["GET"])
+    @protect()
+    @safe
+    def get(self, pk: int, **kwargs: Any) -> Response:
+        """Get a database
+        ---
+        get:
+          description: >-
+            Get a database
+          parameters:
+          - in: path
+            schema:
+              type: integer
+            description: The database id
+            name: pk
+          responses:
+            200:
+              description: Database
+              content:
+                application/json:
+                  schema:
+                    type: object
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            422:
+              $ref: '#/components/responses/422'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        data = self.get_headless(pk, **kwargs)
+        try:
+            if ssh_tunnel := DatabaseDAO.get_ssh_tunnel(pk):
+                payload = data.json
+                payload["result"]["ssh_tunnel"] = ssh_tunnel.data
+                return payload
+            return data
+        except SupersetException as ex:
+            return self.response(ex.status, message=ex.message)
+
     @expose("/", methods=["POST"])
     @protect()
     @safe
@@ -280,6 +328,12 @@ def post(self) -> FlaskResponse:
             if new_model.driver:
                 item["driver"] = new_model.driver
 
+            # Return SSH Tunnel and hide passwords if any
+            if item.get("ssh_tunnel"):
+                item["ssh_tunnel"] = mask_password_info(
+                    new_model.ssh_tunnel  # pylint: disable=no-member
+                )
+
             return self.response(201, id=new_model.id, result=item)
         except DatabaseInvalidError as ex:
             return self.response_422(message=ex.normalized_messages())
@@ -361,6 +415,9 @@ def put(self, pk: int) -> Response:
             item["sqlalchemy_uri"] = changed_model.sqlalchemy_uri
             if changed_model.parameters:
                 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)
             return self.response(200, id=changed_model.id, result=item)
         except DatabaseNotFoundError:
             return self.response_404()
@@ -1206,3 +1263,57 @@ def validate_parameters(self) -> FlaskResponse:
         command = ValidateDatabaseParametersCommand(payload)
         command.run()
         return self.response(200, message="OK")
+
+    @expose("/<int:pk>/ssh_tunnel/", methods=["DELETE"])
+    @protect()
+    @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
@@ -31,6 +31,11 @@
 )
 from superset.databases.commands.test_connection import TestConnectionDatabaseCommand
 from superset.databases.dao import DatabaseDAO
+from superset.databases.ssh_tunnel.commands.create import CreateSSHTunnelCommand
+from superset.databases.ssh_tunnel.commands.exceptions import (
+    SSHTunnelCreateFailedError,
+    SSHTunnelInvalidError,
+)
 from superset.exceptions import SupersetErrorsException
 from superset.extensions import db, event_logger, security_manager
 
@@ -71,12 +76,35 @@ def run(self) -> Model:
             database = DatabaseDAO.create(self._properties, commit=False)
             database.set_sqlalchemy_uri(database.sqlalchemy_uri)
 
+            ssh_tunnel = None
+            if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
+                try:
+                    # So database.id is not None
+                    db.session.flush()
+                    ssh_tunnel = CreateSSHTunnelCommand(
+                        database.id, ssh_tunnel_properties
+                    ).run()
+                except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
+                    event_logger.log_with_context(
+                        action=f"db_creation_failed.{ex.__class__.__name__}",
+                        engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
+                    )
+                    # So we can show the original message
+                    raise ex
+                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],
+                    )
+                    raise DatabaseCreateFailedError() from ex
+
             # 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()

diff --git a/superset/databases/commands/test_connection.py b/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(),