raise CB only for 429/503

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

Author: nikhilsuri-db

src/databricks/sql/exc.py

@@ -131,5 +131,3 @@ class CursorAlreadyClosedError(RequestError):
 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

@@ -179,16 +179,3 @@ def _create_noop_circuit_breaker(cls) -> CircuitBreaker:
             name="noop-circuit-breaker",
         )
         return breaker
-
-
-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)

src/databricks/sql/telemetry/telemetry_client.py

@@ -272,23 +272,14 @@ 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.
-
-        The push client implementation handles circuit breaker logic internally,
-        so this method just forwards the request and handles any errors generically.
-        """
+        """Helper method to send telemetry using the unified HTTP client."""
         try:
             response = self._telemetry_push_client.request(
                 HttpMethod.POST, url, body=data, headers=headers, timeout=timeout
             )
             return response
         except Exception as e:
-            logger.debug(
-                "Failed to send telemetry for connection %s: %s",
-                self._session_id_hex,
-                e,
-            )
+            logger.error("Failed to send telemetry with unified client: %s", e)
             raise
 
     def _telemetry_request_callback(self, future, sent_count: int):

src/databricks/sql/telemetry/telemetry_push_client.py

@@ -19,10 +19,7 @@
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.http import HttpMethod
 from databricks.sql.exc import TelemetryRateLimitError, RequestError
-from databricks.sql.telemetry.circuit_breaker_manager import (
-    CircuitBreakerManager,
-    is_circuit_breaker_error,
-)
+from databricks.sql.telemetry.circuit_breaker_manager import CircuitBreakerManager
 
 logger = logging.getLogger(__name__)
 

tests/unit/test_circuit_breaker_manager.py

@@ -9,7 +9,7 @@
 
 from databricks.sql.telemetry.circuit_breaker_manager import (
     CircuitBreakerManager,
-    is_circuit_breaker_error,
+    CircuitBreakerConfig,
     DEFAULT_MINIMUM_CALLS as MINIMUM_CALLS,
     DEFAULT_RESET_TIMEOUT as RESET_TIMEOUT,
     DEFAULT_NAME as CIRCUIT_BREAKER_NAME,
@@ -24,10 +24,13 @@ 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."""
@@ -83,37 +86,19 @@ def get_breaker(host):
         assert all(b is host0_breakers[0] for b in host0_breakers)
 
 
-class TestCircuitBreakerErrorDetection:
-    """Test cases for circuit breaker error detection."""
-
-    def test_is_circuit_breaker_error_true(self):
-        """Test detecting circuit breaker errors."""
-        error = CircuitBreakerError("Circuit breaker is open")
-        assert is_circuit_breaker_error(error) is True
-
-    def test_is_circuit_breaker_error_false(self):
-        """Test detecting non-circuit breaker errors."""
-        error = ValueError("Some other error")
-        assert is_circuit_breaker_error(error) is False
-
-        error = RuntimeError("Another error")
-        assert is_circuit_breaker_error(error) is False
-
-    def test_is_circuit_breaker_error_none(self):
-        """Test with None input."""
-        assert is_circuit_breaker_error(None) is False
-
-
 class TestCircuitBreakerIntegration:
     """Integration tests for circuit breaker functionality."""
 
     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."""