fix: SSH Tunnel configuration settings (#27186)

(cherry picked from commit 89e89de341)

superset/commands/database/create.py

@@ -19,6 +19,7 @@ from typing import Any, Optional
 
 from flask import current_app
 from flask_appbuilder.models.sqla import Model
+from flask_babel import gettext as _
 from marshmallow import ValidationError
 
 from superset import is_feature_enabled
@@ -33,6 +34,7 @@ from superset.commands.database.exceptions import (
 from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
 from superset.commands.database.ssh_tunnel.exceptions import (
     SSHTunnelCreateFailedError,
+    SSHTunnelDatabasePortError,
     SSHTunnelingNotEnabledError,
     SSHTunnelInvalidError,
 )
@@ -57,7 +59,11 @@ class CreateDatabaseCommand(BaseCommand):
         try:
             # Test connection before starting create transaction
             TestConnectionDatabaseCommand(self._properties).run()
-        except (SupersetErrorsException, SSHTunnelingNotEnabledError) as ex:
+        except (
+            SupersetErrorsException,
+            SSHTunnelingNotEnabledError,
+            SSHTunnelDatabasePortError,
+        ) as ex:
             event_logger.log_with_context(
                 action=f"db_creation_failed.{ex.__class__.__name__}",
                 engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
@@ -103,6 +109,7 @@ class CreateDatabaseCommand(BaseCommand):
             SSHTunnelInvalidError,
             SSHTunnelCreateFailedError,
             SSHTunnelingNotEnabledError,
+            SSHTunnelDatabasePortError,
         ) as ex:
             db.session.rollback()
             event_logger.log_with_context(
@@ -140,6 +147,7 @@ class CreateDatabaseCommand(BaseCommand):
             # Check database_name uniqueness
             if not DatabaseDAO.validate_uniqueness(database_name):
                 exceptions.append(DatabaseExistsValidationError())
+
         if exceptions:
             exception = DatabaseInvalidError()
             exception.extend(exceptions)

superset/commands/database/ssh_tunnel/create.py

@@ -23,11 +23,13 @@ from marshmallow import ValidationError
 from superset.commands.base import BaseCommand
 from superset.commands.database.ssh_tunnel.exceptions import (
     SSHTunnelCreateFailedError,
+    SSHTunnelDatabasePortError,
     SSHTunnelInvalidError,
     SSHTunnelRequiredFieldValidationError,
 )
 from superset.daos.database import SSHTunnelDAO
 from superset.daos.exceptions import DAOCreateFailedError
+from superset.databases.utils import make_url_safe
 from superset.extensions import event_logger
 from superset.models.core import Database
 
@@ -35,9 +37,12 @@ logger = logging.getLogger(__name__)
 
 
 class CreateSSHTunnelCommand(BaseCommand):
+    _database: Database
+
     def __init__(self, database: Database, data: dict[str, Any]):
         self._properties = data.copy()
         self._properties["database"] = database
+        self._database = database
 
     def run(self) -> Model:
         try:
@@ -57,16 +62,22 @@ class CreateSSHTunnelCommand(BaseCommand):
         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")
+        password: Optional[str] = self._properties.get("password")
         private_key: Optional[str] = self._properties.get("private_key")
         private_key_password: Optional[str] = self._properties.get(
             "private_key_password"
         )
+        url = make_url_safe(self._database.sqlalchemy_uri)
+        if not url.port:
+            raise SSHTunnelDatabasePortError()
         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 not private_key and not password:
+            exceptions.append(SSHTunnelRequiredFieldValidationError("password"))
         if private_key_password and private_key is None:
             exceptions.append(SSHTunnelRequiredFieldValidationError("private_key"))
         if exceptions:

superset/commands/database/ssh_tunnel/exceptions.py

@@ -38,6 +38,10 @@ class SSHTunnelInvalidError(CommandInvalidError):
     message = _("SSH Tunnel parameters are invalid.")
 
 
+class SSHTunnelDatabasePortError(CommandInvalidError):
+    message = _("A database port is required when connecting via SSH Tunnel.")
+
+
 class SSHTunnelUpdateFailedError(UpdateFailedError):
     message = _("SSH Tunnel could not be updated.")
 

superset/commands/database/ssh_tunnel/update.py

@@ -21,6 +21,7 @@ from flask_appbuilder.models.sqla import Model
 
 from superset.commands.base import BaseCommand
 from superset.commands.database.ssh_tunnel.exceptions import (
+    SSHTunnelDatabasePortError,
     SSHTunnelInvalidError,
     SSHTunnelNotFoundError,
     SSHTunnelRequiredFieldValidationError,
@@ -29,6 +30,7 @@ from superset.commands.database.ssh_tunnel.exceptions import (
 from superset.daos.database import SSHTunnelDAO
 from superset.daos.exceptions import DAOUpdateFailedError
 from superset.databases.ssh_tunnel.models import SSHTunnel
+from superset.databases.utils import make_url_safe
 
 logger = logging.getLogger(__name__)
 
@@ -39,20 +41,33 @@ class UpdateSSHTunnelCommand(BaseCommand):
         self._model_id = model_id
         self._model: Optional[SSHTunnel] = None
 
-    def run(self) -> Model:
+    def run(self) -> Optional[Model]:
         self.validate()
         try:
-            if self._model is not None:  # So we dont get incompatible types error
-                tunnel = SSHTunnelDAO.update(self._model, self._properties)
+            if self._model is None:
+                return None
+
+            # unset password if private key is provided
+            if self._properties.get("private_key"):
+                self._properties["password"] = None
+
+            # unset private key and password if password is provided
+            if self._properties.get("password"):
+                self._properties["private_key"] = None
+                self._properties["private_key_password"] = None
+
+            tunnel = SSHTunnelDAO.update(self._model, self._properties)
+            return tunnel
         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()
+
+        url = make_url_safe(self._model.database.sqlalchemy_uri)
         private_key: Optional[str] = self._properties.get("private_key")
         private_key_password: Optional[str] = self._properties.get(
             "private_key_password"
@@ -61,3 +76,5 @@ class UpdateSSHTunnelCommand(BaseCommand):
             raise SSHTunnelInvalidError(
                 exceptions=[SSHTunnelRequiredFieldValidationError("private_key")]
             )
+        if not url.port:
+            raise SSHTunnelDatabasePortError()

superset/commands/database/test_connection.py

@@ -32,7 +32,10 @@ from superset.commands.database.exceptions import (
     DatabaseTestConnectionDriverError,
     DatabaseTestConnectionUnexpectedError,
 )
-from superset.commands.database.ssh_tunnel.exceptions import SSHTunnelingNotEnabledError
+from superset.commands.database.ssh_tunnel.exceptions import (
+    SSHTunnelDatabasePortError,
+    SSHTunnelingNotEnabledError,
+)
 from superset.daos.database import DatabaseDAO, SSHTunnelDAO
 from superset.databases.ssh_tunnel.models import SSHTunnel
 from superset.databases.utils import make_url_safe
@@ -61,20 +64,22 @@ def get_log_connection_action(
 
 
 class TestConnectionDatabaseCommand(BaseCommand):
+    _model: Optional[Database] = None
+    _context: dict[str, Any]
+    _uri: str
+
     def __init__(self, data: dict[str, Any]):
         self._properties = data.copy()
-        self._model: Optional[Database] = None
 
-    def run(self) -> None:  # pylint: disable=too-many-statements, too-many-branches
-        self.validate()
-        ex_str = ""
+        if (database_name := self._properties.get("database_name")) is not None:
+            self._model = DatabaseDAO.get_database_by_name(database_name)
+
         uri = self._properties.get("sqlalchemy_uri", "")
         if self._model and uri == self._model.safe_sqlalchemy_uri():
             uri = self._model.sqlalchemy_uri_decrypted
-        ssh_tunnel = self._properties.get("ssh_tunnel")
 
-        # context for error messages
         url = make_url_safe(uri)
+
         context = {
             "hostname": url.host,
             "password": url.password,
@@ -83,6 +88,14 @@ class TestConnectionDatabaseCommand(BaseCommand):
             "database": url.database,
         }
 
+        self._context = context
+        self._uri = uri
+
+    def run(self) -> None:  # pylint: disable=too-many-statements
+        self.validate()
+        ex_str = ""
+        ssh_tunnel = self._properties.get("ssh_tunnel")
+
         serialized_encrypted_extra = self._properties.get(
             "masked_encrypted_extra",
             "{}",
@@ -103,15 +116,12 @@ class TestConnectionDatabaseCommand(BaseCommand):
                 encrypted_extra=serialized_encrypted_extra,
             )
 
-            database.set_sqlalchemy_uri(uri)
+            database.set_sqlalchemy_uri(self._uri)
             database.db_engine_spec.mutate_db_for_connection_test(database)
 
             # Generate tunnel if present in the properties
             if ssh_tunnel:
-                if not is_feature_enabled("SSH_TUNNELING"):
-                    raise SSHTunnelingNotEnabledError()
-                # If there's an existing tunnel for that DB we need to use the stored
-                # password, private_key and private_key_password instead
+                # unmask password while allowing for updated values
                 if ssh_tunnel_id := ssh_tunnel.pop("id", None):
                     if existing_ssh_tunnel := SSHTunnelDAO.find_by_id(ssh_tunnel_id):
                         ssh_tunnel = unmask_password_info(
@@ -186,7 +196,7 @@ class TestConnectionDatabaseCommand(BaseCommand):
                 engine=database.db_engine_spec.__name__,
             )
             # check for custom errors (wrong username, wrong password, etc)
-            errors = database.db_engine_spec.extract_errors(ex, context)
+            errors = database.db_engine_spec.extract_errors(ex, self._context)
             raise SupersetErrorsException(errors) from ex
         except SupersetSecurityException as ex:
             event_logger.log_with_context(
@@ -221,9 +231,12 @@ class TestConnectionDatabaseCommand(BaseCommand):
                 ),
                 engine=database.db_engine_spec.__name__,
             )
-            errors = database.db_engine_spec.extract_errors(ex, context)
+            errors = database.db_engine_spec.extract_errors(ex, self._context)
             raise DatabaseTestConnectionUnexpectedError(errors) from ex
 
     def validate(self) -> None:
-        if (database_name := self._properties.get("database_name")) is not None:
-            self._model = DatabaseDAO.get_database_by_name(database_name)
+        if self._properties.get("ssh_tunnel"):
+            if not is_feature_enabled("SSH_TUNNELING"):
+                raise SSHTunnelingNotEnabledError()
+            if not self._context.get("port"):
+                raise SSHTunnelDatabasePortError()

superset/commands/database/update.py

@@ -18,6 +18,7 @@ import logging
 from typing import Any, Optional
 
 from flask_appbuilder.models.sqla import Model
+from flask_babel import gettext as _
 from marshmallow import ValidationError
 
 from superset import is_feature_enabled
@@ -30,8 +31,11 @@ from superset.commands.database.exceptions import (
     DatabaseUpdateFailedError,
 )
 from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
+from superset.commands.database.ssh_tunnel.delete import DeleteSSHTunnelCommand
 from superset.commands.database.ssh_tunnel.exceptions import (
     SSHTunnelCreateFailedError,
+    SSHTunnelDatabasePortError,
+    SSHTunnelDeleteFailedError,
     SSHTunnelingNotEnabledError,
     SSHTunnelInvalidError,
     SSHTunnelUpdateFailedError,
@@ -47,15 +51,21 @@ logger = logging.getLogger(__name__)
 
 
 class UpdateDatabaseCommand(BaseCommand):
+    _model: Optional[Database]
+
     def __init__(self, model_id: int, data: dict[str, Any]):
         self._properties = data.copy()
         self._model_id = model_id
         self._model: Optional[Database] = None
 
-    def run(self) -> Model:
-        self.validate()
+    def run(self) -> Model:  # pylint: disable=too-many-statements, too-many-branches
+        self._model = DatabaseDAO.find_by_id(self._model_id)
+
         if not self._model:
             raise DatabaseNotFoundError()
+
+        self.validate()
+
         old_database_name = self._model.database_name
 
         # unmask ``encrypted_extra``
@@ -70,36 +80,59 @@ class UpdateDatabaseCommand(BaseCommand):
             database = DatabaseDAO.update(self._model, self._properties, commit=False)
             database.set_sqlalchemy_uri(database.sqlalchemy_uri)
 
-            if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
+            ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id)
+
+            if "ssh_tunnel" in self._properties:
                 if not is_feature_enabled("SSH_TUNNELING"):
                     db.session.rollback()
                     raise SSHTunnelingNotEnabledError()
-                existing_ssh_tunnel_model = DatabaseDAO.get_ssh_tunnel(database.id)
-                if existing_ssh_tunnel_model is None:
-                    # We couldn't found an existing tunnel so we need to create one
+
+                if self._properties.get("ssh_tunnel") is None and ssh_tunnel:
+                    # We need to remove the existing tunnel
                     try:
-                        CreateSSHTunnelCommand(database, ssh_tunnel_properties).run()
-                    except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
-                        # So we can show the original message
-                        raise ex
-                    except Exception as ex:
-                        raise DatabaseUpdateFailedError() from ex
-                else:
-                    # We found an existing tunnel so we need to update it
-                    try:
-                        UpdateSSHTunnelCommand(
-                            existing_ssh_tunnel_model.id, ssh_tunnel_properties
-                        ).run()
-                    except (SSHTunnelInvalidError, SSHTunnelUpdateFailedError) as ex:
-                        # So we can show the original message
+                        DeleteSSHTunnelCommand(ssh_tunnel.id).run()
+                        ssh_tunnel = None
+                    except SSHTunnelDeleteFailedError as ex:
                         raise ex
                     except Exception as ex:
                         raise DatabaseUpdateFailedError() from ex
 
+                if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
+                    if ssh_tunnel is None:
+                        # We couldn't found an existing tunnel so we need to create one
+                        try:
+                            ssh_tunnel = CreateSSHTunnelCommand(
+                                database, ssh_tunnel_properties
+                            ).run()
+                        except (
+                            SSHTunnelInvalidError,
+                            SSHTunnelCreateFailedError,
+                            SSHTunnelDatabasePortError,
+                        ) as ex:
+                            # So we can show the original message
+                            raise ex
+                        except Exception as ex:
+                            raise DatabaseUpdateFailedError() from ex
+                    else:
+                        # We found an existing tunnel so we need to update it
+                        try:
+                            ssh_tunnel_id = ssh_tunnel.id
+                            ssh_tunnel = UpdateSSHTunnelCommand(
+                                ssh_tunnel_id, ssh_tunnel_properties
+                            ).run()
+                        except (
+                            SSHTunnelInvalidError,
+                            SSHTunnelUpdateFailedError,
+                            SSHTunnelDatabasePortError,
+                        ) as ex:
+                            # So we can show the original message
+                            raise ex
+                        except Exception as ex:
+                            raise DatabaseUpdateFailedError() from ex
+
             # adding a new database we always want to force refresh schema list
             # TODO Improve this simplistic implementation for catching DB conn fails
             try:
-                ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id)
                 schemas = database.get_all_schema_names(ssh_tunnel=ssh_tunnel)
             except Exception as ex:
                 db.session.rollback()
@@ -167,10 +200,6 @@ class UpdateDatabaseCommand(BaseCommand):
 
     def validate(self) -> None:
         exceptions: list[ValidationError] = []
-        # Validate/populate model exists
-        self._model = DatabaseDAO.find_by_id(self._model_id)
-        if not self._model:
-            raise DatabaseNotFoundError()
         database_name: Optional[str] = self._properties.get("database_name")
         if database_name:
             # Check database_name uniqueness