chore: Remove database ID dependency for SSH Tunnel creation (#26989)

superset/commands/database/create.py

@@ -41,6 +41,7 @@ from superset.daos.database import DatabaseDAO
 from superset.daos.exceptions import DAOCreateFailedError
 from superset.exceptions import SupersetErrorsException
 from superset.extensions import db, event_logger, security_manager
+from superset.models.core import Database
 
 logger = logging.getLogger(__name__)
 stats_logger = current_app.config["STATS_LOGGER"]
@@ -76,34 +77,20 @@ class CreateDatabaseCommand(BaseCommand):
             "{}",
         )
 
-        try:
-            database = DatabaseDAO.create(attributes=self._properties, commit=False)
-            database.set_sqlalchemy_uri(database.sqlalchemy_uri)
+        ssh_tunnel = None
+
+        try:
+            database = self._create_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()
-                except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
-                    event_logger.log_with_context(
-                        action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel",
-                        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__}.ssh_tunnel",
-                        engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
-                    )
-                    raise DatabaseCreateFailedError() from ex
+
+                ssh_tunnel = CreateSSHTunnelCommand(
+                    database, ssh_tunnel_properties
+                ).run()
+
+            db.session.commit()
 
             # adding a new database we always want to force refresh schema list
             schemas = database.get_all_schema_names(cache=False, ssh_tunnel=ssh_tunnel)
@@ -112,9 +99,23 @@ class CreateDatabaseCommand(BaseCommand):
                     "schema_access", security_manager.get_schema_perm(database, schema)
                 )
 
-            db.session.commit()
-
-        except DAOCreateFailedError as ex:
+        except (
+            SSHTunnelInvalidError,
+            SSHTunnelCreateFailedError,
+            SSHTunnelingNotEnabledError,
+        ) as ex:
+            db.session.rollback()
+            event_logger.log_with_context(
+                action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel",
+                engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
+            )
+            # So we can show the original message
+            raise ex
+        except (
+            DAOCreateFailedError,
+            DatabaseInvalidError,
+            Exception,
+        ) as ex:
             db.session.rollback()
             event_logger.log_with_context(
                 action=f"db_creation_failed.{ex.__class__.__name__}",
@@ -150,3 +151,8 @@ class CreateDatabaseCommand(BaseCommand):
                 )
             )
             raise exception
+
+    def _create_database(self) -> Database:
+        database = DatabaseDAO.create(attributes=self._properties, commit=False)
+        database.set_sqlalchemy_uri(database.sqlalchemy_uri)
+        return database