feat: Add logging for ssh tunneling test_connection attempts (#22625)

superset/databases/commands/test_connection.py

@@ -48,6 +48,17 @@
 logger = logging.getLogger(__name__)
 
 
+def get_log_connection_action(
+    action: str, ssh_tunnel: Optional[Any], exc: Optional[Exception] = None
+) -> str:
+    action_modified = action
+    if exc:
+        action_modified += f".{exc.__class__.__name__}"
+    if ssh_tunnel:
+        action_modified += ".ssh_tunnel"
+    return action_modified
+
+
 class TestConnectionDatabaseCommand(BaseCommand):
     def __init__(self, data: Dict[str, Any]):
         self._properties = data.copy()
@@ -59,6 +70,7 @@ def run(self) -> None:  # pylint: disable=too-many-statements
         uri = self._properties.get("sqlalchemy_uri", "")
         if self._model and uri == self._model.safe_sqlalchemy_uri():
             uri = self._model.sqlalchemy_uri_decrypted
+        ssh_tunnel = self._properties.get("ssh_tunnel")
 
         # context for error messages
         url = make_url_safe(uri)
@@ -94,7 +106,7 @@ def run(self) -> None:  # pylint: disable=too-many-statements
             database.db_engine_spec.mutate_db_for_connection_test(database)
 
             # Generate tunnel if present in the properties
-            if ssh_tunnel := self._properties.get("ssh_tunnel"):
+            if ssh_tunnel:
                 # If there's an existing tunnel for that DB we need to use the stored
                 # password, private_key and private_key_password instead
                 if ssh_tunnel_id := ssh_tunnel.pop("id", None):
@@ -105,7 +117,7 @@ def run(self) -> None:  # pylint: disable=too-many-statements
                 ssh_tunnel = SSHTunnel(**ssh_tunnel)
 
             event_logger.log_with_context(
-                action="test_connection_attempt",
+                action=get_log_connection_action("test_connection_attempt", ssh_tunnel),
                 engine=database.db_engine_spec.__name__,
             )
 
@@ -147,13 +159,15 @@ def ping(engine: Engine) -> bool:
 
             # Log succesful connection test with engine
             event_logger.log_with_context(
-                action="test_connection_success",
+                action=get_log_connection_action("test_connection_success", ssh_tunnel),
                 engine=database.db_engine_spec.__name__,
             )
 
         except (NoSuchModuleError, ModuleNotFoundError) as ex:
             event_logger.log_with_context(
-                action=f"test_connection_error.{ex.__class__.__name__}",
+                action=get_log_connection_action(
+                    "test_connection_error", ssh_tunnel, ex
+                ),
                 engine=database.db_engine_spec.__name__,
             )
             raise DatabaseTestConnectionDriverError(
@@ -163,29 +177,37 @@ def ping(engine: Engine) -> bool:
             ) from ex
         except DBAPIError as ex:
             event_logger.log_with_context(
-                action=f"test_connection_error.{ex.__class__.__name__}",
+                action=get_log_connection_action(
+                    "test_connection_error", ssh_tunnel, ex
+                ),
                 engine=database.db_engine_spec.__name__,
             )
             # check for custom errors (wrong username, wrong password, etc)
             errors = database.db_engine_spec.extract_errors(ex, context)
             raise SupersetErrorsException(errors) from ex
         except SupersetSecurityException as ex:
             event_logger.log_with_context(
-                action=f"test_connection_error.{ex.__class__.__name__}",
+                action=get_log_connection_action(
+                    "test_connection_error", ssh_tunnel, ex
+                ),
                 engine=database.db_engine_spec.__name__,
             )
             raise DatabaseSecurityUnsafeError(message=str(ex)) from ex
         except SupersetTimeoutException as ex:
 
             event_logger.log_with_context(
-                action=f"test_connection_error.{ex.__class__.__name__}",
+                action=get_log_connection_action(
+                    "test_connection_error", ssh_tunnel, ex
+                ),
                 engine=database.db_engine_spec.__name__,
             )
             # bubble up the exception to return a 408
             raise ex
         except Exception as ex:
             event_logger.log_with_context(
-                action=f"test_connection_error.{ex.__class__.__name__}",
+                action=get_log_connection_action(
+                    "test_connection_error", ssh_tunnel, ex
+                ),
                 engine=database.db_engine_spec.__name__,
             )
             errors = database.db_engine_spec.extract_errors(ex, context)

superset/db_engine_specs/impala.py

@@ -92,7 +92,7 @@ def execute(
         cls,
         cursor: Any,
         query: str,
-        **kwargs: Any,  # pylint: disable=unused-argument
+        **kwargs: Any,
     ) -> None:
         try:
             cursor.execute_async(query)

superset/models/helpers.py

@@ -1326,7 +1326,7 @@ def get_sqla_col(self, col: Dict[str, Any]) -> Column:
         col = sa.column(label, type_=col_type)
         return self.make_sqla_column_compatible(col, label)
 
-    def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,unused-argument
+    def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements
         self,
         apply_fetch_values_predicate: bool = False,
         columns: Optional[List[Column]] = None,

tests/unit_tests/databases/commands/test_connection_test.py

@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from parameterized import parameterized
+
+from superset.databases.commands.test_connection import get_log_connection_action
+from superset.databases.ssh_tunnel.models import SSHTunnel
+
+
+@parameterized.expand(
+    [
+        ("foo", None, None, "foo"),
+        ("foo", SSHTunnel, None, "foo.ssh_tunnel"),
+        ("foo", SSHTunnel, Exception("oops"), "foo.Exception.ssh_tunnel"),
+    ],
+)
+def test_get_log_connection_action(action, tunnel, exc, expected_result):
+    assert expected_result == get_log_connection_action(action, tunnel, exc)