removed dup import, renamed func, changed a filter_null_values to lamda

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -37,7 +37,7 @@ class TelemetryHelper:
     _DRIVER_SYSTEM_CONFIGURATION = None
 
     @classmethod
-    def getDriverSystemConfiguration(cls) -> DriverSystemConfiguration:
+    def get_driver_system_configuration(cls) -> DriverSystemConfiguration:
         if cls._DRIVER_SYSTEM_CONFIGURATION is None:
             from databricks.sql import __version__
 
@@ -259,7 +259,7 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
                 entry=FrontendLogEntry(
                     sql_driver_log=TelemetryEvent(
                         session_id=self._connection_uuid,
-                        system_configuration=TelemetryHelper.getDriverSystemConfiguration(),
+                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
                         driver_connection_params=self._driver_connection_params,
                     )
                 ),
@@ -286,7 +286,7 @@ def export_failure_log(self, error_name, error_message):
                 entry=FrontendLogEntry(
                     sql_driver_log=TelemetryEvent(
                         session_id=self._connection_uuid,
-                        system_configuration=TelemetryHelper.getDriverSystemConfiguration(),
+                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
                         driver_connection_params=self._driver_connection_params,
                         error_info=error_info,
                     )
@@ -340,8 +340,6 @@ def _initialize(cls):
     def _install_exception_hook(cls):
         """Install global exception handler for unhandled exceptions"""
         if not cls._excepthook_installed:
-            import sys
-
             cls._original_excepthook = sys.excepthook
             sys.excepthook = cls._handle_unhandled_exception
             cls._excepthook_installed = True

src/databricks/sql/telemetry/utils.py

@@ -16,21 +16,14 @@ def default(self, obj):
         return super().default(obj)
 
 
-def filter_none_values(data):
-    """
-    Recursively remove None values from dictionaries.
-    This reduces telemetry payload size by excluding null fields.
-    """
-    if isinstance(data, dict):
-        return {k: filter_none_values(v) for k, v in data.items() if v is not None}
-    else:
-        return data
-
-
 def to_json_compact(dataclass_obj):
     """
     Convert a dataclass to JSON string, excluding None values.
     """
-    data_dict = asdict(dataclass_obj)
-    filtered_dict = filter_none_values(data_dict)
-    return json.dumps(filtered_dict, cls=EnumEncoder)
+    return json.dumps(
+        asdict(
+            dataclass_obj,
+            dict_factory=lambda data: {k: v for k, v in data if v is not None},
+        ),
+        cls=EnumEncoder,
+    )

tests/unit/test_telemetry.py

@@ -98,7 +98,7 @@ class TestTelemetryClient:
     """Tests for the TelemetryClient class."""
 
     @patch("databricks.sql.telemetry.telemetry_client.TelemetryFrontendLog")
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHelper.getDriverSystemConfiguration")
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHelper.get_driver_system_configuration")
     @patch("databricks.sql.telemetry.telemetry_client.uuid.uuid4")
     @patch("databricks.sql.telemetry.telemetry_client.time.time")
     def test_export_initial_telemetry_log(
@@ -134,7 +134,7 @@ def test_export_initial_telemetry_log(
         client.export_event.assert_called_once_with(mock_frontend_log.return_value)
 
     @patch("databricks.sql.telemetry.telemetry_client.TelemetryFrontendLog")
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHelper.getDriverSystemConfiguration")
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHelper.get_driver_system_configuration")
     @patch("databricks.sql.telemetry.telemetry_client.DriverErrorInfo")
     @patch("databricks.sql.telemetry.telemetry_client.uuid.uuid4")
     @patch("databricks.sql.telemetry.telemetry_client.time.time")