retry

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

Author: saishreeeee

src/databricks/sql/exc.py

@@ -1,8 +1,6 @@
 import json
 import logging
 
-from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
-
 logger = logging.getLogger(__name__)
 
 ### PEP-249 Mandated ###
@@ -21,11 +19,11 @@ def __init__(
         self.context = context or {}
 
         error_name = self.__class__.__name__
-        if session_id_hex:
-            telemetry_client = TelemetryClientFactory.get_telemetry_client(
-                session_id_hex
-            )
-            telemetry_client.export_failure_log(error_name, self.message)
+
+        from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
+
+        telemetry_client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
+        telemetry_client.export_failure_log(error_name, self.message)
 
     def __str__(self):
         return self.message

src/databricks/sql/telemetry/telemetry_client.py

@@ -22,6 +22,8 @@
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
+from requests.adapters import HTTPAdapter
+from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 import sys
 import platform
 import uuid
@@ -31,6 +33,24 @@
 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):
+        # The DatabricksRetryPolicy needs state set before the first attempt.
+        if isinstance(self.max_retries, DatabricksRetryPolicy):
+            # Telemetry requests are idempotent and safe to retry. We use CommandType.OTHER
+            # to signal this to the retry policy, bypassing stricter rules for commands
+            # like ExecuteStatement.
+            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."""
 
@@ -146,6 +166,11 @@ 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 = 0.5  # seconds
+    TELEMETRY_RETRY_DELAY_MAX = 5.0  # seconds
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
+
     # Telemetry endpoint paths
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
@@ -170,6 +195,18 @@ def __init__(
         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,  # Not directly used by telemetry, but required by constructor
+            force_dangerous_codes=[],  # Telemetry doesn't have "dangerous" codes
+        )
+        self._session = requests.Session()
+        adapter = TelemetryHTTPAdapter(max_retries=self._telemetry_retry_policy)
+        self._session.mount("https://", adapter)
+
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
         logger.debug("Exporting event for connection %s", self._session_id_hex)
@@ -215,7 +252,7 @@ def _send_telemetry(self, events):
         try:
             logger.debug("Submitting telemetry request to thread pool")
             future = self._executor.submit(
-                requests.post,
+                self._session.post,
                 url,
                 data=json.dumps(request),
                 headers=headers,
@@ -303,6 +340,7 @@ 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:

tests/unit/test_telemetry.py

@@ -198,7 +198,7 @@ def test_export_event(self, telemetry_client_setup):
         client._flush.assert_called_once()
         assert len(client._events_batch) == 10
 
-    @patch("requests.post")
+    @patch("requests.Session.post")
     def test_send_telemetry_authenticated(self, mock_post, telemetry_client_setup):
         """Test sending telemetry to the server with authentication."""
         client = telemetry_client_setup["client"]
@@ -212,12 +212,12 @@ def test_send_telemetry_authenticated(self, mock_post, telemetry_client_setup):
 
         executor.submit.assert_called_once()
         args, kwargs = executor.submit.call_args
-        assert args[0] == requests.post
+        assert args[0] == client._session.post
         assert kwargs["timeout"] == 10
         assert "Authorization" in kwargs["headers"]
         assert kwargs["headers"]["Authorization"] == "Bearer test-token"
 
-    @patch("requests.post")
+    @patch("requests.Session.post")
     def test_send_telemetry_unauthenticated(self, mock_post, telemetry_client_setup):
         """Test sending telemetry to the server without authentication."""
         host_url = telemetry_client_setup["host_url"]
@@ -239,7 +239,7 @@ def test_send_telemetry_unauthenticated(self, mock_post, telemetry_client_setup)
 
         executor.submit.assert_called_once()
         args, kwargs = executor.submit.call_args
-        assert args[0] == requests.post
+        assert args[0] == unauthenticated_client._session.post
         assert kwargs["timeout"] == 10
         assert "Authorization" not in kwargs["headers"]  # No auth header
         assert kwargs["headers"]["Accept"] == "application/json"
@@ -331,6 +331,34 @@ class TestBaseClient(BaseTelemetryClient):
         with pytest.raises(TypeError):
             TestBaseClient()  # Can't instantiate abstract class
 
+    def test_telemetry_http_adapter_retry_policy(self, telemetry_client_setup):
+        """Test that TelemetryHTTPAdapter properly configures DatabricksRetryPolicy."""
+        from databricks.sql.telemetry.telemetry_client import TelemetryHTTPAdapter
+        from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
+        
+        client = telemetry_client_setup["client"]
+        
+        # Verify that the session has the TelemetryHTTPAdapter mounted
+        adapter = client._session.adapters.get("https://")
+        assert isinstance(adapter, TelemetryHTTPAdapter)
+        assert isinstance(adapter.max_retries, DatabricksRetryPolicy)
+        
+        # Verify that the retry policy has the correct configuration
+        retry_policy = adapter.max_retries
+        assert retry_policy.delay_min == client.TELEMETRY_RETRY_DELAY_MIN
+        assert retry_policy.delay_max == client.TELEMETRY_RETRY_DELAY_MAX
+        assert retry_policy.stop_after_attempts_count == client.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT
+        assert retry_policy.stop_after_attempts_duration == client.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION
+        
+        # Test that the adapter's send method would properly configure the retry policy
+        # by directly testing the logic that sets command_type and starts the timer
+        if isinstance(adapter.max_retries, DatabricksRetryPolicy):
+            adapter.max_retries.command_type = CommandType.OTHER
+            adapter.max_retries.start_retry_timer()
+            
+            # Verify that the retry policy was configured correctly
+            assert retry_policy.command_type == CommandType.OTHER
+
 
 class TestTelemetryHelper:
     """Tests for the TelemetryHelper class."""