51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
import abc
|
|
import logging
|
|
|
|
import requests
|
|
|
|
|
|
class _ErrorCustomizer(abc.ABC):
|
|
"""A customizer for errors from the Databricks REST API."""
|
|
|
|
@abc.abstractmethod
|
|
def customize_error(self, response: requests.Response, kwargs: dict):
|
|
"""Customize the error constructor parameters."""
|
|
|
|
|
|
class _RetryAfterCustomizer(_ErrorCustomizer):
|
|
"""An error customizer that sets the retry_after_secs parameter based on the Retry-After header."""
|
|
|
|
_DEFAULT_RETRY_AFTER_SECONDS = 1
|
|
"""The default number of seconds to wait before retrying a request if the Retry-After header is missing or is not
|
|
a valid integer."""
|
|
|
|
@classmethod
|
|
def _parse_retry_after(cls, response: requests.Response) -> int:
|
|
retry_after = response.headers.get("Retry-After")
|
|
if retry_after is None:
|
|
logging.debug(
|
|
f"No Retry-After header received in response with status code 429 or 503. Defaulting to {cls._DEFAULT_RETRY_AFTER_SECONDS}"
|
|
)
|
|
# 429 requests should include a `Retry-After` header, but if it's missing,
|
|
# we default to 1 second.
|
|
return cls._DEFAULT_RETRY_AFTER_SECONDS
|
|
# If the request is throttled, try parse the `Retry-After` header and sleep
|
|
# for the specified number of seconds. Note that this header can contain either
|
|
# an integer or a RFC1123 datetime string.
|
|
# See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
|
|
#
|
|
# For simplicity, we only try to parse it as an integer, as this is what Databricks
|
|
# platform returns. Otherwise, we fall back and don't sleep.
|
|
try:
|
|
return int(retry_after)
|
|
except ValueError:
|
|
logging.debug(
|
|
f"Invalid Retry-After header received: {retry_after}. Defaulting to {cls._DEFAULT_RETRY_AFTER_SECONDS}"
|
|
)
|
|
# defaulting to 1 sleep second to make self._is_retryable() simpler
|
|
return cls._DEFAULT_RETRY_AFTER_SECONDS
|
|
|
|
def customize_error(self, response: requests.Response, kwargs: dict):
|
|
if response.status_code in (429, 503):
|
|
kwargs["retry_after_secs"] = self._parse_retry_after(response)
|