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..eea69f16 --- /dev/null +++ b/thousandeyes-sdk-client/src/thousandeyes_sdk/client/thousandeyes_retry.py @@ -0,0 +1,64 @@ +import re +import time +from typing import Collection, 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" + } + + RESET_HEADER_PATTERN = re.compile(r"^\s*[0-9]+\s*$") + HTTP_TOO_MANY_REQUESTS = 429 + + def __init__( + self, + total: Union[bool, int, None] = 3, + connect: Optional[int] = None, + read: Optional[int] = None, + redirect: Union[bool, int, None] = None, + status: Optional[int] = 1, + other: Optional[int] = None, + allowed_methods: Optional[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: Collection[str] = Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT, + backoff_jitter: float = 0.0) -> None: + super().__init__(total, connect, read, redirect, status, other, allowed_methods, + status_forcelist, backoff_factor, backoff_max, raise_on_redirect, + raise_on_status, history, respect_retry_after_header, + remove_headers_on_redirect, backoff_jitter) + + def is_retry(self, method: str, status_code: int, has_retry_after: bool = False) -> bool: + # Always retry on 429, regardless of method or status_forcelist + return (status_code == self.HTTP_TOO_MANY_REQUESTS or + super().is_retry(method, status_code, has_retry_after)) + + 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 + + def _parse_reset_header(self, value: Optional[str]) -> Optional[float]: + if value is None or not self.RESET_HEADER_PATTERN.match(value): + return None + + seconds: float = int(value) - time.time() + return max(seconds, 0)