From 721c362d0dccd9ccd6d959981c73764689b887bf Mon Sep 17 00:00:00 2001 From: jack-edmonds-dd Date: Tue, 11 Nov 2025 09:43:42 -0500 Subject: [PATCH] Allow providing a custom retry policy. --- .../src/generator/templates/configuration.j2 | 5 ++ .generator/src/generator/templates/rest.j2 | 4 +- README.md | 23 ++++++++ src/datadog_api_client/configuration.py | 5 ++ src/datadog_api_client/rest.py | 4 +- tests/test_retry.py | 52 +++++++++++++++++++ 6 files changed, 91 insertions(+), 2 deletions(-) diff --git a/.generator/src/generator/templates/configuration.j2 b/.generator/src/generator/templates/configuration.j2 index e6cf20f713..b0f3730142 100644 --- a/.generator/src/generator/templates/configuration.j2 +++ b/.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({ diff --git a/.generator/src/generator/templates/rest.j2 b/.generator/src/generator/templates/rest.j2 index 4b646a20cd..55ec9466c5 100644 --- a/.generator/src/generator/templates/rest.j2 +++ b/.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, diff --git a/README.md b/README.md index 7594ad348e..4af56b6047 100644 --- a/README.md +++ b/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: diff --git a/src/datadog_api_client/configuration.py b/src/datadog_api_client/configuration.py index 4433a08fb0..e02030186f 100644 --- a/src/datadog_api_client/configuration.py +++ b/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( diff --git a/src/datadog_api_client/rest.py b/src/datadog_api_client/rest.py index d592a9ab19..481be7e3d9 100644 --- a/src/datadog_api_client/rest.py +++ b/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, diff --git a/tests/test_retry.py b/tests/test_retry.py index d5ea7ebf16..e403dc0204 100644 --- a/tests/test_retry.py +++ b/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