remove separate http client for telemetry

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/client.py

@@ -279,7 +279,7 @@ def read(self) -> Optional[OAuthToken]:
                 host_url=server_hostname,
                 http_path=http_path,
                 port=kwargs.get("_port", 443),
-                client_context=client_context,
+                http_client=self.http_client,
                 user_agent=self.session.useragent_header
                 if hasattr(self, "session")
                 else None,
@@ -301,7 +301,7 @@ def read(self) -> Optional[OAuthToken]:
             auth_provider=self.session.auth_provider,
             host_url=self.session.host,
             batch_size=self.telemetry_batch_size,
-            client_context=client_context,
+            http_client=self.http_client,
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(

src/databricks/sql/telemetry/telemetry_client.py

@@ -154,44 +154,6 @@ def _flush(self):
         pass
 
 
-class TelemetryHttpClientSingleton:
-    """
-    Singleton HTTP client for telemetry operations.
-
-    This ensures that telemetry has its own dedicated HTTP client that
-    is independent of individual connection lifecycles.
-    """
-
-    _instance = None
-    _lock = threading.RLock()
-
-    def __new__(cls):
-        if cls._instance is None:
-            with cls._lock:
-                if cls._instance is None:
-                    cls._instance = super().__new__(cls)
-                    cls._instance._http_client = None
-                    cls._instance._initialized = False
-        return cls._instance
-
-    def get_http_client(self, client_context):
-        """Get or create the singleton HTTP client."""
-        if not self._initialized and client_context:
-            with self._lock:
-                if not self._initialized:
-                    self._http_client = UnifiedHttpClient(client_context)
-                    self._initialized = True
-        return self._http_client
-
-    def close(self):
-        """Close the singleton HTTP client."""
-        with self._lock:
-            if self._http_client:
-                self._http_client.close()
-                self._http_client = None
-                self._initialized = False
-
-
 class TelemetryClient(BaseTelemetryClient):
     """
     Telemetry client class that handles sending telemetry events in batches to the server.
@@ -210,7 +172,7 @@ def __init__(
         host_url,
         executor,
         batch_size,
-        client_context,
+        http_client,
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
@@ -224,9 +186,8 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
 
-        # Use singleton HTTP client for telemetry instead of connection-specific client
-        self._http_client_singleton = TelemetryHttpClientSingleton()
-        self._http_client = self._http_client_singleton.get_http_client(client_context)
+        # Use the provided HTTP client directly
+        self._http_client = http_client
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
@@ -503,7 +464,7 @@ def initialize_telemetry_client(
         auth_provider,
         host_url,
         batch_size,
-        client_context,
+        http_client,
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
@@ -526,7 +487,7 @@ def initialize_telemetry_client(
                             host_url=host_url,
                             executor=TelemetryClientFactory._executor,
                             batch_size=batch_size,
-                            client_context=client_context,
+                            http_client=http_client,
                         )
                     else:
                         TelemetryClientFactory._clients[
@@ -579,10 +540,10 @@ def connection_failure_log(
         host_url: str,
         http_path: str,
         port: int,
-        client_context,
+        http_client,
         user_agent: Optional[str] = None,
     ):
-        """Send error telemetry when connection creation fails, without requiring a session"""
+        """Send error telemetry when connection creation fails, using existing HTTP client"""
 
         UNAUTH_DUMMY_SESSION_ID = "unauth_session_id"
 
@@ -592,7 +553,7 @@ def connection_failure_log(
             auth_provider=None,
             host_url=host_url,
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            client_context=client_context,
+            http_client=http_client,
         )
 
         telemetry_client = TelemetryClientFactory.get_telemetry_client(

tests/unit/test_telemetry.py

@@ -19,17 +19,12 @@
 
 
 @pytest.fixture
-@patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
-def mock_telemetry_client(mock_singleton_class):
+def mock_telemetry_client():
     """Create a mock telemetry client for testing."""
     session_id = str(uuid.uuid4())
     auth_provider = AccessTokenAuthProvider("test-token")
     executor = MagicMock()
-    mock_client_context = MagicMock()
-    
-    # Mock the singleton to return a mock HTTP client
-    mock_singleton = mock_singleton_class.return_value
-    mock_singleton.get_http_client.return_value = MagicMock()
+    mock_http_client = MagicMock()
 
     return TelemetryClient(
         telemetry_enabled=True,
@@ -38,7 +33,7 @@ def mock_telemetry_client(mock_singleton_class):
         host_url="test-host.com",
         executor=executor,
         batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-        client_context=mock_client_context,
+        http_client=mock_http_client,
     )
 
 
@@ -217,16 +212,11 @@ def telemetry_system_reset(self):
         TelemetryClientFactory._executor = None
         TelemetryClientFactory._initialized = False
 
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
-    def test_client_lifecycle_flow(self, mock_singleton_class):
+    def test_client_lifecycle_flow(self):
         """Test complete client lifecycle: initialize -> use -> close."""
         session_id_hex = "test-session"
         auth_provider = AccessTokenAuthProvider("token")
-        mock_client_context = MagicMock()
-        
-        # Mock the singleton to return a mock HTTP client
-        mock_singleton = mock_singleton_class.return_value
-        mock_singleton.get_http_client.return_value = MagicMock()
+        mock_http_client = MagicMock()
 
         # Initialize enabled client
         TelemetryClientFactory.initialize_telemetry_client(
@@ -235,7 +225,7 @@ def test_client_lifecycle_flow(self, mock_singleton_class):
             auth_provider=auth_provider,
             host_url="test-host.com",
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            client_context=mock_client_context,
+            http_client=mock_http_client,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -248,40 +238,28 @@ def test_client_lifecycle_flow(self, mock_singleton_class):
             mock_close.assert_called_once()
 
         # Should get NoopTelemetryClient after close
-        client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
-        assert isinstance(client, NoopTelemetryClient)
 
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
-    def test_disabled_telemetry_flow(self, mock_singleton_class):
-        """Test that disabled telemetry uses NoopTelemetryClient."""
+    def test_disabled_telemetry_creates_noop_client(self):
+        """Test that disabled telemetry creates NoopTelemetryClient."""
         session_id_hex = "test-session"
-        mock_client_context = MagicMock()
-        
-        # Mock the singleton to return a mock HTTP client
-        mock_singleton = mock_singleton_class.return_value
-        mock_singleton.get_http_client.return_value = MagicMock()
+        mock_http_client = MagicMock()
 
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=False,
             session_id_hex=session_id_hex,
             auth_provider=None,
             host_url="test-host.com",
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            client_context=mock_client_context,
+            http_client=mock_http_client,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, NoopTelemetryClient)
 
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
-    def test_factory_error_handling(self, mock_singleton_class):
+    def test_factory_error_handling(self):
         """Test that factory errors fall back to NoopTelemetryClient."""
         session_id = "test-session"
-        mock_client_context = MagicMock()
-        
-        # Mock the singleton to return a mock HTTP client
-        mock_singleton = mock_singleton_class.return_value
-        mock_singleton.get_http_client.return_value = MagicMock()
+        mock_http_client = MagicMock()
 
         # Simulate initialization error
         with patch(
@@ -294,23 +272,18 @@ def test_factory_error_handling(self, mock_singleton_class):
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-                client_context=mock_client_context,
+                http_client=mock_http_client,
             )
 
         # Should fall back to NoopTelemetryClient
         client = TelemetryClientFactory.get_telemetry_client(session_id)
         assert isinstance(client, NoopTelemetryClient)
 
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
-    def test_factory_shutdown_flow(self, mock_singleton_class):
+    def test_factory_shutdown_flow(self):
         """Test factory shutdown when last client is removed."""
         session1 = "session-1"
         session2 = "session-2"
-        mock_client_context = MagicMock()
-        
-        # Mock the singleton to return a mock HTTP client
-        mock_singleton = mock_singleton_class.return_value
-        mock_singleton.get_http_client.return_value = MagicMock()
+        mock_http_client = MagicMock()
 
         # Initialize multiple clients
         for session in [session1, session2]:
@@ -320,7 +293,7 @@ def test_factory_shutdown_flow(self, mock_singleton_class):
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-                client_context=mock_client_context,
+                http_client=mock_http_client,
             )
 
         # Factory should be initialized