From d041648ad45e881ea38ce6131e3e36e438bae748 Mon Sep 17 00:00:00 2001
From: "Hugh A. Miles II" <hughmil3s@gmail.com>
Date: Thu, 6 Jul 2023 15:43:55 -0400
Subject: [PATCH] fix(ssh): Editting Database w/ SSH Tunneling (#24552)

---
 superset/databases/commands/update.py | 56 ++++++++++++++-------------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/superset/databases/commands/update.py b/superset/databases/commands/update.py
index ea49801aac..d8d86c6d2d 100644
--- a/superset/databases/commands/update.py
+++ b/superset/databases/commands/update.py
@@ -69,10 +69,38 @@ class UpdateDatabaseCommand(BaseCommand):
         try:
             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"):
+                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
+                    try:
+                        CreateSSHTunnelCommand(database.id, 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
+                        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:
-                schemas = database.get_all_schema_names()
+                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()
                 raise DatabaseConnectionFailedError() from ex
@@ -104,32 +132,6 @@ class UpdateDatabaseCommand(BaseCommand):
                     "schema_access", security_manager.get_schema_perm(database, schema)
                 )
 
-            if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
-                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
-                    try:
-                        CreateSSHTunnelCommand(database.id, 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
-                        raise ex
-                    except Exception as ex:
-                        raise DatabaseUpdateFailedError() from ex
-
             db.session.commit()
 
         except (DAOUpdateFailedError, DAOCreateFailedError) as ex: