circuit breaker changes using pybreaker

Signed-off-by: Nikhil Suri <nikhil.suri@databricks.com>

Author: nikhilsuri-db

docs/parameters.md

@@ -254,3 +254,73 @@ You should only set `use_inline_params=True` in the following cases:
 4. Your client code uses [sequences as parameter values](#passing-sequences-as-parameter-values)
 
 We expect limitations (1) and (2) to be addressed in a future Databricks Runtime release.
+
+# Telemetry Circuit Breaker Configuration
+
+The Databricks SQL connector includes a circuit breaker pattern for telemetry requests to prevent telemetry failures from impacting main SQL operations. This feature is enabled by default and can be controlled through a connection parameter.
+
+## Overview
+
+The circuit breaker monitors telemetry request failures and automatically blocks telemetry requests when the failure rate exceeds a configured threshold. This prevents telemetry service issues from affecting your main SQL operations.
+
+## Configuration Parameter
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `telemetry_circuit_breaker_enabled` | bool | `True` | Enable or disable the telemetry circuit breaker |
+
+## Usage Examples
+
+### Default Configuration (Circuit Breaker Enabled)
+
+```python
+from databricks import sql
+
+# Circuit breaker is enabled by default
+with sql.connect(
+    server_hostname="your-host.cloud.databricks.com",
+    http_path="/sql/1.0/warehouses/your-warehouse-id",
+    access_token="your-token"
+) as conn:
+    # Your SQL operations here
+    pass
+```
+
+### Disable Circuit Breaker
+
+```python
+from databricks import sql
+
+# Disable circuit breaker entirely
+with sql.connect(
+    server_hostname="your-host.cloud.databricks.com",
+    http_path="/sql/1.0/warehouses/your-warehouse-id",
+    access_token="your-token",
+    telemetry_circuit_breaker_enabled=False
+) as conn:
+    # Your SQL operations here
+    pass
+```
+
+## Circuit Breaker States
+
+The circuit breaker operates in three states:
+
+1. **Closed**: Normal operation, telemetry requests are allowed
+2. **Open**: Circuit breaker is open, telemetry requests are blocked
+3. **Half-Open**: Testing state, limited telemetry requests are allowed
+
+
+## Performance Impact
+
+The circuit breaker has minimal performance impact on SQL operations:
+
+- Circuit breaker only affects telemetry requests, not SQL queries
+- When circuit breaker is open, telemetry requests are simply skipped
+- No additional latency is added to successful operations
+
+## Best Practices
+
+1. **Keep circuit breaker enabled**: The default configuration works well for most use cases
+2. **Don't disable unless necessary**: Circuit breaker provides important protection against telemetry failures
+3. **Monitor application logs**: Circuit breaker state changes are logged for troubleshooting

pyproject.toml

@@ -26,6 +26,7 @@ pyarrow = [
     { version = ">=18.0.0", python = ">=3.13", optional=true }
 ]
 pyjwt = "^2.0.0"
+pybreaker = "^1.0.0"
 requests-kerberos = {version = "^0.15.0", optional = true}
 
 

src/databricks/sql/auth/common.py

@@ -51,6 +51,8 @@ def __init__(
         pool_connections: Optional[int] = None,
         pool_maxsize: Optional[int] = None,
         user_agent: Optional[str] = None,
+        # Telemetry circuit breaker configuration
+        telemetry_circuit_breaker_enabled: Optional[bool] = None,
     ):
         self.hostname = hostname
         self.access_token = access_token
@@ -83,6 +85,9 @@ def __init__(
         self.pool_connections = pool_connections or 10
         self.pool_maxsize = pool_maxsize or 20
         self.user_agent = user_agent
+        
+        # Telemetry circuit breaker configuration
+        self.telemetry_circuit_breaker_enabled = telemetry_circuit_breaker_enabled if telemetry_circuit_breaker_enabled is not None else True
 
 
 def get_effective_azure_login_app_id(hostname) -> str:

src/databricks/sql/telemetry/circuit_breaker_manager.py

@@ -0,0 +1,231 @@
+"""
+Circuit breaker implementation for telemetry requests.
+
+This module provides circuit breaker functionality to prevent telemetry failures
+from impacting the main SQL operations. It uses pybreaker library to implement
+the circuit breaker pattern with configurable thresholds and timeouts.
+"""
+
+import logging
+import threading
+from typing import Dict, Optional, Any
+from dataclasses import dataclass
+
+import pybreaker
+from pybreaker import CircuitBreaker, CircuitBreakerError
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CircuitBreakerConfig:
+    """Configuration for circuit breaker behavior."""
+    
+    # Failure threshold percentage (0.0 to 1.0)
+    failure_threshold: float = 0.5
+    
+    # Minimum number of calls before circuit can open
+    minimum_calls: int = 20
+    
+    # Time window for counting failures (in seconds)
+    timeout: int = 30
+    
+    # Time to wait before trying to close circuit (in seconds)
+    reset_timeout: int = 30
+    
+    # Expected exception types that should trigger circuit breaker
+    expected_exception: tuple = (Exception,)
+    
+    # Name for the circuit breaker (for logging)
+    name: str = "telemetry-circuit-breaker"
+
+
+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.
+    """
+    
+    _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:
+        """
+        Get or create a circuit breaker instance for the specified host.
+        
+        Args:
+            host: The hostname for which to get the circuit breaker
+            
+        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)
+                logger.debug("Created circuit breaker for host: %s", host)
+            
+            return cls._instances[host]
+    
+    @classmethod
+    def _create_circuit_breaker(cls, host: str) -> CircuitBreaker:
+        """
+        Create a new circuit breaker instance for the specified host.
+        
+        Args:
+            host: The hostname for the circuit breaker
+            
+        Returns:
+            New CircuitBreaker instance
+        """
+        config = cls._config
+        
+        # Create circuit breaker with configuration
+        breaker = CircuitBreaker(
+            fail_max=config.minimum_calls,
+            reset_timeout=config.reset_timeout,
+            name=f"{config.name}-{host}"
+        )
+        
+        # Set failure threshold
+        breaker.failure_threshold = config.failure_threshold
+        
+        # Add state change listeners for logging
+        breaker.add_listener(cls._on_state_change)
+        
+        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"
+        )
+        breaker.failure_threshold = 1.0  # 100% failure threshold
+        return breaker
+    
+    @classmethod
+    def _on_state_change(cls, old_state: str, new_state: str, breaker: CircuitBreaker) -> None:
+        """
+        Handle circuit breaker state changes.
+        
+        Args:
+            old_state: Previous state of the circuit breaker
+            new_state: New state of the circuit breaker
+            breaker: The circuit breaker instance
+        """
+        logger.info(
+            "Circuit breaker state changed from %s to %s for %s",
+            old_state, new_state, breaker.name
+        )
+        
+        if new_state == "open":
+            logger.warning(
+                "Circuit breaker opened for %s - telemetry requests will be blocked",
+                breaker.name
+            )
+        elif new_state == "closed":
+            logger.info(
+                "Circuit breaker closed for %s - telemetry requests will be allowed",
+                breaker.name
+            )
+        elif new_state == "half-open":
+            logger.info(
+                "Circuit breaker half-open for %s - testing telemetry requests",
+                breaker.name
+            )
+    
+    @classmethod
+    def get_circuit_breaker_state(cls, host: str) -> str:
+        """
+        Get the current state of the circuit breaker for a host.
+        
+        Args:
+            host: The hostname
+            
+        Returns:
+            Current state of the circuit breaker
+        """
+        if not cls._config:
+            return "disabled"
+        
+        with cls._lock:
+            if host not in cls._instances:
+                return "not_initialized"
+            
+            breaker = cls._instances[host]
+            return breaker.current_state
+    
+    @classmethod
+    def reset_circuit_breaker(cls, host: str) -> None:
+        """
+        Reset the circuit breaker for a host to closed state.
+        
+        Args:
+            host: The hostname
+        """
+        with cls._lock:
+            if host in cls._instances:
+                # pybreaker doesn't have a reset method, we need to recreate the breaker
+                del cls._instances[host]
+                logger.info("Reset circuit breaker for host: %s", host)
+    
+    @classmethod
+    def clear_circuit_breaker(cls, host: str) -> None:
+        """
+        Remove the circuit breaker instance for a host.
+        
+        Args:
+            host: The hostname
+        """
+        with cls._lock:
+            if host in cls._instances:
+                del cls._instances[host]
+                logger.debug("Cleared circuit breaker for host: %s", host)
+    
+    @classmethod
+    def clear_all_circuit_breakers(cls) -> None:
+        """Clear all circuit breaker instances."""
+        with cls._lock:
+            cls._instances.clear()
+            logger.debug("Cleared all circuit breakers")
+
+
+def is_circuit_breaker_error(exception: Exception) -> bool:
+    """
+    Check if an exception is a circuit breaker error.
+    
+    Args:
+        exception: The exception to check
+        
+    Returns:
+        True if the exception is a circuit breaker error
+    """
+    return isinstance(exception, CircuitBreakerError)