From 0a7016df2e56289888fc29e9c40da7ff5c09641f Mon Sep 17 00:00:00 2001
From: "Hugh A. Miles II" <hughmil3s@gmail.com>
Date: Tue, 28 Feb 2023 08:21:40 -0800
Subject: [PATCH] chore: Add statsd logger for ssh tunneling creation (#23225)

---
 superset/databases/commands/create.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/superset/databases/commands/create.py b/superset/databases/commands/create.py
index 0ed2354960..eb0582a980 100644
--- a/superset/databases/commands/create.py
+++ b/superset/databases/commands/create.py
@@ -17,6 +17,7 @@
 import logging
 from typing import Any, Dict, List, Optional
 
+from flask import current_app
 from flask_appbuilder.models.sqla import Model
 from marshmallow import ValidationError
 
@@ -42,6 +43,7 @@ from superset.exceptions import SupersetErrorsException
 from superset.extensions import db, event_logger, security_manager
 
 logger = logging.getLogger(__name__)
+stats_logger = current_app.config["STATS_LOGGER"]
 
 
 class CreateDatabaseCommand(BaseCommand):
@@ -91,14 +93,14 @@ class CreateDatabaseCommand(BaseCommand):
                     ).run()
                 except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
                     event_logger.log_with_context(
-                        action=f"db_creation_failed.{ex.__class__.__name__}",
+                        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__}",
+                        action=f"db_creation_failed.{ex.__class__.__name__}.ssh_tunnel",
                         engine=self._properties.get("sqlalchemy_uri", "").split(":")[0],
                     )
                     raise DatabaseCreateFailedError() from ex
@@ -111,6 +113,7 @@ class CreateDatabaseCommand(BaseCommand):
                 )
 
             db.session.commit()
+
         except DAOCreateFailedError as ex:
             db.session.rollback()
             event_logger.log_with_context(
@@ -118,6 +121,10 @@ class CreateDatabaseCommand(BaseCommand):
                 engine=database.db_engine_spec.__name__,
             )
             raise DatabaseCreateFailedError() from ex
+
+        if ssh_tunnel:
+            stats_logger.incr("db_creation_success.ssh_tunnel")
+
         return database
 
     def validate(self) -> None: