shifted telemetry http client to http.py

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: saishreeeee

src/databricks/sql/common/http.py

@@ -5,8 +5,10 @@
 import threading
 from dataclasses import dataclass
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Optional
 import logging
+from requests.adapters import HTTPAdapter
+from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 
 logger = logging.getLogger(__name__)
 
@@ -81,3 +83,70 @@ def execute(
 
     def close(self):
         self.session.close()
+
+
+class TelemetryHTTPAdapter(HTTPAdapter):
+    """
+    Custom HTTP adapter to prepare our DatabricksRetryPolicy before each request.
+    This ensures the retry timer is started and the command type is set correctly,
+    allowing the policy to manage its state for the duration of the request retries.
+    """
+
+    def send(self, request, **kwargs):
+        self.max_retries.command_type = CommandType.OTHER
+        self.max_retries.start_retry_timer()
+        return super().send(request, **kwargs)
+
+
+class TelemetryHttpClient: # TODO: Unify all the http clients in the PySQL Connector
+    """Singleton HTTP client for sending telemetry data."""
+
+    _instance: Optional["TelemetryHttpClient"] = None
+    _lock = threading.Lock()
+
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT = 3
+    TELEMETRY_RETRY_DELAY_MIN = 1.0
+    TELEMETRY_RETRY_DELAY_MAX = 10.0
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
+
+    def __init__(self):
+        """Initializes the session and mounts the custom retry adapter."""
+        retry_policy = DatabricksRetryPolicy(
+            delay_min=self.TELEMETRY_RETRY_DELAY_MIN,
+            delay_max=self.TELEMETRY_RETRY_DELAY_MAX,
+            stop_after_attempts_count=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT,
+            stop_after_attempts_duration=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION,
+            delay_default=1.0,
+            force_dangerous_codes=[],
+        )
+        adapter = TelemetryHTTPAdapter(max_retries=retry_policy)
+        self.session = requests.Session()
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+
+    @classmethod
+    def get_instance(cls) -> "TelemetryHttpClient":
+        """Get the singleton instance of the TelemetryHttpClient."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    logger.debug("Initializing singleton TelemetryHttpClient")
+                    cls._instance = TelemetryHttpClient()
+        return cls._instance
+
+    def post(self, url: str, **kwargs) -> requests.Response:
+        """
+        Executes a POST request using the configured session.
+
+        This is a blocking call intended to be run in a background thread.
+        """
+        logger.debug("Executing telemetry POST request to: %s", url)
+        return self.session.post(url, **kwargs)
+
+    def close(self):
+        """Closes the underlying requests.Session."""
+        logger.debug("Closing TelemetryHttpClient session.")
+        self.session.close()
+        # Clear the instance to allow for re-initialization if needed
+        with TelemetryHttpClient._lock:
+            TelemetryHttpClient._instance = None

src/databricks/sql/telemetry/telemetry_client.py

@@ -4,6 +4,7 @@
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Optional
+from databricks.sql.common.http import TelemetryHttpClient
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -25,8 +26,6 @@
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
-from requests.adapters import HTTPAdapter
-from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 import sys
 import platform
 import uuid
@@ -36,19 +35,6 @@
 logger = logging.getLogger(__name__)
 
 
-class TelemetryHTTPAdapter(HTTPAdapter):
-    """
-    Custom HTTP adapter to prepare our DatabricksRetryPolicy before each request.
-    This ensures the retry timer is started and the command type is set correctly,
-    allowing the policy to manage its state for the duration of the request retries.
-    """
-
-    def send(self, request, **kwargs):
-        self.max_retries.command_type = CommandType.OTHER
-        self.max_retries.start_retry_timer()
-        return super().send(request, **kwargs)
-
-
 class TelemetryHelper:
     """Helper class for getting telemetry related information."""
 
@@ -143,11 +129,6 @@ class TelemetryClient(BaseTelemetryClient):
     It uses a thread pool to handle asynchronous operations, that it gets from the TelemetryClientFactory.
     """
 
-    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT = 3
-    TELEMETRY_RETRY_DELAY_MIN = 1.0
-    TELEMETRY_RETRY_DELAY_MAX = 10.0
-    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
-
     # Telemetry endpoint paths
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
@@ -173,21 +154,7 @@ def __init__(
         self._driver_connection_params = None
         self._host_url = host_url
         self._executor = executor
-
-        self._telemetry_retry_policy = DatabricksRetryPolicy(
-            delay_min=self.TELEMETRY_RETRY_DELAY_MIN,
-            delay_max=self.TELEMETRY_RETRY_DELAY_MAX,
-            stop_after_attempts_count=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT,
-            stop_after_attempts_duration=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION,
-            delay_default=1.0,
-            force_dangerous_codes=[],
-        )
-        self._session = (
-            requests.Session()
-        )  # TODO: Use DatabricksHttpClient instead (unify all http clients)
-        adapter = TelemetryHTTPAdapter(max_retries=self._telemetry_retry_policy)
-        self._session.mount("https://", adapter)
-        self._session.mount("http://", adapter)
+        self._http_client = TelemetryHttpClient.get_instance()
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
@@ -236,7 +203,7 @@ def _send_telemetry(self, events):
         try:
             logger.debug("Submitting telemetry request to thread pool")
             future = self._executor.submit(
-                self._session.post,
+                self._http_client.post,
                 url,
                 data=request.to_json(),
                 headers=headers,
@@ -341,7 +308,6 @@ def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
         self._flush()
-        self._session.close()
 
 
 class TelemetryClientFactory:
@@ -463,6 +429,7 @@ def close(session_id_hex):
                 )
                 try:
                     TelemetryClientFactory._executor.shutdown(wait=True)
+                    TelemetryHttpClient.get_instance().close()
                 except Exception as e:
                     logger.debug("Failed to shutdown thread pool executor: %s", e)
                 TelemetryClientFactory._executor = None

tests/unit/test_telemetry.py

@@ -91,7 +91,7 @@ def test_network_request_flow(self, mock_post, mock_telemetry_client):
         args, kwargs = client._executor.submit.call_args
 
         # Verify correct function and URL
-        assert args[0] == client._session.post
+        assert args[0] == client._http_client.post
         assert args[1] == "https://test-host.com/telemetry-ext"
         assert kwargs["headers"]["Authorization"] == "Bearer test-token"