Reverting retry behavior on 429s/503s to how it worked in 2.9.3 (#349)

Signed-off-by: Ben Cassell <[email protected]>

Author: benc-db

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+# 3.0.4 (TBD)
+
+- Revert retry-after behavior to be exponential backoff
+
 # 3.0.3 (2024-02-02)
 
 - Add support in-house OAuth on GCP (#338)
@@ -50,15 +54,15 @@
 
 ## 2.9.2 (2023-08-17)
 
-__Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed. The log changes are incorporated into version 2.9.3 and greater.__
+**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed. The log changes are incorporated into version 2.9.3 and greater.**
 
 - Other: Add `examples/v3_retries_query_execute.py` (#199)
 - Other: suppress log message when `_enable_v3_retries` is not `True` (#199)
 - Other: make this connector backwards compatible with `urllib3>=1.0.0` (#197)
 
 ## 2.9.1 (2023-08-11)
 
-__Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed.__
+**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed.**
 
 - Other: Explicitly pin urllib3 to ^2.0.0 (#191)
 
@@ -111,6 +115,7 @@ __Note: this release was yanked from Pypi on 13 September 2023 due to compatibil
 - Other: Relax sqlalchemy required version as it was unecessarily strict.
 
 ## 2.5.0 (2023-04-14)
+
 - Add support for External Auth providers
 - Fix: Python HTTP proxies were broken
 - Other: All Thrift requests that timeout during connection will be automatically retried
@@ -132,8 +137,8 @@ __Note: this release was yanked from Pypi on 13 September 2023 due to compatibil
 
 ## 2.2.2 (2023-01-03)
 
-- Support custom oauth client id and redirect port 
-- Fix: Add none check on _oauth_persistence in DatabricksOAuthProvider
+- Support custom oauth client id and redirect port
+- Fix: Add none check on \_oauth_persistence in DatabricksOAuthProvider
 
 ## 2.2.1 (2022-11-29)
 
@@ -165,57 +170,71 @@ Huge thanks to @dbaxa for contributing this change!
 
 - Add retry logic for `GetOperationStatus` requests that fail with an `OSError`
 - Reorganised code to use Poetry for dependency management.
+
 ## 2.0.2 (2022-05-04)
+
 - Better exception handling in automatic connection close
 
 ## 2.0.1 (2022-04-21)
+
 - Fixed Pandas dependency in setup.cfg to be >= 1.2.0
 
 ## 2.0.0 (2022-04-19)
+
 - Initial stable release of V2
-- Added better support for complex types, so that in Databricks runtime 10.3+, Arrays, Maps and Structs will get 
+- Added better support for complex types, so that in Databricks runtime 10.3+, Arrays, Maps and Structs will get
   deserialized as lists, lists of tuples and dicts, respectively.
 - Changed the name of the metadata arg to http_headers
 
 ## 2.0.b2 (2022-04-04)
+
 - Change import of collections.Iterable to collections.abc.Iterable to make the library compatible with Python 3.10
 - Fixed bug with .tables method so that .tables works as expected with Unity-Catalog enabled endpoints
 
 ## 2.0.0b1 (2022-03-04)
+
 - Fix packaging issue (dependencies were not being installed properly)
 - Fetching timestamp results will now return aware instead of naive timestamps
 - The client will now default to using simplified error messages
 
 ## 2.0.0b (2022-02-08)
+
 - Initial beta release of V2. V2 is an internal re-write of large parts of the connector to use Databricks edge features. All public APIs from V1 remain.
-- Added Unity Catalog support (pass catalog and / or  schema key word args to the .connect method to select initial schema and catalog)
+- Added Unity Catalog support (pass catalog and / or schema key word args to the .connect method to select initial schema and catalog)
 
 ---
 
 **Note**: The code for versions prior to `v2.0.0b` is not contained in this repository. The below entries are included for reference only.
 
 ---
+
 ## 1.0.0 (2022-01-20)
+
 - Add operations for retrieving metadata
 - Add the ability to access columns by name on result rows
 - Add the ability to provide configuration settings on connect
 
 ## 0.9.4 (2022-01-10)
+
 - Improved logging and error messages.
 
 ## 0.9.3 (2021-12-08)
+
 - Add retries for 429 and 503 HTTP responses.
 
 ## 0.9.2 (2021-12-02)
+
 - (Bug fix) Increased Thrift requirement from 0.10.0 to 0.13.0 as 0.10.0 was in fact incompatible
 - (Bug fix) Fixed error message after query execution failed -SQLSTATE and Error message were misplaced
 
 ## 0.9.1 (2021-09-01)
+
 - Public Preview release, Experimental tag removed
 - minor updates in internal build/packaging
 - no functional changes
 
 ## 0.9.0 (2021-08-04)
+
 - initial (Experimental) release of pyhive-forked connector
 - Python DBAPI 2.0 (PEP-0249), thrift based
 - see docs for more info: https://docs.databricks.com/dev-tools/python-sql-connector.html

src/databricks/sql/auth/retry.py

@@ -283,8 +283,10 @@ def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:  # type: ignore
         """
         retry_after = self.get_retry_after(response)
         if retry_after:
-            self.check_proposed_wait(retry_after)
-            time.sleep(retry_after)
+            backoff = self.get_backoff_time()
+            proposed_wait = max(backoff, retry_after)
+            self.check_proposed_wait(proposed_wait)
+            time.sleep(proposed_wait)
             return True
 
         return False

tests/e2e/common/retry_test_mixins.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+import time
 from typing import List
 from unittest.mock import MagicMock, PropertyMock, patch
 
@@ -166,6 +167,33 @@ def test_retry_max_count_not_exceeded(self):
                     pass
             assert mock_obj.return_value.getresponse.call_count == 6
 
+    def test_retry_exponential_backoff(self):
+        """GIVEN the retry policy is configured for reasonable exponential backoff
+        WHEN the server sends nothing but 429 responses with retry-afters
+        THEN the connector will use those retry-afters as a floor
+        """
+        retry_policy = self._retry_policy.copy()
+        retry_policy["_retry_delay_min"] = 1
+
+        time_start = time.time()
+        with mocked_server_response(status=429, headers={"Retry-After": "3"}) as mock_obj:
+            with pytest.raises(RequestError) as cm:
+                with self.connection(extra_params=retry_policy) as conn:
+                    pass
+
+            duration = time.time() - time_start
+            assert isinstance(cm.value.args[1], MaxRetryDurationError)
+
+            # With setting delay_min to 1, the expected retry delays should be:
+            # 3, 3, 4
+            # The first 2 retries are allowed, the 3rd retry puts the total duration over the limit
+            # of 10 seconds
+            assert mock_obj.return_value.getresponse.call_count == 3
+            assert duration > 6
+
+            # Should be less than 7, but this is a safe margin for CI/CD slowness
+            assert duration < 10
+
     def test_retry_max_duration_not_exceeded(self):
         """GIVEN the max attempt duration of 10 seconds
         WHEN the server sends a Retry-After header of 60 seconds

tests/unit/test_retry.py

@@ -0,0 +1,55 @@
+from os import error
+import time
+from unittest.mock import Mock, patch
+import pytest
+from requests import Request
+from urllib3 import HTTPResponse
+from databricks.sql.auth.retry import DatabricksRetryPolicy, RequestHistory
+
+
+class TestRetry:
+
+    @pytest.fixture()
+    def retry_policy(self) -> DatabricksRetryPolicy:
+        return DatabricksRetryPolicy(
+            delay_min=1,
+            delay_max=30,
+            stop_after_attempts_count=3,
+            stop_after_attempts_duration=900,
+            delay_default=2,
+            force_dangerous_codes=[],
+        )
+
+    @pytest.fixture()
+    def error_history(self) -> RequestHistory:
+        return RequestHistory(
+            method="POST", url=None, error=None, status=503, redirect_location=None
+        )
+
+    @patch("time.sleep")
+    def test_sleep__no_retry_after(self, t_mock, retry_policy, error_history):
+        retry_policy._retry_start_time = time.time()
+        retry_policy.history = [error_history, error_history]
+        retry_policy.sleep(HTTPResponse(status=503))
+        t_mock.assert_called_with(2)
+
+    @patch("time.sleep")
+    def test_sleep__retry_after_is_binding(self, t_mock, retry_policy, error_history):
+        retry_policy._retry_start_time = time.time()
+        retry_policy.history = [error_history, error_history]
+        retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
+        t_mock.assert_called_with(3)
+
+    @patch("time.sleep")
+    def test_sleep__retry_after_present_but_not_binding(self, t_mock, retry_policy, error_history):
+        retry_policy._retry_start_time = time.time()
+        retry_policy.history = [error_history, error_history]
+        retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "1"}))
+        t_mock.assert_called_with(2)
+
+    @patch("time.sleep")
+    def test_sleep__retry_after_surpassed(self, t_mock, retry_policy, error_history):
+        retry_policy._retry_start_time = time.time()
+        retry_policy.history = [error_history, error_history, error_history]
+        retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
+        t_mock.assert_called_with(4)