remove cb congig class to constants

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

Author: nikhilsuri-db

src/databricks/sql/common/unified_http_client.py

@@ -28,6 +28,42 @@
 logger = logging.getLogger(__name__)
 
 
+def _extract_http_status_from_max_retry_error(e: MaxRetryError) -> Optional[int]:
+    """
+    Extract HTTP status code from MaxRetryError if available.
+
+    urllib3 structures MaxRetryError in different ways depending on the failure scenario:
+    - e.reason.response.status: Most common case when retries are exhausted
+    - e.response.status: Alternate structure in some scenarios
+
+    Args:
+        e: MaxRetryError exception from urllib3
+
+    Returns:
+        HTTP status code as int if found, None otherwise
+    """
+    # Try primary structure: e.reason.response.status
+    if (
+        hasattr(e, "reason")
+        and e.reason is not None
+        and hasattr(e.reason, "response")
+        and e.reason.response is not None
+    ):
+        http_code = getattr(e.reason.response, "status", None)
+        if http_code is not None:
+            return http_code
+
+    # Try alternate structure: e.response.status
+    if (
+        hasattr(e, "response")
+        and e.response is not None
+        and hasattr(e.response, "status")
+    ):
+        return e.response.status
+
+    return None
+
+
 class UnifiedHttpClient:
     """
     Unified HTTP client for all Databricks SQL connector HTTP operations.
@@ -265,23 +301,8 @@ def request_context(
         except MaxRetryError as e:
             logger.error("HTTP request failed after retries: %s", e)
 
-            # Try to extract HTTP status code from the MaxRetryError
-            http_code = None
-            if (
-                hasattr(e, "reason")
-                and e.reason is not None
-                and hasattr(e.reason, "response")
-                and e.reason.response is not None
-            ):
-                # The reason may contain a response object with status
-                http_code = getattr(e.reason.response, "status", None)
-            elif (
-                hasattr(e, "response")
-                and e.response is not None
-                and hasattr(e.response, "status")
-            ):
-                # Or the error itself may have a response
-                http_code = e.response.status
+            # Extract HTTP status code from MaxRetryError if available
+            http_code = _extract_http_status_from_max_retry_error(e)
 
             context = {}
             if http_code is not None:

src/databricks/sql/telemetry/circuit_breaker_manager.py

@@ -3,13 +3,12 @@
 
 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.
+the circuit breaker pattern.
 """
 
 import logging
 import threading
-from typing import Dict, Optional, Any
-from dataclasses import dataclass
+from typing import Dict
 
 import pybreaker
 from pybreaker import CircuitBreaker, CircuitBreakerError, CircuitBreakerListener
@@ -18,10 +17,10 @@
 
 logger = logging.getLogger(__name__)
 
-# Circuit Breaker Configuration Constants
-DEFAULT_MINIMUM_CALLS = 20
-DEFAULT_RESET_TIMEOUT = 30
-DEFAULT_NAME = "telemetry-circuit-breaker"
+# Circuit Breaker Constants
+MINIMUM_CALLS = 20  # Number of failures before circuit opens
+RESET_TIMEOUT = 30  # Seconds to wait before trying to close circuit
+NAME_PREFIX = "telemetry-circuit-breaker"
 
 # Circuit Breaker State Constants (used in logging)
 CIRCUIT_BREAKER_STATE_OPEN = "open"
@@ -73,47 +72,16 @@ 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.
+    Creates and caches circuit breaker instances per host to ensure 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:
@@ -126,56 +94,16 @@ 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)
+                breaker = CircuitBreaker(
+                    fail_max=MINIMUM_CALLS,
+                    reset_timeout=RESET_TIMEOUT,
+                    name=f"{NAME_PREFIX}-{host}",
+                )
+                # Add state change listener for logging
+                breaker.add_listener(CircuitBreakerStateListener())
+                cls._instances[host] = breaker
                 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
