(superset) 01/01: Change order of operations

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

diegopucci pushed a commit to branch diego/ch78526/ssh-tunnel-db-id-dep
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 38f009e7a3ccf09aecd6a4fab8d5782ca69c284d
Author: geido <di...@gmail.com>
AuthorDate: Fri Feb 2 08:53:52 2024 +0100

    Change order of operations
---
 superset/commands/database/create.py            | 6 ++++--
 superset/commands/database/ssh_tunnel/create.py | 9 ---------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/superset/commands/database/create.py b/superset/commands/database/create.py
index a012e9b2a5..b6ad4772dc 100644
--- a/superset/commands/database/create.py
+++ b/superset/commands/database/create.py
@@ -80,14 +80,14 @@ class CreateDatabaseCommand(BaseCommand):
             database = DatabaseDAO.create(attributes=self._properties, commit=False)
             database.set_sqlalchemy_uri(database.sqlalchemy_uri)
 
+            db.session.add(database)
+
             ssh_tunnel = None
             if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
                 if not is_feature_enabled("SSH_TUNNELING"):
                     db.session.rollback()
                     raise SSHTunnelingNotEnabledError()
                 try:
-                    # So database.id is not None
-                    db.session.flush()
                     ssh_tunnel = CreateSSHTunnelCommand(
                         database.id, ssh_tunnel_properties
                     ).run()
@@ -96,6 +96,7 @@ class CreateDatabaseCommand(BaseCommand):
                         action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel",
                         engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
                     )
+                    db.session.rollback()
                     # So we can show the original message
                     raise ex
                 except Exception as ex:
@@ -103,6 +104,7 @@ class CreateDatabaseCommand(BaseCommand):
                         action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel",
                         engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
                     )
+                    db.session.rollback()
                     raise DatabaseCreateFailedError() from ex
 
             # adding a new database we always want to force refresh schema list
diff --git a/superset/commands/database/ssh_tunnel/create.py b/superset/commands/database/ssh_tunnel/create.py
index 07209f010b..fc4d667e35 100644
--- a/superset/commands/database/ssh_tunnel/create.py
+++ b/superset/commands/database/ssh_tunnel/create.py
@@ -40,20 +40,11 @@ class CreateSSHTunnelCommand(BaseCommand):
 
     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()
             return SSHTunnelDAO.create(attributes=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
 
     def validate(self) -> None: