removed unused params from driver connection params, initialize_telemetry_client does not return a telemetry client

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -303,13 +303,17 @@ def read(self) -> Optional[OAuthToken]:
             kwargs.get("use_inline_params", False)
         )
 
-        self.telemetry_client = TelemetryClientFactory.initialize_telemetry_client(
+        TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
             connection_uuid=self.get_session_id_hex(),
             auth_provider=auth_provider,
             host_url=self.host,
         )
 
+        self.telemetry_client = TelemetryClientFactory.get_telemetry_client(
+            connection_uuid=self.get_session_id_hex()
+        )
+
         driver_connection_params = DriverConnectionParameters(
             http_path=http_path,
             mode=DatabricksClientType.THRIFT,

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

@@ -54,10 +54,6 @@ class DriverConnectionParameters:
         host_info (HostDetails): Details about the host connection
         auth_mech (AuthMech): The authentication mechanism used
         auth_flow (AuthFlow): The authentication flow type
-        auth_scope (str): The scope of authentication
-        discovery_url (str): URL for service discovery
-        allowed_volume_ingestion_paths (str): JSON string of allowed paths for volume operations
-        azure_tenant_id (str): Azure tenant ID for Azure authentication
         socket_timeout (int): Connection timeout in milliseconds
     """
 
@@ -66,10 +62,6 @@ class DriverConnectionParameters:
     host_info: HostDetails
     auth_mech: Optional[AuthMech] = None
     auth_flow: Optional[AuthFlow] = None
-    auth_scope: Optional[str] = None
-    discovery_url: Optional[str] = None
-    allowed_volume_ingestion_paths: Optional[str] = None
-    azure_tenant_id: Optional[str] = None
     socket_timeout: Optional[int] = None
 
     def to_json(self):

src/databricks/sql/telemetry/telemetry_client.py

@@ -60,11 +60,11 @@ def get_auth_mechanism(auth_provider):
         if isinstance(auth_provider, AccessTokenAuthProvider):
             return AuthMech.PAT  # Personal Access Token authentication
         elif isinstance(auth_provider, DatabricksOAuthProvider):
-            return AuthMech.OAUTH  # Databricks-managed OAuth flow
+            return AuthMech.DATABRICKS_OAUTH  # Databricks-managed OAuth flow
         elif isinstance(auth_provider, ExternalAuthProvider):
             return (
-                AuthMech.EXTERNAL
-            )  # External identity provider (AWS IAM, Azure AD, etc.)
+                AuthMech.EXTERNAL_AUTH
+            )  # External identity provider (AWS, Azure, etc.)
         return AuthMech.OTHER  # Custom or unknown authentication provider
 
     @staticmethod
@@ -269,29 +269,33 @@ def initialize_telemetry_client(
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         TelemetryClientFactory._initialize()
 
-        if telemetry_enabled:
-            with TelemetryClientFactory._lock:
-                if connection_uuid not in TelemetryClientFactory._clients:
-                    logger.info(
-                        f"Creating new TelemetryClient for connection {connection_uuid}"
-                    )
+        with TelemetryClientFactory._lock:
+            if connection_uuid not in TelemetryClientFactory._clients:
+                logger.info(
+                    f"Creating new TelemetryClient for connection {connection_uuid}"
+                )
+                if telemetry_enabled:
                     TelemetryClientFactory._clients[connection_uuid] = TelemetryClient(
                         telemetry_enabled=telemetry_enabled,
                         connection_uuid=connection_uuid,
                         auth_provider=auth_provider,
                         host_url=host_url,
                         executor=TelemetryClientFactory._executor,
                     )
-                return TelemetryClientFactory._clients[connection_uuid]
-        else:
-            return NoopTelemetryClient()
+                else:
+                    TelemetryClientFactory._clients[
+                        connection_uuid
+                    ] = NoopTelemetryClient()
 
     @staticmethod
     def get_telemetry_client(connection_uuid):
         """Get the telemetry client for a specific connection"""
         if connection_uuid in TelemetryClientFactory._clients:
             return TelemetryClientFactory._clients[connection_uuid]
         else:
+            logger.error(
+                f"Telemetry client not initialized for connection {connection_uuid}"
+            )
             return NoopTelemetryClient()
 
     @staticmethod

tests/unit/test_telemetry.py

@@ -182,7 +182,7 @@ def test_initialize_telemetry_client_enabled(self, mock_client_class):
         mock_client = MagicMock()
         mock_client_class.return_value = mock_client
 
-        client = TelemetryClientFactory.initialize_telemetry_client(
+        TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=True,
             connection_uuid=connection_uuid,
             auth_provider=auth_provider,
@@ -197,7 +197,6 @@ def test_initialize_telemetry_client_enabled(self, mock_client_class):
             host_url=host_url,
             executor=TelemetryClientFactory._executor,
         )
-        self.assertEqual(client, mock_client)
         self.assertEqual(TelemetryClientFactory._clients[connection_uuid], mock_client)
 
         # Call again with the same connection_uuid
@@ -209,18 +208,18 @@ def test_initialize_telemetry_client_enabled(self, mock_client_class):
 
     def test_initialize_telemetry_client_disabled(self):
         """Test initializing a telemetry client when telemetry is disabled."""
-        client = TelemetryClientFactory.initialize_telemetry_client(
+        connection_uuid = "test-uuid"
+        TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=False,
-            connection_uuid="test-uuid",
+            connection_uuid=connection_uuid,
             auth_provider=MagicMock(),
             host_url="test-host",
         )
 
-        # Verify a NoopTelemetryClient was returned
-        self.assertIsInstance(client, NoopTelemetryClient)
-        self.assertEqual(TelemetryClientFactory._clients, {})  # No client was stored
+        # Verify a NoopTelemetryClient was stored
+        self.assertIsInstance(TelemetryClientFactory._clients[connection_uuid], NoopTelemetryClient)
 
-        client2 = TelemetryClientFactory.get_telemetry_client("test-uuid")
+        client2 = TelemetryClientFactory.get_telemetry_client(connection_uuid)
         self.assertIsInstance(client2, NoopTelemetryClient)
 
     def test_get_telemetry_client_existing(self):
@@ -267,7 +266,7 @@ def test_close(self, mock_client_class, mock_executor_class):
         mock_executor_class.return_value = mock_executor2
         mock_client_class.return_value = mock_client2
 
-        client = TelemetryClientFactory.initialize_telemetry_client(
+        TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=True,
             connection_uuid=connection_uuid2,
             auth_provider=MagicMock(),
@@ -279,7 +278,7 @@ def test_close(self, mock_client_class, mock_executor_class):
         self.assertIsNotNone(TelemetryClientFactory._executor)
         self.assertEqual(TelemetryClientFactory._executor, mock_executor2)
         self.assertIn(connection_uuid2, TelemetryClientFactory._clients)
-        self.assertEqual(client, mock_client2)
+        self.assertEqual(TelemetryClientFactory._clients[connection_uuid2], mock_client2)
 
         # Verify new ThreadPoolExecutor was created
         self.assertEqual(mock_executor_class.call_count, 1)