modified get_telemetry_client and some comments

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -247,9 +247,6 @@ def read(self) -> Optional[OAuthToken]:
         self.telemetry_enabled = (
             self.client_telemetry_enabled and self.server_telemetry_enabled
         )
-        telemetry_batch_size = kwargs.get(
-            "telemetry_batch_size", 100
-        )  # TODO: Decide on batch size
 
         user_agent_entry = kwargs.get("user_agent_entry")
         if user_agent_entry is None:
@@ -315,9 +312,8 @@ def read(self) -> Optional[OAuthToken]:
             discovery_url=TelemetryHelper.get_discovery_url(auth_provider),
             socket_timeout=kwargs.get("_socket_timeout", None),
         )
-        self.telemetry_client = telemetry_client_factory.get_telemetry_client(
+        self.telemetry_client = telemetry_client_factory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
-            batch_size=telemetry_batch_size,
             connection_uuid=self.get_session_id_hex(),
             auth_provider=auth_provider,
             user_agent=useragent_header,

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

@@ -9,8 +9,8 @@ class AuthFlow(Enum):
 class AuthMech(Enum):
     OTHER = "OTHER"
     PAT = "PAT"
-    OAUTH = "OAUTH"
-    EXTERNAL = "EXTERNAL"
+    DATABRICKS_OAUTH = "DATABRICKS_OAUTH"
+    EXTERNAL_AUTH = "EXTERNAL_AUTH"
 
 
 class DatabricksClientType(Enum):

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

@@ -16,7 +16,7 @@ class EnumEncoder(json.JSONEncoder):
     """
     Custom JSON encoder to handle Enum values.
     This is used to convert Enum values to their string representations.
-    Default JSON encoder does not handle Enum values.
+    Default JSON encoder raises a TypeError for Enums.
     """
 
     def default(self, obj):

src/databricks/sql/telemetry/telemetry_client.py

@@ -58,12 +58,14 @@ def get_auth_mechanism(auth_provider):
         if not auth_provider:
             return None
         if isinstance(auth_provider, AccessTokenAuthProvider):
-            return AuthMech.PAT
+            return AuthMech.PAT  # Personal Access Token authentication
         elif isinstance(auth_provider, DatabricksOAuthProvider):
-            return AuthMech.OAUTH
+            return AuthMech.OAUTH  # Databricks-managed OAuth flow
         elif isinstance(auth_provider, ExternalAuthProvider):
-            return AuthMech.EXTERNAL
-        return AuthMech.OTHER
+            return (
+                AuthMech.EXTERNAL
+            )  # External identity provider (AWS IAM, Azure AD, etc.)
+        return AuthMech.OTHER  # Custom or unknown authentication provider
 
     @staticmethod
     def get_auth_flow(auth_provider):
@@ -72,14 +74,15 @@ def get_auth_flow(auth_provider):
             return None
 
         if isinstance(auth_provider, DatabricksOAuthProvider):
-            if (
-                hasattr(auth_provider, "_refresh_token")
-                and auth_provider._refresh_token
-            ):
-                return AuthFlow.TOKEN_PASSTHROUGH
+            if auth_provider._access_token and auth_provider._refresh_token:
+                return (
+                    AuthFlow.TOKEN_PASSTHROUGH
+                )  # Has existing tokens, no user interaction needed
 
             if hasattr(auth_provider, "oauth_manager"):
-                return AuthFlow.BROWSER_BASED_AUTHENTICATION
+                return (
+                    AuthFlow.BROWSER_BASED_AUTHENTICATION
+                )  # Will initiate OAuth flow requiring browser
 
         return None
 
@@ -140,15 +143,14 @@ class TelemetryClient(BaseTelemetryClient):
     def __init__(
         self,
         telemetry_enabled,
-        batch_size,
         connection_uuid,
         auth_provider,
         user_agent,
         driver_connection_params,
         executor,
     ):
         self.telemetry_enabled = telemetry_enabled
-        self.batch_size = batch_size
+        self.batch_size = 10  # TODO: Decide on batch size
         self.connection_uuid = connection_uuid
         self.auth_provider = auth_provider
         self.user_agent = user_agent
@@ -176,15 +178,12 @@ def flush(self):
 
     def _send_telemetry(self, events):
         """Send telemetry events to the server"""
-        try:
-            request = {
-                "uploadTime": int(time.time() * 1000),
-                "items": [],
-                "protoLogs": [event.to_json() for event in events],
-            }
-        except Exception as e:
-            print(f"[DEBUG] Error creating telemetry request: {e}", flush=True)
-            raise e
+
+        request = {
+            "uploadTime": int(time.time() * 1000),
+            "items": [],
+            "protoLogs": [event.to_json() for event in events],
+        }
 
         path = "/telemetry-ext" if self.auth_provider else "/telemetry-unauth"
         url = f"https://{self.host_url}{path}"
@@ -248,10 +247,9 @@ def __init__(self):
         )  # Thread pool for async operations TODO: Decide on max workers
         self._initialized = True
 
