CP-2126 Add handling for ThousandEyes Rate Limit headers

thousandeyes-sdk-client/src/thousandeyes_sdk/client/configuration.py

@@ -8,6 +8,8 @@ from typing import Optional
 
 import urllib3
 
+from thousandeyes_sdk.client.thousandeyes_retry import ThousandEyesRetry
+
 JSON_SCHEMA_VALIDATION_KEYWORDS = {
     'multipleOf', 'maximum', 'exclusiveMaximum',
     'minimum', 'exclusiveMinimum', 'maxLength',
@@ -54,6 +56,7 @@ class Configuration:
                  server_index=None, server_variables=None,
                  server_operation_index=None, server_operation_variables=None,
                  ssl_ca_cert=None,
+                 retries=None
                  ) -> None:
         """Constructor
         """
@@ -154,7 +157,9 @@ class Configuration:
         self.safe_chars_for_path_param = ''
         """Safe chars for path_param
         """
-        self.retries = None
+        self.retries = ThousandEyesRetry()
+        if retries:
+            self.retries = retries
         """Adding retries to override urllib3 default value 3
         """
         # Enable client side validation

thousandeyes-sdk-client/src/thousandeyes_sdk/client/exceptions.py

@@ -139,6 +139,9 @@ class ApiException(OpenApiException):
         if http_resp.status == 404:
             raise NotFoundException(http_resp=http_resp, body=body, data=data)
 
+        if http_resp.status == 429:
+            raise TooManyRequestsException(http_resp=http_resp, body=body, data=data)
+
         if 500 <= http_resp.status <= 599:
             raise ServiceException(http_resp=http_resp, body=body, data=data)
         raise ApiException(http_resp=http_resp, body=body, data=data)
@@ -173,6 +176,10 @@ class ForbiddenException(ApiException):
     pass
 
 
+class TooManyRequestsException(ApiException):
+    pass
+
+
 class ServiceException(ApiException):
     pass
 

thousandeyes-sdk-client/src/thousandeyes_sdk/client/thousandeyes_retry.py

@@ -0,0 +1,61 @@
+import re
+import time
+import typing
+from typing import Optional, Union
+
+from urllib3 import BaseHTTPResponse
+from urllib3.util.retry import RequestHistory, Retry
+
+
+class ThousandEyesRetry(Retry):
+    RATE_LIMIT_RESET_HEADERS = {
+        "x-organization-rate-limit-reset",
+        "x-instant-test-rate-limit-reset"
+    }
+
+    def __init__(
+            self,
+            total: Union[bool, int, None] = 10,
+            connect: Optional[int] = None,
+            read: Optional[int] = None,
+            redirect: Union[bool, int, None] = None,
+            status: Optional[int] = None,
+            other: Optional[int] = None,
+            allowed_methods: Optional[typing.Collection[str]] = Retry.DEFAULT_ALLOWED_METHODS,
+            status_forcelist=None,
+            backoff_factor: float = 0,
+            backoff_max: float = Retry.DEFAULT_BACKOFF_MAX,
+            raise_on_redirect: bool = False,
+            raise_on_status: bool = False,
+            history: Optional[tuple[RequestHistory, ...]] = None,
+            respect_retry_after_header: bool = True,
+            remove_headers_on_redirect: typing.Collection[
+                str
+            ] = Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT,
+            backoff_jitter: float = 0.0) -> None:
+        super().__init__(total, connect, read, redirect, status, other, allowed_methods,
+                         [429] if status_forcelist is None else status_forcelist,
+                         backoff_factor, backoff_max, raise_on_redirect,
+                         raise_on_status, history, respect_retry_after_header,
+                         remove_headers_on_redirect, backoff_jitter)
+
+    def get_retry_after(self, response: BaseHTTPResponse) -> Optional[float]:
+        retry_after: Optional[float] = super().get_retry_after(response)
+
+        if retry_after:
+            return retry_after
+
+        for header in self.RATE_LIMIT_RESET_HEADERS:
+            value = self._parse_reset_header(response.headers.get(header))
+            if value and (retry_after is None or value > retry_after):
+                retry_after = value
+
+        return retry_after
+
+    @staticmethod
+    def _parse_reset_header(value: Optional[str]) -> Optional[float]:
+        if value is None or not re.match(r"^\s*[0-9]+\s*$", value):
+            return None
+
+        seconds: float = int(value) - time.time()
+        return max(seconds, 0)