JsonSerializableMixin

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

Author: saishreeeee

src/databricks/sql/telemetry/models/event.py

@@ -1,5 +1,4 @@
-import json
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass
 from databricks.sql.telemetry.models.enums import (
     AuthMech,
     AuthFlow,
@@ -9,11 +8,11 @@
     ExecutionResultFormat,
 )
 from typing import Optional
-from databricks.sql.telemetry.utils import to_json_compact
+from databricks.sql.telemetry.utils import JsonSerializableMixin
 
 
 @dataclass
-class HostDetails:
+class HostDetails(JsonSerializableMixin):
     """
     Represents the host connection details for a Databricks workspace.
 
@@ -25,12 +24,9 @@ class HostDetails:
     host_url: str
     port: int
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class DriverConnectionParameters:
+class DriverConnectionParameters(JsonSerializableMixin):
     """
     Contains all connection parameters used to establish a connection to Databricks SQL.
     This includes authentication details, host information, and connection settings.
@@ -51,12 +47,9 @@ class DriverConnectionParameters:
     auth_flow: Optional[AuthFlow] = None
     socket_timeout: Optional[int] = None
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class DriverSystemConfiguration:
+class DriverSystemConfiguration(JsonSerializableMixin):
     """
     Contains system-level configuration information about the client environment.
     This includes details about the operating system, runtime, and driver version.
@@ -87,12 +80,9 @@ class DriverSystemConfiguration:
     client_app_name: Optional[str] = None
     locale_name: Optional[str] = None
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class DriverVolumeOperation:
+class DriverVolumeOperation(JsonSerializableMixin):
     """
     Represents a volume operation performed by the driver.
     Used for tracking volume-related operations in telemetry.
@@ -105,12 +95,9 @@ class DriverVolumeOperation:
     volume_operation_type: DriverVolumeOperationType
     volume_path: str
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class DriverErrorInfo:
+class DriverErrorInfo(JsonSerializableMixin):
     """
     Contains detailed information about errors that occur during driver operations.
     Used for error tracking and debugging in telemetry.
@@ -123,12 +110,9 @@ class DriverErrorInfo:
     error_name: str
     stack_trace: str
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class SqlExecutionEvent:
+class SqlExecutionEvent(JsonSerializableMixin):
     """
     Represents a SQL query execution event.
     Contains details about the query execution, including type, compression, and result format.
@@ -145,12 +129,9 @@ class SqlExecutionEvent:
     execution_result: ExecutionResultFormat
     retry_count: int
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class TelemetryEvent:
+class TelemetryEvent(JsonSerializableMixin):
     """
     Main telemetry event class that aggregates all telemetry data.
     Contains information about the session, system configuration, connection parameters,
@@ -177,6 +158,3 @@ class TelemetryEvent:
     sql_operation: Optional[SqlExecutionEvent] = None
     error_info: Optional[DriverErrorInfo] = None
     operation_latency_ms: Optional[int] = None
-
-    def to_json(self):
-        return to_json_compact(self)

src/databricks/sql/telemetry/models/frontend_logs.py

@@ -1,12 +1,11 @@
-import json
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass
 from databricks.sql.telemetry.models.event import TelemetryEvent
-from databricks.sql.telemetry.utils import to_json_compact
+from databricks.sql.telemetry.utils import JsonSerializableMixin
 from typing import Optional
 
 
 @dataclass
-class TelemetryClientContext:
+class TelemetryClientContext(JsonSerializableMixin):
     """
     Contains client-side context information for telemetry events.
     This includes timestamp and user agent information for tracking when and how the client is being used.
@@ -19,12 +18,9 @@ class TelemetryClientContext:
     timestamp_millis: int
     user_agent: str
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class FrontendLogContext:
+class FrontendLogContext(JsonSerializableMixin):
     """
     Wrapper for client context information in frontend logs.
     Provides additional context about the client environment for telemetry events.
@@ -35,12 +31,9 @@ class FrontendLogContext:
 
     client_context: TelemetryClientContext
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class FrontendLogEntry:
+class FrontendLogEntry(JsonSerializableMixin):
     """
     Contains the actual telemetry event data in a frontend log.
     Wraps the SQL driver log information for frontend processing.
