Add retry (#1384)

* add retry and its configs

* auto reformat

* format

* redo is_retry and rewrite retry config options

* Reformat

* Fix syntax & add test

* Fixe test & adjusting retry impl

* pre-commit fixes

* Update template

* Fix import

* pre-commit fixes

* change retry import

* remove retry import

* remove saving content of the response

* update cassette

* delete env file

* update cassette

* use mock for test_retry

* pre-commit fixes

* use urllib3 for requests

* pre-commit fixes

* fix return format

* pre-commit fixes

* add response header

* pre-commit fixes

* change retry test

* pre-commit fixes

* update test_retry

* pre-commit fixes

* use call count

* pre-commit fixes

* Update .generator/src/generator/templates/rest.j2

Co-authored-by: Thomas Hervé <[email protected]>

* pre-commit fixes

* Add back test

* pre-commit fixes

* Fix content encoding

* touch up rest client and test

* pre-commit fixes

* add test without reset header

* pre-commit fixes

* change allowed method check

* Minor tweaks

* Use cassette for errors as well

* Docs

---------

Co-authored-by: ci.datadog-api-spec <[email protected]>
Co-authored-by: Thomas Hervé <[email protected]>

Author: 3 people

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

@@ -158,6 +158,9 @@ class Configuration:
         check_input_type=True,
         check_return_type=True,
         spec_property_naming=False,
+        enable_retry=False,
+        retry_backoff_factor=2,
+        max_retries=3,
     ):
         """Constructor."""
         self._base_path = "https://api.datadoghq.com" if host is None else host
@@ -200,7 +203,6 @@ class Configuration:
         self.proxy = None
         self.proxy_headers = None
         self.safe_chars_for_path_param = ""
-        self.retries = None
         # Enable client side validation
         self.client_side_validation = True
 
@@ -217,6 +219,11 @@ class Configuration:
         self.check_return_type = check_return_type
         self.spec_property_naming = spec_property_naming
 
+        # Options for http retry
+        self.enable_retry = enable_retry
+        self.retry_backoff_factor = retry_backoff_factor
+        self.max_retries = max_retries
+
         # Keep track of unstable operations
         self.unstable_operations = _UnstableOperations({
         {%- for version, api in apis.items() %}

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

@@ -6,7 +6,6 @@ import re
 import ssl
 from urllib.parse import urlencode
 import zlib
-
 import urllib3  # type: ignore
 
 from {{ package }}.exceptions import (
@@ -22,6 +21,30 @@ from {{ package }}.exceptions import (
 logger = logging.getLogger(__name__)
 
 
+class ClientRetry(urllib3.util.Retry):
+    RETRY_AFTER_STATUS_CODES = frozenset([413, 429, 500, 501, 502, 503, 504, 505, 506, 507, 509, 510, 511])
+    DEFAULT_ALLOWED_METHODS = frozenset(["GET", "PUT", "DELETE", "POST", "PATCH"])
+
+    def get_retry_after(self, response):
+        """
+        This method overrides the default "Retry-after" header and uses dd's X-Ratelimit-Reset header
+        and gets the value of X-Ratelimit-Reset in seconds.
+        """
+        retry_after = response.headers.get("X-Ratelimit-Reset")
+
+        if retry_after is None:
+            return None
+        return self.parse_retry_after(retry_after)
+
+    def is_retry(self, method, status_code, has_retry_after=False):
+        if method not in self.DEFAULT_ALLOWED_METHODS:
+            return False
+
+        if self.status_forcelist and status_code in self.status_forcelist:
+            return True
+        return self.total and self.respect_retry_after_header and (status_code in self.RETRY_AFTER_STATUS_CODES)
+
+
 class RESTClientObject:
     def __init__(self, configuration, pools_size=4, maxsize=4):
         # urllib3.PoolManager will pass all kw parameters to connectionpool
@@ -40,8 +63,12 @@ class RESTClientObject:
         if configuration.assert_hostname is not None:
             addition_pool_args["assert_hostname"] = configuration.assert_hostname
 
-        if configuration.retries is not None:
-            addition_pool_args["retries"] = configuration.retries
+        if configuration.enable_retry:
+            retries = ClientRetry(
+                total=configuration.max_retries,
+                backoff_factor=configuration.retry_backoff_factor,
+            )
+            addition_pool_args["retries"] = retries
 
         if configuration.socket_options is not None:
             addition_pool_args["socket_options"] = configuration.socket_options

README.md

@@ -85,6 +85,20 @@ If you want to enable requests logging, set the `debug` flag on your configurati
 configuration.debug = True
 ```
 
+### Enable retry
+
+If you want to enable retry when getting status code `429` rate-limited, set `enable_retry` to `True`
+
+```python
+    configuration.enable_retry = True
+```
+
+The default max retry is `3`, you can change it with `max_retries`
+
+```python
+    configuration.max_retries = 5
+```
+
 ### Configure proxy
 
 You can configure the client to use proxy by setting the `proxy` key on configuration object:

src/datadog_api_client/configuration.py

@@ -159,6 +159,9 @@ def __init__(
         check_input_type=True,
         check_return_type=True,
         spec_property_naming=False,
+        enable_retry=False,
+        retry_backoff_factor=2,
+        max_retries=3,
     ):
         """Constructor."""
         self._base_path = "https://api.datadoghq.com" if host is None else host
@@ -201,7 +204,6 @@ def __init__(
         self.proxy = None
         self.proxy_headers = None
         self.safe_chars_for_path_param = ""
-        self.retries = None
         # Enable client side validation
         self.client_side_validation = True
 
@@ -218,6 +220,11 @@ def __init__(
         self.check_return_type = check_return_type
         self.spec_property_naming = spec_property_naming
 
+        # Options for http retry
+        self.enable_retry = enable_retry
+        self.retry_backoff_factor = retry_backoff_factor
+        self.max_retries = max_retries
+
         # Keep track of unstable operations
         self.unstable_operations = _UnstableOperations(
             {

src/datadog_api_client/rest.py

@@ -8,7 +8,6 @@
 import ssl
 from urllib.parse import urlencode
 import zlib
-
 import urllib3  # type: ignore
 
 from datadog_api_client.exceptions import (
@@ -24,6 +23,30 @@
 logger = logging.getLogger(__name__)
 
 
+class ClientRetry(urllib3.util.Retry):
+    RETRY_AFTER_STATUS_CODES = frozenset([413, 429, 500, 501, 502, 503, 504, 505, 506, 507, 509, 510, 511])
+    DEFAULT_ALLOWED_METHODS = frozenset(["GET", "PUT", "DELETE", "POST", "PATCH"])
+
+    def get_retry_after(self, response):
+        """
+        This method overrides the default "Retry-after" header and uses dd's X-Ratelimit-Reset header
+        and gets the value of X-Ratelimit-Reset in seconds.
+        """
+        retry_after = response.headers.get("X-Ratelimit-Reset")
+
+        if retry_after is None:
+            return None
+        return self.parse_retry_after(retry_after)
+
+    def is_retry(self, method, status_code, has_retry_after=False):
+        if method not in self.DEFAULT_ALLOWED_METHODS:
+            return False
+
+        if self.status_forcelist and status_code in self.status_forcelist:
+            return True
+        return self.total and self.respect_retry_after_header and (status_code in self.RETRY_AFTER_STATUS_CODES)
+
+
 class RESTClientObject:
     def __init__(self, configuration, pools_size=4, maxsize=4):
         # urllib3.PoolManager will pass all kw parameters to connectionpool
@@ -42,8 +65,12 @@ 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.retries is not None:
-            addition_pool_args["retries"] = configuration.retries
+        if configuration.enable_retry:
+            retries = ClientRetry(
+                total=configuration.max_retries,
+                backoff_factor=configuration.retry_backoff_factor,
+            )
+            addition_pool_args["retries"] = retries
 
         if configuration.socket_options is not None:
             addition_pool_args["socket_options"] = configuration.socket_options