|
| 1 | +import functools |
| 2 | +import logging |
| 3 | +import sys |
| 4 | +import time |
| 5 | + |
| 6 | +from requests.exceptions import RequestException, HTTPError |
| 7 | + |
| 8 | +logger = logging.getLogger(__name__) |
| 9 | +logging.basicConfig(stream=sys.stdout, level=logging.INFO) |
| 10 | + |
| 11 | + |
| 12 | +def retry_on_error( |
| 13 | + max_retries=3, |
| 14 | + initial_delay=1, |
| 15 | + backoff_factor=2, |
| 16 | + # retry on bad gateway, service unavailable, gateway timeout, too many requests |
| 17 | + retryable_status_codes=(502, 503, 504, 429), |
| 18 | + retryable_exceptions=(RequestException,), |
| 19 | +): |
| 20 | + """ |
| 21 | + A decorator for retrying functions that might fail due to transient network issues. |
| 22 | +
|
| 23 | + Args: |
| 24 | + max_retries: Maximum number of retry attempts |
| 25 | + initial_delay: Initial delay between retries in seconds |
| 26 | + backoff_factor: Factor by which the delay increases with each retry |
| 27 | + retryable_status_codes: HTTP status codes that trigger a retry |
| 28 | + retryable_exceptions: Exception types that trigger a retry |
| 29 | +
|
| 30 | + Returns: |
| 31 | + Decorated function with retry logic |
| 32 | + """ |
| 33 | + |
| 34 | + def decorator(func): |
| 35 | + @functools.wraps(func) |
| 36 | + def wrapper(*args, **kwargs): |
| 37 | + delay = initial_delay |
| 38 | + last_exception = None |
| 39 | + |
| 40 | + for retry_count in range(max_retries + 1): |
| 41 | + try: |
| 42 | + if retry_count > 0: |
| 43 | + logger.info( |
| 44 | + f"Retry attempt {retry_count}/{max_retries} for {func.__name__}" |
| 45 | + ) |
| 46 | + |
| 47 | + response = func(*args, **kwargs) |
| 48 | + |
| 49 | + # Check for retryable status codes in the response |
| 50 | + if ( |
| 51 | + hasattr(response, "status_code") |
| 52 | + and response.status_code in retryable_status_codes |
| 53 | + ): |
| 54 | + status_code = response.status_code |
| 55 | + logger.warning( |
| 56 | + f"Received status code {status_code} from {func.__name__}, retrying..." |
| 57 | + ) |
| 58 | + last_exception = HTTPError(f"HTTP Error {status_code}") |
| 59 | + else: |
| 60 | + # Success, return the response |
| 61 | + return response |
| 62 | + |
| 63 | + except retryable_exceptions as e: |
| 64 | + logger.warning(f"Request failed in {func.__name__}: {str(e)}") |
| 65 | + last_exception = e |
| 66 | + |
| 67 | + # Don't sleep if this was the last attempt |
| 68 | + if retry_count < max_retries: |
| 69 | + sleep_time = delay * (backoff_factor**retry_count) |
| 70 | + logger.info(f"Waiting {sleep_time:.2f} seconds before retry") |
| 71 | + time.sleep(sleep_time) |
| 72 | + |
| 73 | + # If we got here, all retries failed |
| 74 | + logger.error(f"All {max_retries} retries failed for {func.__name__}") |
| 75 | + if last_exception: |
| 76 | + raise last_exception |
| 77 | + return None |
| 78 | + |
| 79 | + return wrapper |
| 80 | + |
| 81 | + return decorator |
0 commit comments