-    def get_telemetry_client(
+    def initialize_telemetry_client(
         self,
         telemetry_enabled,
-        batch_size,
         connection_uuid,
         auth_provider,
         user_agent,
@@ -262,7 +260,6 @@ def get_telemetry_client(
             if connection_uuid not in self._clients:
                 self._clients[connection_uuid] = TelemetryClient(
                     telemetry_enabled=telemetry_enabled,
-                    batch_size=batch_size,
                     connection_uuid=connection_uuid,
                     auth_provider=auth_provider,
                     user_agent=user_agent,
@@ -273,6 +270,13 @@ def get_telemetry_client(
         else:
             return NoopTelemetryClient()
 
+    def get_telemetry_client(self, connection_uuid):
+        """Get the telemetry client for a specific connection"""
+        if connection_uuid in self._clients:
+            return self._clients[connection_uuid]
+        else:
+            return NoopTelemetryClient()
+
     def close(self, connection_uuid):
         if connection_uuid in self._clients:
             del self._clients[connection_uuid]

tests/unit/test_telemetry.py

@@ -60,7 +60,6 @@ def setUp(self):
 
         self.client = TelemetryClient(
             telemetry_enabled=True,
-            batch_size=10,
             connection_uuid=self.connection_uuid,
             auth_provider=self.auth_provider,
             user_agent=self.user_agent,
@@ -123,7 +122,6 @@ def test_send_telemetry_unauthenticated(self, mock_post):
         """Test sending telemetry to the server without authentication."""
         unauthenticated_client = TelemetryClient(
             telemetry_enabled=True,
-            batch_size=10,
             connection_uuid=str(uuid.uuid4()),
             auth_provider=None,  # No auth provider
             user_agent=self.user_agent,
@@ -190,9 +188,8 @@ def test_get_telemetry_client_enabled(self, mock_client_class):
         mock_client = MagicMock()
         mock_client_class.return_value = mock_client
 
-        client = self.factory.get_telemetry_client(
+        client = self.factory.initialize_telemetry_client(
             telemetry_enabled=True,
-            batch_size=10,
             connection_uuid=connection_uuid,
             auth_provider=auth_provider,
             user_agent=user_agent,
@@ -202,7 +199,6 @@ def test_get_telemetry_client_enabled(self, mock_client_class):
         # Verify a new client was created and stored
         mock_client_class.assert_called_once_with(
             telemetry_enabled=True,
-            batch_size=10,
             connection_uuid=connection_uuid,
             auth_provider=auth_provider,
             user_agent=user_agent,
@@ -213,24 +209,16 @@ def test_get_telemetry_client_enabled(self, mock_client_class):
         self.assertEqual(self.factory._clients[connection_uuid], mock_client)
 
         # Call again with the same connection_uuid
-        client2 = self.factory.get_telemetry_client(
-            telemetry_enabled=True,
-            batch_size=10,
-            connection_uuid=connection_uuid,
-            auth_provider=auth_provider,
-            user_agent=user_agent,
-            driver_connection_params=driver_connection_params,
-        )
+        client2 = self.factory.get_telemetry_client(connection_uuid=connection_uuid)
 
         # Verify the same client was returned and no new client was created
         self.assertEqual(client2, mock_client)
         mock_client_class.assert_called_once()  # Still only called once
 
     def test_get_telemetry_client_disabled(self):
         """Test getting a telemetry client when telemetry is disabled."""
-        client = self.factory.get_telemetry_client(
+        client = self.factory.initialize_telemetry_client(
             telemetry_enabled=False,
-            batch_size=10,
             connection_uuid="test-uuid",
             auth_provider=MagicMock(),
             user_agent="test-user-agent",
@@ -241,6 +229,9 @@ def test_get_telemetry_client_disabled(self):
         self.assertIsInstance(client, NoopTelemetryClient)
         self.assertEqual(self.factory._clients, {})  # No client was stored
 
+        client2 = self.factory.get_telemetry_client("test-uuid")
+        self.assertIsInstance(client2, NoopTelemetryClient)
+
     @patch("databricks.sql.telemetry.telemetry_client.ThreadPoolExecutor")
     def test_close(self, mock_executor_class):
         """Test closing a client."""