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

Co-authored-by: Antonio Rivero Martinez <38889534+Antonio-RiveroMartnez@users.noreply.github.com>
Co-authored-by: Elizabeth Thompson <eschutho@gmail.com>

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,92 @@
+# 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 db, 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:
+        try:
+            # Start nested transaction since we are always creating the tunnel
+            # through a DB command (Create or Update). Without this, we cannot
+            # safely rollback changes to databases if any, i.e, things like
+            # test_do_not_create_database_if_ssh_tunnel_creation_fails test will fail
+            db.session.begin_nested()
+            self.validate()
+            tunnel = SSHTunnelDAO.create(self._properties, commit=False)
+        except DAOCreateFailedError as ex:
+            # Rollback nested transaction
+            db.session.rollback()
+            raise SSHTunnelCreateFailedError() from ex
+        except SSHTunnelInvalidError as ex:
+            # Rollback nested transaction
+            db.session.rollback()
+            raise 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()

superset/databases/ssh_tunnel/commands/exceptions.py

@@ -0,0 +1,54 @@
+# 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 flask_babel import lazy_gettext as _
+from marshmallow import ValidationError
+
+from superset.commands.exceptions import (
+    CommandException,
+    CommandInvalidError,
+    DeleteFailedError,
+    UpdateFailedError,
+)
+
+
+class SSHTunnelDeleteFailedError(DeleteFailedError):
+    message = _("SSH Tunnel could not be deleted.")
+
+
+class SSHTunnelNotFoundError(CommandException):
+    status = 404
+    message = _("SSH Tunnel not found.")
+
+
+class SSHTunnelInvalidError(CommandInvalidError):
+    message = _("SSH Tunnel parameters are invalid.")
+
+
+class SSHTunnelUpdateFailedError(UpdateFailedError):
+    message = _("SSH Tunnel could not be updated.")
+
+
+class SSHTunnelCreateFailedError(CommandException):
+    message = _("Creating SSH Tunnel failed for an unknown reason")
+
+
+class SSHTunnelRequiredFieldValidationError(ValidationError):
+    def __init__(self, field_name: str) -> None:
+        super().__init__(
+            [_("Field is required")],
+            field_name=field_name,
+        )

superset/databases/ssh_tunnel/commands/update.py

@@ -0,0 +1,62 @@
+# 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, Optional
+
+from flask_appbuilder.models.sqla import Model
+
+from superset.commands.base import BaseCommand
+from superset.dao.exceptions import DAOUpdateFailedError
+from superset.databases.ssh_tunnel.commands.exceptions import (
+    SSHTunnelInvalidError,
+    SSHTunnelNotFoundError,
+    SSHTunnelRequiredFieldValidationError,
+    SSHTunnelUpdateFailedError,
+)
+from superset.databases.ssh_tunnel.dao import SSHTunnelDAO
+from superset.databases.ssh_tunnel.models import SSHTunnel
+
+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
+
+    def run(self) -> Model:
+        self.validate()
+        try:
+            tunnel = SSHTunnelDAO.update(self._model, self._properties)
+        except DAOUpdateFailedError as ex:
+            raise SSHTunnelUpdateFailedError() from ex
+        return tunnel
+
+    def validate(self) -> None:
+        # Validate/populate model exists
+        self._model = SSHTunnelDAO.find_by_id(self._model_id)
+        if not self._model:
+            raise SSHTunnelNotFoundError()
+        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:
+            exception = SSHTunnelInvalidError()
+            exception.add(SSHTunnelRequiredFieldValidationError("private_key"))
+            raise exception

superset/databases/ssh_tunnel/dao.py

@@ -0,0 +1,26 @@
+# 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 superset.dao.base import BaseDAO
+from superset.databases.ssh_tunnel.models import SSHTunnel
+
+logger = logging.getLogger(__name__)
+
+
+class SSHTunnelDAO(BaseDAO):
+    model_cls = SSHTunnel

superset/databases/ssh_tunnel/models.py

@@ -0,0 +1,76 @@
+# 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 typing import Any, Dict
+
+import sqlalchemy as sa
+from flask import current_app
+from flask_appbuilder import Model
+from sqlalchemy.orm import backref, relationship
+from sqlalchemy_utils import EncryptedType
+
+from superset.models.core import Database
+from superset.models.helpers import (
+    AuditMixinNullable,
+    ExtraJSONMixin,
+    ImportExportMixin,
+)
+
+app_config = current_app.config
+
+
+class SSHTunnel(Model, AuditMixinNullable, ExtraJSONMixin, ImportExportMixin):
+    """
+    A ssh tunnel configuration in a database.
+    """
+
+    __tablename__ = "ssh_tunnels"
+
+    id = sa.Column(sa.Integer, primary_key=True)
+    database_id = sa.Column(
+        sa.Integer, sa.ForeignKey("dbs.id"), nullable=False, unique=True
+    )
+    database: Database = relationship(
+        "Database",
+        backref=backref("ssh_tunnels", uselist=False, cascade="all, delete-orphan"),
+        foreign_keys=[database_id],
+    )
+
+    server_address = sa.Column(sa.Text)
+    server_port = sa.Column(sa.Integer)
+    username = sa.Column(EncryptedType(sa.String, app_config["SECRET_KEY"]))
+
+    # basic authentication
+    password = sa.Column(
+        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
+    )
+
+    # password protected pkey authentication
+    private_key = sa.Column(
+        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
+    )
+    private_key_password = sa.Column(
+        EncryptedType(sa.String, app_config["SECRET_KEY"]), nullable=True
+    )
+
+    @property
+    def data(self) -> Dict[str, Any]:
+        return {
+            "server_address": self.server_address,
+            "server_port": self.server_port,
+            "username": self.username,
+        }