diff --git a/caldav/davclient.py b/caldav/davclient.py index f13ecb86..dfd29e0b 100644 --- a/caldav/davclient.py +++ b/caldav/davclient.py @@ -10,7 +10,10 @@ import logging import sys +import time import warnings +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime from types import TracebackType from typing import TYPE_CHECKING, Any, Optional from urllib.parse import unquote @@ -206,6 +209,9 @@ def __init__( features: FeatureSet | dict | str = None, enable_rfc6764: bool = True, require_tls: bool = True, + rate_limit_handle: bool = False, + rate_limit_default_sleep: int = None, + rate_limit_max_sleep: int = None, ) -> None: """ Sets up a HTTPConnection object towards the server in the url. @@ -243,6 +249,13 @@ def __init__( redirect to unencrypted HTTP. Set to False ONLY if you need to support non-TLS servers and trust your DNS infrastructure. This parameter has no effect if enable_rfc6764=False. + rate_limit_handle: boolean, a parameter that determines whether the rate limit response + should be handled. Default: False. + rate_limit_default_sleep: integer, the default number of seconds to sleep if the server + response cannot be parsed, or if no retry-after is specified + and the HTTP response status code is 429. Default: None. + rate_limit_max_sleep: integer, the maximum number of seconds the script will sleep + when encountering a rate limit. Default: None. The niquests library will honor a .netrc-file, if such a file exists username and password may be omitted. @@ -341,6 +354,10 @@ def __init__( self._principal = None + self.rate_limit_handle = rate_limit_handle + self.rate_limit_default_sleep = rate_limit_default_sleep + self.rate_limit_max_sleep = rate_limit_max_sleep + def __enter__(self) -> Self: ## Used for tests, to set up a temporarily test server if hasattr(self, "setup"): @@ -931,7 +948,21 @@ def request( Returns: DAVResponse """ - return self._sync_request(url, method, body, headers) + try: + return self._sync_request(url, method, body, headers) + except error.RateLimitError as e: + if self.rate_limit_handle: + retry_after_seconds = self.rate_limit_default_sleep + if e.retry_after_seconds is not None: + retry_after_seconds = e.retry_after_seconds + if self.rate_limit_max_sleep: + retry_after_seconds = min(retry_after_seconds or 0, self.rate_limit_max_sleep) + if retry_after_seconds <= 0: + raise e + time.sleep(retry_after_seconds) + return self._sync_request(url, method, body, headers) + + raise e def _sync_request( self, @@ -974,8 +1005,32 @@ def _sync_request( cert=self.ssl_cert, ) - # Handle 401 responses for auth negotiation r_headers = CaseInsensitiveDict(r.headers) + + # Handle 429, 503 responses for retry negotiation + if r.status_code in (429, 503): + retry_after_header: Optional[str] = r_headers.get("Retry-After") + retry_after_value: Optional[str] = None + retry_seconds: Optional[float] = None + if retry_after_header: + retry_after_value = retry_after_header + try: + retry_seconds = int(retry_after_header) + except ValueError: + try: + retry_date = parsedate_to_datetime(retry_after) + now = datetime.now(timezone.utc) + retry_seconds = max(0, (retry_date - now).total_seconds()) + except: + pass + if r.status_code == 429 or retry_after_header is not None: + raise error.RateLimitError( + f"Rate limited or service unavailable. Retry after: {retry_after}", + retry_after=retry_after, + retry_after_seconds=retry_seconds, + ) + + # Handle 401 responses for auth negotiation if ( r.status_code == 401 and "WWW-Authenticate" in r_headers diff --git a/caldav/lib/error.py b/caldav/lib/error.py index 3eafb364..5391c62b 100644 --- a/caldav/lib/error.py +++ b/caldav/lib/error.py @@ -133,6 +133,15 @@ class ResponseError(DAVError): pass +class RateLimitError(DAVError): + """Returns in case 429 Too Many Requests or 503 Service Unavailable w Retry-After.""" + + def __init__(self, msg, retry_after=None, retry_after_seconds=None): + super().__init__(msg) + self.retry_after = retry_after + self.retry_after_seconds = retry_after_seconds + + exception_by_method: dict[str, DAVError] = defaultdict(lambda: DAVError) for method in ( "delete",