improved circuit breaker for handling only 429/503

Signed-off-by: Nikhil Suri <[email protected]>

Author: nikhilsuri-db

src/databricks/sql/common/unified_http_client.py

@@ -264,7 +264,25 @@ def request_context(
             yield response
         except MaxRetryError as e:
             logger.error("HTTP request failed after retries: %s", e)
-            raise RequestError(f"HTTP request failed: {e}")
+            
+            # Try to extract HTTP status code from the MaxRetryError
+            http_code = None
+            if hasattr(e, 'reason') and hasattr(e.reason, 'response'):
+                # The reason may contain a response object with status
+                http_code = getattr(e.reason.response, 'status', None)
+            elif hasattr(e, 'response') and hasattr(e.response, 'status'):
+                # Or the error itself may have a response
+                http_code = e.response.status
+            
+            context = {}
+            if http_code is not None:
+                context["http-code"] = http_code
+                logger.error("HTTP request failed with status code: %d", http_code)
+            
+            raise RequestError(
+                f"HTTP request failed: {e}",
+                context=context
+            )
         except Exception as e:
             logger.error("HTTP request error: %s", e)
             raise RequestError(f"HTTP request error: {e}")

src/databricks/sql/exc.py

@@ -126,3 +126,9 @@ class SessionAlreadyClosedError(RequestError):
 
 class CursorAlreadyClosedError(RequestError):
     """Thrown if CancelOperation receives a code 404. ThriftBackend should gracefully proceed as this is expected."""
+
+
+class TelemetryRateLimitError(Exception):
+    """Raised when telemetry endpoint returns 429 or 503, indicating rate limiting or service unavailable.
+    This exception is used exclusively by the circuit breaker to track telemetry rate limiting events."""
+    pass

src/databricks/sql/telemetry/circuit_breaker_manager.py

@@ -14,18 +14,19 @@
 import pybreaker
 from pybreaker import CircuitBreaker, CircuitBreakerError, CircuitBreakerListener
 
+from databricks.sql.exc import TelemetryRateLimitError
+
 logger = logging.getLogger(__name__)
 
 # Circuit Breaker Configuration Constants
-MINIMUM_CALLS = 20
-RESET_TIMEOUT = 30
-CIRCUIT_BREAKER_NAME = "telemetry-circuit-breaker"
+DEFAULT_MINIMUM_CALLS = 20
+DEFAULT_RESET_TIMEOUT = 30
+DEFAULT_NAME = "telemetry-circuit-breaker"
 
-# Circuit Breaker State Constants
+# Circuit Breaker State Constants (used in logging)
 CIRCUIT_BREAKER_STATE_OPEN = "open"
 CIRCUIT_BREAKER_STATE_CLOSED = "closed"
 CIRCUIT_BREAKER_STATE_HALF_OPEN = "half-open"
-CIRCUIT_BREAKER_STATE_DISABLED = "disabled"
 
 # Logging Message Constants
 LOG_CIRCUIT_BREAKER_STATE_CHANGED = "Circuit breaker state changed from %s to %s for %s"
@@ -72,18 +73,47 @@ def state_change(self, cb: CircuitBreaker, old_state, new_state) -> None:
             logger.info(LOG_CIRCUIT_BREAKER_HALF_OPEN, cb.name)
 
 
+@dataclass(frozen=True)
+class CircuitBreakerConfig:
+    """Configuration for circuit breaker behavior.
+
+    This class is immutable to prevent modification of circuit breaker settings.
+    All configuration values are set to constants defined at the module level.
+    """
+
+    # Minimum number of calls before circuit can open
+    minimum_calls: int = DEFAULT_MINIMUM_CALLS
+
+    # Time to wait before trying to close circuit (in seconds)
+    reset_timeout: int = DEFAULT_RESET_TIMEOUT
+
+    # Name for the circuit breaker (for logging)
+    name: str = DEFAULT_NAME
+
+
 class CircuitBreakerManager:
     """
     Manages circuit breaker instances for telemetry requests.
 
     This class provides a singleton pattern to manage circuit breaker instances
     per host, ensuring that telemetry failures don't impact main SQL operations.
-
-    Circuit breaker configuration is fixed and cannot be overridden.
     """
 
     _instances: Dict[str, CircuitBreaker] = {}
     _lock = threading.RLock()
+    _config: Optional[CircuitBreakerConfig] = None
+
+    @classmethod
+    def initialize(cls, config: CircuitBreakerConfig) -> None:
+        """
+        Initialize the circuit breaker manager with configuration.
+
+        Args:
+            config: Circuit breaker configuration
+        """
+        with cls._lock:
+            cls._config = config
+            logger.debug("CircuitBreakerManager initialized with config: %s", config)
 
     @classmethod
     def get_circuit_breaker(cls, host: str) -> CircuitBreaker:
@@ -96,6 +126,10 @@ def get_circuit_breaker(cls, host: str) -> CircuitBreaker:
         Returns:
             CircuitBreaker instance for the host
         """
+        if not cls._config:
+            # Return a no-op circuit breaker if not initialized
+            return cls._create_noop_circuit_breaker()
+
         with cls._lock:
             if host not in cls._instances:
                 cls._instances[host] = cls._create_circuit_breaker(host)
@@ -114,16 +148,39 @@ def _create_circuit_breaker(cls, host: str) -> CircuitBreaker:
         Returns:
             New CircuitBreaker instance
         """
-        # Create circuit breaker with fixed configuration
+        config = cls._config
+        if config is None:
+            raise RuntimeError("CircuitBreakerManager not initialized")
+
+        # Create circuit breaker with configuration
         breaker = CircuitBreaker(
-            fail_max=MINIMUM_CALLS,
-            reset_timeout=RESET_TIMEOUT,
-            name=f"{CIRCUIT_BREAKER_NAME}-{host}",
+            fail_max=config.minimum_calls,  # Number of failures before circuit opens
+            reset_timeout=config.reset_timeout,
+            name=f"{config.name}-{host}",
         )
+
+        # Add state change listeners for logging
         breaker.add_listener(CircuitBreakerStateListener())
 
         return breaker
 
+    @classmethod
+    def _create_noop_circuit_breaker(cls) -> CircuitBreaker:
+        """
+        Create a no-op circuit breaker that always allows calls.
+
+        Returns:
+            CircuitBreaker that never opens
+        """
+        # Create a circuit breaker with very high thresholds so it never opens
+        breaker = CircuitBreaker(
+            fail_max=1000000,  # Very high threshold
+            reset_timeout=1,  # Short reset time
+            name="noop-circuit-breaker",
+        )
+        return breaker
+
+
 
 def is_circuit_breaker_error(exception: Exception) -> bool:
     """

src/databricks/sql/telemetry/telemetry_client.py

@@ -46,9 +46,6 @@
     TelemetryPushClient,
     CircuitBreakerTelemetryPushClient,
 )
-from databricks.sql.telemetry.circuit_breaker_manager import (
-    is_circuit_breaker_error,
-)
 
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
@@ -275,21 +272,23 @@ def _send_telemetry(self, events):
             logger.debug("Failed to submit telemetry request: %s", e)
 
     def _send_with_unified_client(self, url, data, headers, timeout=900):
-        """Helper method to send telemetry using the telemetry push client."""
+        """
+        Helper method to send telemetry using the telemetry push client.
+        
+        The push client implementation handles circuit breaker logic internally,
+        so this method just forwards the request and handles any errors generically.
+        """
         try:
             response = self._telemetry_push_client.request(
                 HttpMethod.POST, url, body=data, headers=headers, timeout=timeout
             )
             return response
         except Exception as e:
-            if is_circuit_breaker_error(e):
-                logger.warning(
-                    "Telemetry request blocked by circuit breaker for connection %s: %s",
-                    self._session_id_hex,
-                    e,
-                )
-            else:
-                logger.error("Failed to send telemetry: %s", e)
+            logger.debug(
+                "Failed to send telemetry for connection %s: %s",
+                self._session_id_hex,
+                e,
+            )
             raise
 
     def _telemetry_request_callback(self, future, sent_count: int):