Allow providing a custom retry policy. (#2946)

Author: jack-edmonds-dd

.generator/src/generator/templates/configuration.j2

@@ -142,6 +142,9 @@ class Configuration:
     :type retry_backoff_factor: float
     :param max_retries: The maximum number of times a single request can be retried.
     :type max_retries: int
+    :param retry_policy: Custom retry policy instance (e.g., urllib3.util.Retry). If provided, this overrides
+        the default retry behavior and the enable_retry, retry_backoff_factor, and max_retries settings.
+    :type retry_policy: urllib3.util.Retry
     :param delegated_auth_provider: The delegated authentication provider (e.g., 'aws' for AWS).
     :type delegated_auth_provider: str
     :param delegated_auth_org_uuid: The organization UUID for delegated authentication.
@@ -173,6 +176,7 @@ class Configuration:
         enable_retry=False,
         retry_backoff_factor=2,
         max_retries=3,
+        retry_policy=None,
         delegated_auth_provider=None,
         delegated_auth_org_uuid=None,
     ):
@@ -237,6 +241,7 @@ class Configuration:
         self.enable_retry = enable_retry
         self.retry_backoff_factor = retry_backoff_factor
         self.max_retries = max_retries
+        self.retry_policy = retry_policy
 
         # Keep track of unstable operations
         self.unstable_operations = _UnstableOperations({

.generator/src/generator/templates/rest.j2

@@ -67,7 +67,9 @@ class RESTClientObject:
         if configuration.assert_hostname is not None:
             addition_pool_args["assert_hostname"] = configuration.assert_hostname
 
-        if configuration.enable_retry:
+        if configuration.retry_policy is not None:
+            addition_pool_args["retries"] = configuration.retry_policy
+        elif configuration.enable_retry:
             retries = ClientRetry(
                 total = configuration.max_retries,
                 backoff_factor = configuration.retry_backoff_factor,

README.md

@@ -99,6 +99,29 @@ The default max retry is `3`, you can change it with `max_retries`
     configuration.max_retries = 5
 ```
 
+#### Custom retry policy
+
+You can provide a custom retry policy by passing a `urllib3.util.Retry` instance to the configuration.
+When a custom retry policy is provided, it takes precedence over the `enable_retry`, `retry_backoff_factor`,
+and `max_retries` settings.
+
+```python
+import urllib3
+from datadog_api_client import Configuration
+
+# Create a custom retry policy
+custom_retry = urllib3.util.Retry(
+    total=5,
+    backoff_factor=2,
+    status_forcelist=[429, 500, 502, 503, 504],
+    allowed_methods=["GET", "POST", "PUT", "DELETE"]
+)
+
+configuration = Configuration(retry_policy=custom_retry)
+```
+
+See [urllib3.util.Retry documentation](https://urllib3.readthedocs.io/en/stable/reference/urllib3.util.html#urllib3.util.Retry) for more details.
+
 ### Configure proxy
 
 You can configure the client to use proxy by setting the `proxy` key on configuration object:

src/datadog_api_client/configuration.py

@@ -143,6 +143,9 @@ class Configuration:
     :type retry_backoff_factor: float
     :param max_retries: The maximum number of times a single request can be retried.
     :type max_retries: int
+    :param retry_policy: Custom retry policy instance (e.g., urllib3.util.Retry). If provided, this overrides
+        the default retry behavior and the enable_retry, retry_backoff_factor, and max_retries settings.
+    :type retry_policy: urllib3.util.Retry
     :param delegated_auth_provider: The delegated authentication provider (e.g., 'aws' for AWS).
     :type delegated_auth_provider: str
     :param delegated_auth_org_uuid: The organization UUID for delegated authentication.
@@ -174,6 +177,7 @@ def __init__(
         enable_retry=False,
         retry_backoff_factor=2,
         max_retries=3,
+        retry_policy=None,
         delegated_auth_provider=None,
         delegated_auth_org_uuid=None,
     ):
@@ -238,6 +242,7 @@ def __init__(
         self.enable_retry = enable_retry
         self.retry_backoff_factor = retry_backoff_factor
         self.max_retries = max_retries
+        self.retry_policy = retry_policy
 
         # Keep track of unstable operations
         self.unstable_operations = _UnstableOperations(

src/datadog_api_client/rest.py

@@ -69,7 +69,9 @@ def __init__(self, configuration, pools_size=4, maxsize=4):
         if configuration.assert_hostname is not None:
             addition_pool_args["assert_hostname"] = configuration.assert_hostname
 
-        if configuration.enable_retry:
+        if configuration.retry_policy is not None:
+            addition_pool_args["retries"] = configuration.retry_policy
+        elif configuration.enable_retry:
             retries = ClientRetry(
                 total=configuration.max_retries,
                 backoff_factor=configuration.retry_backoff_factor,

tests/test_retry.py

@@ -42,3 +42,55 @@ def test_retry_backoff_factor_validation():
         configuration.retry_backoff_factor = 1
 
     configuration.retry_backoff_factor = 3
+
+
+@mock.patch("time.sleep", return_value=None)
+def test_custom_retry_policy(sleep_mock):
+    """Test that a custom retry policy is used when provided"""
+    import urllib3
+
+    # Create a custom retry policy with different settings
+    custom_retry = urllib3.util.Retry(
+        total=5,  # Different from default 3
+        backoff_factor=1,  # Different from default 2
+        status_forcelist=[500, 502, 503, 504],
+        allowed_methods=["GET", "POST"],
+    )
+
+    configuration = Configuration(retry_policy=custom_retry)
+
+    with vcr.use_cassette("tests/cassettes/test_retry/test_retry_errors.yaml", record_mode=vcr.mode.NONE):
+        with ApiClient(configuration) as api_client:
+            api_instance = logs_api.LogsApi(api_client)
+            logs = api_instance.list_logs_get()
+            assert len(logs.data) == 10
+            # With backoff_factor=1, sleep times are: 2, 4
+            assert sleep_mock.call_count == 2
+            assert sleep_mock.call_args_list[0][0][0] == 2
+            assert sleep_mock.call_args_list[1][0][0] == 4
+
+
+def test_custom_retry_policy_overrides_enable_retry():
+    """Test that retry_policy takes precedence over enable_retry"""
+    import urllib3
+
+    custom_retry = urllib3.util.Retry(total=10)
+    configuration = Configuration(
+        enable_retry=True,  # This should be ignored
+        max_retries=3,  # This should be ignored
+        retry_policy=custom_retry,
+    )
+
+    # Verify the configuration accepts the custom policy
+    assert configuration.retry_policy is custom_retry
+    assert configuration.retry_policy.total == 10
+
+
+def test_default_retry_when_no_custom_policy():
+    """Test that default retry behavior works when no custom policy is provided"""
+    configuration = Configuration(enable_retry=True, max_retries=5)
+
+    # Verify no custom policy is set
+    assert configuration.retry_policy is None
+    assert configuration.enable_retry is True
+    assert configuration.max_retries == 5