removed caplog

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/exc.py

@@ -19,11 +19,13 @@ def __init__(
         self.context = context or {}
 
         error_name = self.__class__.__name__
+        if session_id_hex:
+            from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
-        from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
-
-        telemetry_client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
-        telemetry_client.export_failure_log(error_name, self.message)
+            telemetry_client = TelemetryClientFactory.get_telemetry_client(
+                session_id_hex
+            )
+            telemetry_client.export_failure_log(error_name, self.message)
 
     def __str__(self):
         return self.message

src/databricks/sql/telemetry/telemetry_client.py

@@ -41,13 +41,8 @@ class TelemetryHTTPAdapter(HTTPAdapter):
     """
 
     def send(self, request, **kwargs):
-        # The DatabricksRetryPolicy needs state set before the first attempt.
-        if isinstance(self.max_retries, DatabricksRetryPolicy):
-            # Telemetry requests are idempotent and safe to retry. We use CommandType.OTHER
-            # to signal this to the retry policy, bypassing stricter rules for commands
-            # like ExecuteStatement.
-            self.max_retries.command_type = CommandType.OTHER
-            self.max_retries.start_retry_timer()
+        self.max_retries.command_type = CommandType.OTHER
+        self.max_retries.start_retry_timer()
         return super().send(request, **kwargs)
 
 

tests/e2e/test_telemetry_retry.py

@@ -1,50 +1,34 @@
-# tests/e2e/test_telemetry_retry.py
-
 import pytest
-import logging
 from unittest.mock import patch, MagicMock
-from functools import wraps
-import time
-from concurrent.futures import Future
+import io
 
-# Imports for the code being tested
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.telemetry.models.event import DriverConnectionParameters, HostDetails, DatabricksClientType
 from databricks.sql.telemetry.models.enums import AuthMech
-from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
-
-# Imports for mocking the network layer correctly
-from urllib3.connectionpool import HTTPSConnectionPool
-from urllib3.exceptions import MaxRetryError
-from requests.exceptions import ConnectionError as RequestsConnectionError
+from databricks.sql.auth.retry import DatabricksRetryPolicy
 
 PATCH_TARGET = 'urllib3.connectionpool.HTTPSConnectionPool._get_conn'
 
-# Helper to create a mock that looks and acts like a urllib3.response.HTTPResponse.
-def create_urllib3_response(status, headers=None, body=b'{}'):
-    """Create a proper mock response that simulates urllib3's HTTPResponse"""
-    mock_response = MagicMock()
-    mock_response.status = status
-    mock_response.headers = headers or {}
-    mock_response.msg = headers or {}  # For urllib3~=1.0 compatibility
-    mock_response.data = body
-    mock_response.read.return_value = body
-    mock_response.get_redirect_location.return_value = False
-    mock_response.closed = False
-    mock_response.isclosed.return_value = False
-    return mock_response
+def create_mock_conn(responses):
+    """Creates a mock connection object whose getresponse() method yields a series of responses."""
+    mock_conn = MagicMock()
+    mock_http_responses = []
+    for resp in responses:
+        mock_http_response = MagicMock()
+        mock_http_response.status = resp.get("status")
+        mock_http_response.headers = resp.get("headers", {})
+        body = resp.get("body", b'{}')
+        mock_http_response.fp = io.BytesIO(body)
+        def release():
+            mock_http_response.fp.close()
+        mock_http_response.release_conn = release
+        mock_http_responses.append(mock_http_response)
+    mock_conn.getresponse.side_effect = mock_http_responses
+    return mock_conn
 
-@pytest.mark.usefixtures("caplog")
 class TestTelemetryClientRetries:
-    """
-    Test suite for verifying the retry mechanism of the TelemetryClient.
-    This suite patches the low-level urllib3 connection to correctly
-    trigger and test the retry logic configured in the requests adapter.
-    """
-
     @pytest.fixture(autouse=True)
-    def setup_and_teardown(self, caplog):
-        caplog.set_level(logging.DEBUG)
+    def setup_and_teardown(self):
         TelemetryClientFactory._initialized = False
         TelemetryClientFactory._clients = {}
         TelemetryClientFactory._executor = None
@@ -55,159 +39,94 @@ def setup_and_teardown(self, caplog):
         TelemetryClientFactory._clients = {}
         TelemetryClientFactory._executor = None
 
-    def get_client(self, session_id, total_retries=3):
+    def get_client(self, session_id, num_retries=3):
+        """
+        Configures a client with a specific number of retries.
+        """
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=True,
             session_id_hex=session_id,
             auth_provider=None,
             host_url="test.databricks.com",
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)
-
+        
         retry_policy = DatabricksRetryPolicy(
             delay_min=0.01,
             delay_max=0.02,
             stop_after_attempts_duration=2.0,
-            stop_after_attempts_count=total_retries,
+            stop_after_attempts_count=num_retries, 
             delay_default=0.1,
             force_dangerous_codes=[],
-            urllib3_kwargs={'total': total_retries}
+            urllib3_kwargs={'total': num_retries}
         )
         adapter = client._session.adapters.get("https://")
         adapter.max_retries = retry_policy
         return client, adapter
 