@@ -51,12 +44,9 @@ class FrontendLogEntry:
 
     sql_driver_log: TelemetryEvent
 
-    def to_json(self):
-        return to_json_compact(self)
-
 
 @dataclass
-class TelemetryFrontendLog:
+class TelemetryFrontendLog(JsonSerializableMixin):
     """
     Main container for frontend telemetry data.
     Aggregates workspace information, event ID, context, and the actual log entry.
@@ -73,6 +63,3 @@ class TelemetryFrontendLog:
     context: FrontendLogContext
     entry: FrontendLogEntry
     workspace_id: Optional[int] = None
-
-    def to_json(self):
-        return to_json_compact(self)

src/databricks/sql/telemetry/telemetry_client.py

@@ -265,7 +265,8 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
                 ),
             )
 
-            self.export_event(telemetry_frontend_log)
+            self._export_event(telemetry_frontend_log)
+
         except Exception as e:
             logger.debug("Failed to export initial telemetry log: %s", e)
 
@@ -292,7 +293,7 @@ def export_failure_log(self, error_name, error_message):
                     )
                 ),
             )
-            self.export_event(telemetry_frontend_log)
+            self._export_event(telemetry_frontend_log)
         except Exception as e:
             logger.debug("Failed to export failure log: %s", e)
 
@@ -347,9 +348,9 @@ def _handle_unhandled_exception(cls, exc_type, exc_value, exc_traceback):
         """Handle unhandled exceptions by sending telemetry and flushing thread pool"""
         logger.debug("Handling unhandled exception: %s", exc_type.__name__)
 
-        # Flush existing thread pool work and wait for completion
-        for uuid, _ in cls._clients.items():
-            cls.close(uuid)
+        clients_to_close = list(cls._clients.values())
+        for client in clients_to_close:
+            client.close()
 
         # Call the original exception handler to maintain normal behavior
         if cls._original_excepthook:

src/databricks/sql/telemetry/utils.py

@@ -1,6 +1,25 @@
 import json
 from enum import Enum
 from dataclasses import asdict
+from abc import ABC
+from typing import Any
+
+
+class JsonSerializableMixin(ABC):
+    """Mixin class to provide JSON serialization capabilities to dataclasses."""
+
+    def to_json(self) -> str:
+        """
+        Convert the object to a JSON string, excluding None values.
+        Handles Enum serialization and filters out None values from the output.
+        """
+        return json.dumps(
+            asdict(
+                self,
+                dict_factory=lambda data: {k: v for k, v in data if v is not None},
+            ),
+            cls=EnumEncoder,
+        )
 
 
 class EnumEncoder(json.JSONEncoder):
@@ -14,16 +33,3 @@ def default(self, obj):
         if isinstance(obj, Enum):
             return obj.value
         return super().default(obj)
-
-
-def to_json_compact(dataclass_obj):
-    """
-    Convert a dataclass to JSON string, excluding None values.
-    """
-    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

@@ -117,7 +117,7 @@ def test_export_initial_telemetry_log(
 
         client = telemetry_client_setup["client"]
         host_url = telemetry_client_setup["host_url"]
-        client.export_event = MagicMock()
+        client._export_event = MagicMock()
 
         driver_connection_params = DriverConnectionParameters(
             http_path="test-path",
@@ -131,7 +131,7 @@ def test_export_initial_telemetry_log(
         client.export_initial_telemetry_log(driver_connection_params, user_agent)
 
         mock_frontend_log.assert_called_once()
-        client.export_event.assert_called_once_with(mock_frontend_log.return_value)
+        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.get_driver_system_configuration")
@@ -155,7 +155,7 @@ def test_export_failure_log(
         mock_frontend_log.return_value = MagicMock()
 
         client = telemetry_client_setup["client"]
-        client.export_event = MagicMock()
+        client._export_event = MagicMock()
 
         client._driver_connection_params = "test-connection-params"
         client._user_agent = "test-user-agent"
@@ -172,7 +172,7 @@ def test_export_failure_log(
 
         mock_frontend_log.assert_called_once()
 
-        client.export_event.assert_called_once_with(mock_frontend_log.return_value)
+        client._export_event.assert_called_once_with(mock_frontend_log.return_value)
 
     def test_export_event(self, telemetry_client_setup):
         """Test exporting an event."""