import functools
import re
import warnings
import time
from typing import Optional
from googleapiclient.errors import HttpError
MAX_RETRY_ATTRIBUTE = "max_retry"
SLEEP_ATTRIBUTE = "sleep"
[docs]def retry_on_type_and_msg(exception: Exception, msg_pattern: Optional[str] = None):
def wrapper(method):
@functools.wraps(wrapped=method)
def wrapped_function(self, *args, **kwargs):
max_retry = getattr(self, MAX_RETRY_ATTRIBUTE, 0)
max_retry = max(max_retry, 0)
sleep = getattr(self, SLEEP_ATTRIBUTE, 5)
if sleep < 0:
raise AttributeError(f"{SLEEP_ATTRIBUTE} must be positive. Got {sleep}")
pattern = None
if msg_pattern is not None:
pattern = re.compile(msg_pattern)
while True:
max_retry -= 1
try:
result = method(self, *args, **kwargs)
return result
except exception as e:
if max_retry < 0 or pattern is not None and not pattern.match(str(e)):
# max retry reached or pattern does not match.
raise e
warnings.warn(f"handled exception {e}. remaining retry: {max_retry}", UserWarning)
time.sleep(sleep)
sleep = sleep*2
continue
return wrapped_function
return wrapper
[docs]def retry_on_out_of_quota():
"""A decorator to give wrapped function ability to retry on quota exceeded related HttpError raise by Google API.
It only works on class method and requires "max_retry" and "sleep" attribute in the class. If `max_retry` is
non-positive, no retry will be attempt. `sleep` is the base number of seconds between consecutive retries. The number
of wait seconds will double after each sleep.
:return:
"""
return retry_on_type_and_msg(HttpError, msg_pattern=".*(User Rate Limit Exceeded|Quota exceeded)+.*")