-    def wait_for_async_request(self, timeout=2.0):
-        """Wait for async telemetry request to complete"""
-        start_time = time.time()
-        while time.time() - start_time < timeout:
-            if TelemetryClientFactory._executor and TelemetryClientFactory._executor._threads:
-                # Wait a bit more for threads to complete
-                time.sleep(0.1)
-            else:
-                break
-        time.sleep(0.1)  # Extra buffer for completion
-
     def test_success_no_retry(self):
         client, _ = self.get_client("session-success")
         params = DriverConnectionParameters(
-            http_path="test-path",
-            mode=DatabricksClientType.THRIFT,
+            http_path="test-path", mode=DatabricksClientType.THRIFT,
             host_info=HostDetails(host_url="test.databricks.com", port=443),
             auth_mech=AuthMech.PAT
         )
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(200)
-            
+        mock_responses = [{"status": 200}]
+        
+        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
             client.export_initial_telemetry_log(params, "test-agent")
-            self.wait_for_async_request()
             TelemetryClientFactory.close(client._session_id_hex)
 
             mock_get_conn.return_value.getresponse.assert_called_once()
-
-    def test_retry_on_503_then_succeeds(self):
-        client, _ = self.get_client("session-retry-once")
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.side_effect = [
-                create_urllib3_response(503),
-                create_urllib3_response(200),
-            ]
-            
-            client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
-            TelemetryClientFactory.close(client._session_id_hex)
-            
-            assert mock_get_conn.return_value.getresponse.call_count == 2
-
-    def test_respects_retry_after_header(self, caplog):
-        client, _ = self.get_client("session-retry-after")
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.side_effect = [
-                create_urllib3_response(429, headers={'Retry-After': '1'}),  # Use integer seconds to avoid parsing issues
-                create_urllib3_response(200)
-            ]
-            
+        client, _ = self.get_client("session-retry-once", num_retries=1)
+        mock_responses = [{"status": 503}, {"status": 200}]
+        
+        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
             TelemetryClientFactory.close(client._session_id_hex)
 
-            # Check that the request was retried (should be 2 calls: initial + 1 retry)
             assert mock_get_conn.return_value.getresponse.call_count == 2
-            assert "Retrying after" in caplog.text
 
-    def test_exceeds_retry_count_limit(self, caplog):
-        client, _ = self.get_client("session-exceed-limit", total_retries=3)
-        expected_call_count = 4
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(503)
-            
+    @pytest.mark.parametrize(
+    "status_code, description",
+    [
+        (401, "Unauthorized"),
+        (403, "Forbidden"),
+        (501, "Not Implemented"),
+    ],
+    )
+    def test_non_retryable_status_codes_are_not_retried(self, status_code, description):
+        """
+        Verifies that terminal error codes (401, 403, 501, etc.) are not retried.
+        """
+        # Use the status code in the session ID for easier debugging if it fails
+        client, _ = self.get_client(f"session-{status_code}")
+        mock_responses = [{"status": status_code}]
+
+        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
             TelemetryClientFactory.close(client._session_id_hex)
-            
-            assert mock_get_conn.return_value.getresponse.call_count == expected_call_count
-            assert "Telemetry request failed with exception" in caplog.text
-            assert "Max retries exceeded" in caplog.text
 