-        if config is None:
-            raise RuntimeError("CircuitBreakerManager not initialized")
-
-        # Create circuit breaker with configuration
-        breaker = CircuitBreaker(
-            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

src/databricks/sql/telemetry/telemetry_client.py

@@ -196,7 +196,7 @@ def __init__(
 
         # Create telemetry push client based on circuit breaker enabled flag
         if client_context.telemetry_circuit_breaker_enabled:
-            # Create circuit breaker telemetry push client with fixed configuration
+            # Create circuit breaker telemetry push client (circuit breakers created on-demand)
             self._telemetry_push_client: ITelemetryPushClient = (
                 CircuitBreakerTelemetryPushClient(
                     TelemetryPushClient(self._http_client),

src/databricks/sql/telemetry/telemetry_push_client.py

@@ -91,12 +91,18 @@ def _create_mock_success_response(self) -> BaseHTTPResponse:
 
         This allows telemetry to fail silently without raising exceptions.
         """
-        from unittest.mock import Mock
+        # Create a simple object that mimics BaseHTTPResponse interface
+        class _MockTelemetryResponse:
+            """Simple response object for silently handling circuit breaker state."""
 
-        mock_response = Mock(spec=BaseHTTPResponse)
-        mock_response.status = 200
-        mock_response.data = b'{"numProtoSuccess": 0, "errors": []}'
-        return mock_response
+            status = 200
+            # Include all required fields for TelemetryResponse dataclass
+            data = b'{"numProtoSuccess": 0, "numSuccess": 0, "numRealtimeSuccess": 0, "errors": []}'
+
+            def close(self):
+                pass
+
+        return _MockTelemetryResponse()
 
     def request(
         self,

tests/unit/test_circuit_breaker_http_client.py

@@ -91,6 +91,7 @@ def test_request_enabled_other_error(self):
         response = self.client.request(HttpMethod.POST, "https://test.com", {})
         assert response is not None
         assert response.status == 200
+        assert b"numProtoSuccess" in response.data
 
     def test_is_circuit_breaker_enabled(self):
         """Test checking if circuit breaker is enabled."""
@@ -146,14 +147,12 @@ def test_circuit_breaker_opens_after_failures(self):
         """Test that circuit breaker opens after repeated failures (429/503 errors)."""
         from databricks.sql.telemetry.circuit_breaker_manager import (
             CircuitBreakerManager,
-            CircuitBreakerConfig,
-            DEFAULT_MINIMUM_CALLS as MINIMUM_CALLS,
+            MINIMUM_CALLS,
         )
         from databricks.sql.exc import TelemetryRateLimitError
 
         # Clear any existing state
         CircuitBreakerManager._instances.clear()
-        CircuitBreakerManager.initialize(CircuitBreakerConfig())
 
         client = CircuitBreakerTelemetryPushClient(self.mock_delegate, self.host)
 
@@ -177,15 +176,13 @@ def test_circuit_breaker_recovers_after_success(self):
         """Test that circuit breaker recovers after successful calls."""
         from databricks.sql.telemetry.circuit_breaker_manager import (
             CircuitBreakerManager,
-            CircuitBreakerConfig,
-            DEFAULT_MINIMUM_CALLS as MINIMUM_CALLS,
-            DEFAULT_RESET_TIMEOUT as RESET_TIMEOUT,
+            MINIMUM_CALLS,
+            RESET_TIMEOUT,
         )
         import time
 
         # Clear any existing state
         CircuitBreakerManager._instances.clear()
-        CircuitBreakerManager.initialize(CircuitBreakerConfig())
 
         client = CircuitBreakerTelemetryPushClient(self.mock_delegate, self.host)
 
@@ -198,11 +195,13 @@ def test_circuit_breaker_recovers_after_success(self):
         for i in range(MINIMUM_CALLS + 5):
             response = client.request(HttpMethod.POST, "https://test.com", {})
             assert response.status == 200  # Returns mock success
+            assert b"numProtoSuccess" in response.data
 
         # Circuit should be open now - still returns mock response
         response = client.request(HttpMethod.POST, "https://test.com", {})
         assert response is not None
         assert response.status == 200  # Mock success response
+        assert b"numProtoSuccess" in response.data
 
         # Wait for reset timeout
         time.sleep(RESET_TIMEOUT + 1.0)

tests/unit/test_circuit_breaker_manager.py

@@ -9,10 +9,9 @@
 
 from databricks.sql.telemetry.circuit_breaker_manager import (
     CircuitBreakerManager,
-    CircuitBreakerConfig,
-    DEFAULT_MINIMUM_CALLS as MINIMUM_CALLS,
-    DEFAULT_RESET_TIMEOUT as RESET_TIMEOUT,
-    DEFAULT_NAME as CIRCUIT_BREAKER_NAME,
+    MINIMUM_CALLS,
+    RESET_TIMEOUT,
+    NAME_PREFIX as CIRCUIT_BREAKER_NAME,
 )
 from pybreaker import CircuitBreakerError
 
@@ -24,13 +23,10 @@ def setup_method(self):
         """Set up test fixtures."""
         # Clear any existing instances
         CircuitBreakerManager._instances.clear()
-        # Initialize with default config
-        CircuitBreakerManager.initialize(CircuitBreakerConfig())
 
     def teardown_method(self):
         """Clean up after tests."""
         CircuitBreakerManager._instances.clear()
-        CircuitBreakerManager._config = None
 
     def test_get_circuit_breaker_creates_instance(self):
         """Test getting circuit breaker creates instance with correct config."""
@@ -60,6 +56,16 @@ def test_get_circuit_breaker_creates_breaker(self):
         assert breaker is not None
         assert breaker.current_state in ["closed", "open", "half-open"]
 
+    def test_circuit_breaker_reused_for_same_host(self):
+        """Test that circuit breakers are reused for the same host."""
+        # Get circuit breaker for a host
+        breaker1 = CircuitBreakerManager.get_circuit_breaker("host1.example.com")
+        assert breaker1 is not None
+
+        # Get circuit breaker again for the same host - should be SAME instance
+        breaker2 = CircuitBreakerManager.get_circuit_breaker("host1.example.com")
+        assert breaker2 is breaker1  # Same instance, state preserved across calls
+
     def test_thread_safety(self):
         """Test thread safety of circuit breaker manager."""
         results = []
@@ -92,13 +98,10 @@ class TestCircuitBreakerIntegration:
     def setup_method(self):
         """Set up test fixtures."""
         CircuitBreakerManager._instances.clear()
-        # Initialize with default config
-        CircuitBreakerManager.initialize(CircuitBreakerConfig())
 
     def teardown_method(self):
         """Clean up after tests."""
         CircuitBreakerManager._instances.clear()
-        CircuitBreakerManager._config = None
 
     def test_circuit_breaker_state_transitions(self):
         """Test circuit breaker state transitions."""