telemetry circuit breaker

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

Author: saishreeeee

poetry.lock

@@ -21,6 +21,7 @@ oauthlib = "^3.1.0"
 openpyxl = "^3.0.10"
 urllib3 = ">=1.26"
 python-dateutil = "^2.8.0"
+pybreaker = "^0.6.0"
 pyarrow = [
     { version = ">=14.0.1", python = ">=3.8,<3.13", optional=true },
     { version = ">=18.0.0", python = ">=3.13", optional=true }

pyproject.toml

@@ -9,6 +9,7 @@
 import logging
 from requests.adapters import HTTPAdapter
 from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
+from pybreaker import CircuitBreaker, CircuitBreakerError
 
 logger = logging.getLogger(__name__)
 
@@ -109,6 +110,9 @@ class TelemetryHttpClient:  # TODO: Unify all the http clients in the PySQL Conn
     TELEMETRY_RETRY_DELAY_MAX = 10.0
     TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
 
+    CIRCUIT_BREAKER_FAIL_MAX = 5
+    CIRCUIT_BREAKER_RESET_TIMEOUT = 60
+
     def __init__(self):
         """Initializes the session and mounts the custom retry adapter."""
         retry_policy = DatabricksRetryPolicy(
@@ -123,6 +127,10 @@ def __init__(self):
         self.session = requests.Session()
         self.session.mount("https://", adapter)
         self.session.mount("http://", adapter)
+        self.breaker = CircuitBreaker(
+            fail_max=self.CIRCUIT_BREAKER_FAIL_MAX,
+            reset_timeout=self.CIRCUIT_BREAKER_RESET_TIMEOUT,
+        )
 
     @classmethod
     def get_instance(cls) -> "TelemetryHttpClient":
@@ -141,7 +149,14 @@ def post(self, url: str, **kwargs) -> requests.Response:
         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)
+        try:
+            return self.breaker.call(self.session.post, url, **kwargs)
+        except CircuitBreakerError as e:
+            logger.error("Circuit breaker error: %s", e)
+            raise e
+        except Exception as e:
+            logger.error("Error executing telemetry POST request: %s", e)
+            raise e
 
     def close(self):
         """Closes the underlying requests.Session."""

src/databricks/sql/common/http.py

@@ -1,13 +1,16 @@
 import uuid
 import pytest
 from unittest.mock import patch, MagicMock
+import time
+import requests
+from pybreaker import CircuitBreakerError
+from databricks.sql.common.http import TelemetryHttpClient
 
 from databricks.sql.telemetry.telemetry_client import (
     TelemetryClient,
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
@@ -316,3 +319,93 @@ def test_connection_failure_sends_correct_telemetry_payload(
         call_arguments = mock_export_failure_log.call_args
         assert call_arguments[0][0] == "Exception"
         assert call_arguments[0][1] == error_message
+
+
+class TestTelemetryHttpClient:
+    """Tests for the TelemetryHttpClient, including retry and circuit breaker logic."""
+
+    @pytest.fixture
+    def http_client(self):
+        """
+        Provides a fresh TelemetryHttpClient instance for each test,
+        ensuring the singleton state is reset.
+        """
+        if TelemetryHttpClient._instance:
+            TelemetryHttpClient.get_instance().close()
+
+        client = TelemetryHttpClient.get_instance()
+        yield client
+
+        client.close()
+
+    def test_circuit_breaker_opens_after_failures(self, http_client):
+        """Verify the circuit opens after N consecutive failures and rejects new calls."""
+        fail_max = 3
+        http_client.breaker.fail_max = fail_max
+
+        with patch.object(http_client.session, "post") as mock_post:
+            mock_post.side_effect = requests.exceptions.RequestException("Connection failed")
+
+            for _ in range(fail_max - 1):
+                with pytest.raises(requests.exceptions.RequestException):
+                    http_client.post("https://test.com/telemetry")
+
+            with pytest.raises(CircuitBreakerError):
+                http_client.post("https://test.com/telemetry")
+
+            assert http_client.breaker.current_state == "open"
+            assert mock_post.call_count == fail_max
+
+            with pytest.raises(CircuitBreakerError):
+                http_client.post("https://test.com/telemetry")
+            assert mock_post.call_count == fail_max
+
+    def test_circuit_breaker_closes_after_timeout_and_success(self, http_client):
+        """Verify the circuit moves to half-open and then closes after a successful probe."""
+        fail_max = 2
+        reset_timeout = 0.1
+        http_client.breaker.fail_max = fail_max
+        http_client.breaker.reset_timeout = reset_timeout
+
+        with patch.object(http_client.session, "post") as mock_post:
+            mock_post.side_effect = [
+                requests.exceptions.RequestException("Fail 1"),
+                requests.exceptions.RequestException("Fail 2"),
+                MagicMock(ok=True)
+            ]
+
+            with pytest.raises(requests.exceptions.RequestException):
+                 http_client.post("https://test.com")
+            with pytest.raises(CircuitBreakerError):
+                 http_client.post("https://test.com")
+            
+            assert http_client.breaker.current_state == "open"
+            time.sleep(reset_timeout)
+
+            http_client.post("https://test.com")
+            assert http_client.breaker.current_state == "closed"
+            assert mock_post.call_count == 3
+
+    def test_circuit_breaker_reopens_if_probe_fails(self, http_client):
+        """Verify the circuit moves to half-open and then back to open if the probe fails."""
+        fail_max = 2
+        reset_timeout = 0.1
+        http_client.breaker.fail_max = fail_max
+        http_client.breaker.reset_timeout = reset_timeout
+
+        with patch.object(http_client.session, "post") as mock_post:
+            mock_post.side_effect = requests.exceptions.RequestException("Always fails")
+
+            with pytest.raises(requests.exceptions.RequestException):
+                 http_client.post("https://test.com")
+            with pytest.raises(CircuitBreakerError):
+                 http_client.post("https://test.com")
+            
+            assert http_client.breaker.current_state == "open"
+            time.sleep(reset_timeout)
+            
+            with pytest.raises(CircuitBreakerError):
+                http_client.post("https://test.com")
+            
+            assert http_client.breaker.current_state == "open"
+            assert mock_post.call_count == 3