-    def test_no_retry_on_401_unauthorized(self, caplog):
-        """Test that 401 responses are not retried (per retry policy)"""
-        client, _ = self.get_client("session-401")
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(401)
-            
-            client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
-            TelemetryClientFactory.close(client._session_id_hex)
-            
-            # 401 should not be retried based on the retry policy
             mock_get_conn.return_value.getresponse.assert_called_once()
-            assert "Telemetry request failed with status code: 401" in caplog.text
 
-    def test_retries_on_400_bad_request(self, caplog):
-        """Test that 400 responses are retried (this is the current behavior for telemetry)"""
-        client, _ = self.get_client("session-400")
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(400)
-            
-            client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
-            TelemetryClientFactory.close(client._session_id_hex)
-            
-            # Based on the logs, 400 IS being retried (this is the actual behavior for CommandType.OTHER)
-            expected_call_count = 4  # total + 1 (initial + 3 retries)
-            assert mock_get_conn.return_value.getresponse.call_count == expected_call_count
-            assert "Telemetry request failed with exception" in caplog.text
-            assert "Max retries exceeded" in caplog.text
-
-    def test_no_retry_on_403_forbidden(self, caplog):
-        """Test that 403 responses are not retried (per retry policy)"""
-        client, _ = self.get_client("session-403")
-        with patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(403)
-            
+    def test_respects_retry_after_header(self):
+        client, _ = self.get_client("session-retry-after", num_retries=1)
+        mock_responses = [{"status": 429, "headers": {'Retry-After': '1'}}, {"status": 200}]
+        
+        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
             TelemetryClientFactory.close(client._session_id_hex)
 
-            # 403 should not be retried based on the retry policy
-            mock_get_conn.return_value.getresponse.assert_called_once()
-            assert "Telemetry request failed with status code: 403" in caplog.text
+            assert mock_get_conn.return_value.getresponse.call_count == 2
 
-    def test_retry_policy_command_type_is_set_to_other(self):
-        client, adapter = self.get_client("session-command-type")
+    def test_exceeds_retry_count_limit(self):
+        num_retries = 3
+        expected_total_calls = num_retries + 1 
+        client, _ = self.get_client("session-exceed-limit", num_retries=num_retries)
+        mock_responses = [{"status": 503}] * expected_total_calls
 
-        original_send = adapter.send
-        @wraps(original_send)
-        def wrapper(request, **kwargs):
-            assert adapter.max_retries.command_type == CommandType.OTHER
-            return original_send(request, **kwargs)
-
-        with patch.object(adapter, 'send', side_effect=wrapper, autospec=True), \
-             patch(PATCH_TARGET) as mock_get_conn:
-            mock_get_conn.return_value.getresponse.return_value = create_urllib3_response(200)
-            
+        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
-            self.wait_for_async_request()
             TelemetryClientFactory.close(client._session_id_hex)
 
-            assert adapter.send.call_count == 1
+            assert mock_get_conn.return_value.getresponse.call_count == expected_total_calls

tests/unit/test_telemetry.py

@@ -331,26 +331,6 @@ class TestBaseClient(BaseTelemetryClient):
         with pytest.raises(TypeError):
             TestBaseClient()  # Can't instantiate abstract class
 
-    def test_telemetry_http_adapter_configuration(self, telemetry_client_setup):
-        """Test that TelemetryHTTPAdapter is properly configured with correct retry parameters."""
-        from databricks.sql.telemetry.telemetry_client import TelemetryHTTPAdapter
-        from databricks.sql.auth.retry import DatabricksRetryPolicy
-        
-        client = telemetry_client_setup["client"]
-        
-        # Verify that the session has the TelemetryHTTPAdapter mounted
-        adapter = client._session.adapters.get("https://")
-        assert isinstance(adapter, TelemetryHTTPAdapter)
-        assert isinstance(adapter.max_retries, DatabricksRetryPolicy)
-        
-        # Verify that the retry policy has the correct static configuration
-        retry_policy = adapter.max_retries
-        assert retry_policy.delay_min == client.TELEMETRY_RETRY_DELAY_MIN
-        assert retry_policy.delay_max == client.TELEMETRY_RETRY_DELAY_MAX
-        assert retry_policy.stop_after_attempts_count == client.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT
-        assert retry_policy.stop_after_attempts_duration == client.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION
-
-
 class TestTelemetryHelper:
     """Tests for the TelemetryHelper class."""