JsonSerializableMixin, TelemetryRequest

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

Author: saishreeeee

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

@@ -1,10 +1,11 @@
 import json
 from dataclasses import dataclass, asdict
 from typing import List, Optional
+from databricks.sql.telemetry.utils import JsonSerializableMixin
 
 
 @dataclass
-class TelemetryRequest:
+class TelemetryRequest(JsonSerializableMixin):
     """
     Represents a request to send telemetry data to the server side.
     Contains the telemetry items to be uploaded and optional protocol buffer logs.
@@ -19,12 +20,9 @@ class TelemetryRequest:
     items: List[str]
     protoLogs: Optional[List[str]]
 
-    def to_json(self):
-        return json.dumps(asdict(self))
-
 
 @dataclass
-class TelemetryResponse:
+class TelemetryResponse(JsonSerializableMixin):
     """
     Represents the response from the telemetry backend after processing a request.
     Contains information about the success or failure of the telemetry upload.
@@ -38,6 +36,3 @@ class TelemetryResponse:
     errors: List[str]
     numSuccess: int
     numProtoSuccess: int
-
-    def to_json(self):
-        return json.dumps(asdict(self))

src/databricks/sql/telemetry/telemetry_client.py

@@ -17,6 +17,10 @@
     FrontendLogEntry,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
+from databricks.sql.telemetry.models.endpoint_models import (
+    TelemetryRequest,
+    TelemetryResponse,
+)
 from databricks.sql.auth.authenticators import (
     AccessTokenAuthProvider,
     DatabricksOAuthProvider,
@@ -202,11 +206,11 @@ def _flush(self):
     def _send_telemetry(self, events):
         """Send telemetry events to the server"""
 
-        request = {
-            "uploadTime": int(time.time() * 1000),
-            "items": [],
-            "protoLogs": [event.to_json() for event in events],
-        }
+        request = TelemetryRequest(
+            uploadTime=int(time.time() * 1000),
+            items=[],
+            protoLogs=[event.to_json() for event in events],
+        )
 
         path = (
             self.TELEMETRY_AUTHENTICATED_PATH
@@ -225,7 +229,7 @@ def _send_telemetry(self, events):
             future = self._executor.submit(
                 requests.post,
                 url,
-                data=json.dumps(request),
+                data=request.to_json(),
                 headers=headers,
                 timeout=10,
             )