chore: Add statsd logger for ssh tunneling creation (#23225)

This commit is contained in:
Hugh A. Miles II 2023-02-28 08:21:40 -08:00 committed by GitHub
parent b479e93b49
commit 0a7016df2e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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: