fix minor issues & improvement

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

Author: nikhilsuri-db

src/databricks/sql/auth/common.py

@@ -84,11 +84,7 @@ def __init__(
         self.pool_connections = pool_connections or 10
         self.pool_maxsize = pool_maxsize or 20
         self.user_agent = user_agent
-        self.telemetry_circuit_breaker_enabled = (
-            telemetry_circuit_breaker_enabled
-            if telemetry_circuit_breaker_enabled is not None
-            else False
-        )
+        self.telemetry_circuit_breaker_enabled = bool(telemetry_circuit_breaker_enabled)
 
 
 def get_effective_azure_login_app_id(hostname) -> str:

src/databricks/sql/telemetry/telemetry_push_client.py

@@ -129,10 +129,8 @@ def request(
             )
         except CircuitBreakerError as e:
             logger.warning(
-                "Circuit breaker is open for host %s, blocking telemetry request to %s: %s",
+                "Circuit breaker is open for host %s, blocking telemetry request",
                 self._host,
-                url,
-                e,
             )
             raise
         except Exception as e:
@@ -150,21 +148,18 @@ def request_context(
     ):
         """Context manager for making HTTP requests with circuit breaker protection."""
         try:
-            # Use circuit breaker to protect the request
-            def _make_request():
-                with self._delegate.request_context(
-                    method, url, headers, **kwargs
-                ) as response:
-                    return response
-
-            response = self._circuit_breaker.call(_make_request)
-            yield response
+            # Keep the context manager open while yielding the response
+            # Circuit breaker will track failures through the exception handling
+            with self._delegate.request_context(
+                method, url, headers, **kwargs
+            ) as response:
+                # Record success with circuit breaker before yielding
+                self._circuit_breaker.call(lambda: None)
+                yield response
         except CircuitBreakerError as e:
             logger.warning(
-                "Circuit breaker is open for host %s, blocking telemetry request to %s: %s",
+                "Circuit breaker is open for host %s, blocking telemetry request",
                 self._host,
-                url,
-                e,
             )
             raise
         except Exception as e:

tests/unit/test_circuit_breaker_http_client.py

@@ -223,7 +223,7 @@ def test_circuit_breaker_recovers_after_success(self):
             client.request(HttpMethod.POST, "https://test.com", {})
 
         # Wait for reset timeout
-        time.sleep(RESET_TIMEOUT + 0.1)
+        time.sleep(RESET_TIMEOUT + 1.0)
 
         # Simulate successful calls
         self.mock_delegate.request.side_effect = None

tests/unit/test_circuit_breaker_manager.py

@@ -155,7 +155,7 @@ def failing_func():
         assert breaker.current_state == "open"
 
         # Wait for reset timeout
-        time.sleep(RESET_TIMEOUT + 0.1)
+        time.sleep(RESET_TIMEOUT + 1.0)
 
         # Try successful call to close circuit breaker
         def successful_func():

tests/unit/test_telemetry_circuit_breaker_integration.py

@@ -21,9 +21,6 @@ def setup_method(self):
         # Create mock client context with circuit breaker config
         self.client_context = Mock(spec=ClientContext)
         self.client_context.telemetry_circuit_breaker_enabled = True
-        self.client_context.telemetry_circuit_breaker_failure_threshold = (
-            0.1  # 10% failure rate
-        )
         self.client_context.telemetry_circuit_breaker_minimum_calls = 2
         self.client_context.telemetry_circuit_breaker_timeout = 30
         self.client_context.telemetry_circuit_breaker_reset_timeout = (
@@ -274,7 +271,6 @@ def setup_method(self):
         """Set up test fixtures."""
         self.client_context = Mock(spec=ClientContext)
         self.client_context.telemetry_circuit_breaker_enabled = True
-        self.client_context.telemetry_circuit_breaker_failure_threshold = 0.1
         self.client_context.telemetry_circuit_breaker_minimum_calls = 2
         self.client_context.telemetry_circuit_breaker_timeout = 30
         self.client_context.telemetry_circuit_breaker_reset_timeout = 1

tests/unit/test_telemetry_push_client.py

@@ -256,7 +256,7 @@ def test_circuit_breaker_recovers_after_success(self):
             client.request(HttpMethod.POST, "https://test.com", {})
 
         # Wait for reset timeout
-        time.sleep(RESET_TIMEOUT + 0.1)
+        time.sleep(RESET_TIMEOUT + 1.0)
 
         # Simulate successful calls
         self.mock_delegate.request.side_effect = None