diff --git a/thousandeyes-sdk-client/src/thousandeyes_sdk/client/configuration.py b/thousandeyes-sdk-client/src/thousandeyes_sdk/client/configuration.py index 579dea2f..28fc2ef6 100644 --- a/thousandeyes-sdk-client/src/thousandeyes_sdk/client/configuration.py +++ b/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 diff --git a/thousandeyes-sdk-client/src/thousandeyes_sdk/client/exceptions.py b/thousandeyes-sdk-client/src/thousandeyes_sdk/client/exceptions.py index 0452f86e..43050fb1 100644 --- a/thousandeyes-sdk-client/src/thousandeyes_sdk/client/exceptions.py +++ b/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 diff --git a/thousandeyes-sdk-client/src/thousandeyes_sdk/client/thousandeyes_retry.py b/thousandeyes-sdk-client/src/thousandeyes_sdk/client/thousandeyes_retry.py new file mode 100644 index 00000000..53d6098c --- /dev/null +++ b/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)