Add retry handling for 429 and transient errors

atomic_sdk/__init__.py

@@ -17,6 +17,7 @@
     ConflictError,
     InvalidRequestError,
     NotFoundError,
+    RateLimitError,
     ServerError,
 )
 
@@ -37,6 +38,7 @@
     "ConflictError",
     "InvalidRequestError",
     "NotFoundError",
+    "RateLimitError",
     "ServerError",
     "Job",
     "Backup",

atomic_sdk/api/base.py

@@ -1,24 +1,53 @@
-import requests
+import logging
+import random
+import time
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
 from typing import Iterator, Optional, Tuple, Union
 
-from ..exceptions import AtomicAPIError, ConflictError, InvalidRequestError, NotFoundError, ServerError
+import requests
+
+from ..exceptions import AtomicAPIError, ConflictError, InvalidRequestError, NotFoundError, RateLimitError, ServerError
+
+
+logger = logging.getLogger("atomic_sdk.retry")
+
+
+RETRYABLE_REQUEST_EXCEPTIONS = (
+    requests.exceptions.Timeout,
+    requests.exceptions.ChunkedEncodingError,
+)
 
 
 class ResourceClient:
     """A base client for a group of API resources."""
 
-    def __init__(self, session: requests.Session, base_url: str, client_id_or_name: str):
+    def __init__(
+        self,
+        session: requests.Session,
+        base_url: str,
+        client_id_or_name: str,
+        max_retries: int = 3,
+        backoff_base: float = 0.5,
+        timeout: Optional[Union[float, Tuple[float, float]]] = 30,
+    ):
         """
         Initializes the ResourceClient.
 
         Args:
             session: A requests.Session object configured with authentication.
             base_url: The base URL for the Atomic API.
             client_id_or_name: The client's identifier (name or ID).
+            max_retries: Number of retries for 429, 5xx, and connection errors.
+            backoff_base: Base delay in seconds for exponential backoff with jitter.
+            timeout: Default timeout passed to requests calls.
         """
         self._session = session
         self._base_url = base_url
         self._client_id_or_name = client_id_or_name
+        self._timeout = timeout
+        self._max_retries = max_retries
+        self._backoff_base = backoff_base
 
     def _get(self, endpoint: str, params: Optional[dict] = None) -> dict:
         """
@@ -61,15 +90,31 @@ def _get_raw(self, endpoint: str, params: Optional[dict] = None) -> bytes:
             The raw response content as bytes.
         """
         url = self._base_url.rstrip('/') + endpoint
+        attempt = 0
+        while True:
+            try:
+                response = self._session.get(url, params=params, stream=True, timeout=300) # Longer timeout for downloads
+                response.raise_for_status()
+                break
+            except requests.exceptions.HTTPError as e:
+                if self._retry_http_error(e, attempt):
+                    attempt += 1
+                    continue
+                self._raise_for_http_error(e)
+            except requests.exceptions.RequestException as e:
+                if self._retry_request_exception(e, url, attempt):
+                    attempt += 1
+                    continue
+                raise AtomicAPIError(f"Request failed for {url}: {e}") from e
+
         try:
-            response = self._session.get(url, params=params, timeout=300) # Longer timeout for downloads
-            response.raise_for_status()
             return response.content
-        except requests.exceptions.HTTPError as e:
-            # Re-raise with a more specific custom exception if needed
-            raise AtomicAPIError(f"HTTP Error for {url}: {e.response.status_code} {e.response.text}") from e
         except requests.exceptions.RequestException as e:
             raise AtomicAPIError(f"Request failed for {url}: {e}") from e
+        finally:
+            close = getattr(response, "close", None)
+            if close:
+                close()
 
     def _get_stream(
         self,
@@ -101,6 +146,7 @@ def _get_stream(
             ServerError: For 5xx responses.
         """
         url = self._base_url.rstrip('/') + endpoint
+        timeout = self._timeout if timeout is None else timeout
         try:
             with self._session.get(url, params=params, stream=True, timeout=timeout) as response:
                 response.raise_for_status()
@@ -125,6 +171,12 @@ def _raise_for_http_error(self, error: requests.exceptions.HTTPError) -> None:
 
         if status_code == 404:
             raise NotFoundError(message, status_code) from error
+        if status_code == 429:
+            raise RateLimitError(
+                message,
+                status_code,
+                retry_after=self._parse_retry_after(error.response),
+            ) from error
         if status_code == 409:
             raise ConflictError(message, status_code) from error
         if 400 <= status_code < 500:
@@ -134,6 +186,86 @@ def _raise_for_http_error(self, error: requests.exceptions.HTTPError) -> None:
 
         raise AtomicAPIError(message, status_code) from error
 
+    @staticmethod
+    def _parse_retry_after(response: requests.Response) -> Optional[int]:
+        """Parse a Retry-After header as seconds, or as an HTTP-date."""
+        value = response.headers.get("Retry-After")
+        if not value:
+            return None
+
+        try:
+            return max(0, int(float(value)))
+        except ValueError:
+            pass
+
+        try:
+            retry_at = parsedate_to_datetime(value)
+        except (TypeError, ValueError):
+            return None
+
+        if retry_at.tzinfo is None:
+            retry_at = retry_at.replace(tzinfo=timezone.utc)
+        return max(0, int((retry_at - datetime.now(timezone.utc)).total_seconds()))
+
+    def _retry_http_error(self, error: requests.exceptions.HTTPError, attempt: int) -> bool:
+        """Sleep and return True when an HTTP error should be retried."""
+        status_code = error.response.status_code
+        if status_code != 429 and not 500 <= status_code < 600:
+            return False
+        if attempt >= self._max_retries:
+            return False
+
+        retry_after = self._parse_retry_after(error.response) if status_code == 429 else None
+        delay = retry_after if retry_after is not None else self._backoff_delay(attempt)
+
+        if status_code == 429 and retry_after is not None:
+            logger.warning(
+                "429 Retry-After=%ss, retrying (attempt %s/%s)",
+                retry_after,
+                attempt + 1,
+                self._max_retries,
+            )
+        else:
+            logger.warning(
+                "%s %s, backing off %.2fs (attempt %s/%s)",
+                status_code,
+                error.response.reason,
+                delay,
+                attempt + 1,
+                self._max_retries,
+            )
+        time.sleep(delay)
+        return True
+
+    def _retry_request_exception(self, error: requests.exceptions.RequestException, url: str, attempt: int) -> bool:
+        """Sleep and return True when a connection-level request error should be retried."""
+        if not self._is_retryable_request_exception(error):
+            return False
+        if attempt >= self._max_retries:
+            return False
+
+        delay = self._backoff_delay(attempt)
+        logger.warning(
+            "Request failed for %s: %s, backing off %.2fs (attempt %s/%s)",
+            url,
+            error,
+            delay,
+            attempt + 1,
+            self._max_retries,
+        )
+        time.sleep(delay)
+        return True
+
+    @staticmethod
+    def _is_retryable_request_exception(error: requests.exceptions.RequestException) -> bool:
+        return (
+            type(error) is requests.exceptions.ConnectionError
+            or isinstance(error, RETRYABLE_REQUEST_EXCEPTIONS)
+        )
+
+    def _backoff_delay(self, attempt: int) -> float:
+        return random.uniform(0, self._backoff_base * (2 ** attempt))
+
     def _request(self, method: str, endpoint: str, **kwargs) -> dict:
         """
         Makes an HTTP request to the specified endpoint and handles JSON response.
@@ -151,16 +283,25 @@ def _request(self, method: str, endpoint: str, **kwargs) -> dict:
             InvalidRequestError: For 4xx client errors with a message.
         """
         url = self._base_url.rstrip('/') + endpoint
-        try:
-            response = self._session.request(method, url, **kwargs)
-            response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
-            return response.json()
-
-        except requests.exceptions.HTTPError as e:
-            self._raise_for_http_error(e)
-
-        except requests.exceptions.RequestException as e:
-            raise AtomicAPIError(f"Request failed for {url}: {e}") from e
+        kwargs.setdefault("timeout", self._timeout)
+        attempt = 0
+        while True:
+            try:
+                response = self._session.request(method, url, **kwargs)
+                response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
+                return response.json()
+
+            except requests.exceptions.HTTPError as e:
+                if self._retry_http_error(e, attempt):
+                    attempt += 1
+                    continue
+                self._raise_for_http_error(e)
+
+            except requests.exceptions.RequestException as e:
+                if self._retry_request_exception(e, url, attempt):
+                    attempt += 1
+                    continue
+                raise AtomicAPIError(f"Request failed for {url}: {e}") from e
 
     def _get_service_and_identifier(self, site_id: Optional[int], domain: Optional[str]) -> Tuple[str, Union[int, str]]:
         """

atomic_sdk/client.py

@@ -29,23 +29,38 @@ class AtomicClient:
 
     BASE_URL = "https://atomic-api.wordpress.com/api/v1.0/"
 
-    def __init__(self, api_key: str, client_id_or_name: str, timeout: int = 30):
+    def __init__(
+        self,
+        api_key: str,
+        client_id_or_name: str,
+        timeout: int = 30,
+        max_retries: int = 3,
+        backoff_base: float = 0.5,
+    ):
         """
         Initializes the Atomic API client.
 
         Args:
             api_key: Your platform or developer API key for authentication.
             client_id_or_name: Your unique client identifier (e.g., 'your-client-name').
             timeout: The timeout in seconds for API requests. Defaults to 30.
+            max_retries: Number of retries for 429, 5xx, and connection errors. Defaults to 3.
+            backoff_base: Base delay in seconds for exponential backoff with jitter.
         """
         if not api_key:
             raise ValueError("An API key is required.")
         if not client_id_or_name:
             raise ValueError("A client identifier (name or ID) is required.")
+        if max_retries < 0:
+            raise ValueError("max_retries must be greater than or equal to 0.")
+        if backoff_base < 0:
+            raise ValueError("backoff_base must be greater than or equal to 0.")
 
         self.api_key = api_key
         self.client_id_or_name = client_id_or_name
         self.timeout = timeout
+        self.max_retries = max_retries
+        self.backoff_base = backoff_base
 
         # Get the package version at runtime to avoid circular imports
         try:
@@ -60,24 +75,32 @@ def __init__(self, api_key: str, client_id_or_name: str, timeout: int = 30):
             "User-Agent": f"Python AtomicSDK/{sdk_version}",
             "Accept": "application/json",
         })
-        self._session.timeout = self.timeout
+
+        resource_args = (
+            self._session,
+            self.BASE_URL,
+            self.client_id_or_name,
+            self.max_retries,
+            self.backoff_base,
+            self.timeout,
+        )
 
         # Instantiate and attach all the resource-specific clients
-        self.backups = BackupsClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.client = ClientClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.cron = CronClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.custom_certificates = CustomCertificatesClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.edge_cache = EdgeCacheClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.email = EmailClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.metrics = MetricsClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.security = SecurityClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.servers = ServersClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.sites = SitesClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.ssh = SSHClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.tasks = TasksClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.utility = UtilityClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.migrations = MigrationsClient(self._session, self.BASE_URL, self.client_id_or_name)
-        self.response_tickets = ResponseTicketsClient(self._session, self.BASE_URL, self.client_id_or_name)
+        self.backups = BackupsClient(*resource_args)
+        self.client = ClientClient(*resource_args)
+        self.cron = CronClient(*resource_args)
+        self.custom_certificates = CustomCertificatesClient(*resource_args)
+        self.edge_cache = EdgeCacheClient(*resource_args)
+        self.email = EmailClient(*resource_args)
+        self.metrics = MetricsClient(*resource_args)
+        self.security = SecurityClient(*resource_args)
+        self.servers = ServersClient(*resource_args)
+        self.sites = SitesClient(*resource_args)
+        self.ssh = SSHClient(*resource_args)
+        self.tasks = TasksClient(*resource_args)
+        self.utility = UtilityClient(*resource_args)
+        self.migrations = MigrationsClient(*resource_args)
+        self.response_tickets = ResponseTicketsClient(*resource_args)
 
         # Pass a reference of the main client to resource clients that return Job objects,
         # so Job.status() can call self._client.sites.get_job_status().

atomic_sdk/exceptions.py

@@ -23,6 +23,13 @@ def __init__(self, message="Invalid request.", status_code=None):
         super().__init__(message, status_code)
 
 
+class RateLimitError(AtomicAPIError):
+    """Raised when the API returns HTTP 429 (rate limited)."""
+    def __init__(self, message="Rate limit exceeded.", status_code=429, retry_after=None):
+        super().__init__(message, status_code)
+        self.retry_after = retry_after
+
+
 class ConflictError(InvalidRequestError):
     """Raised when a request conflicts with existing resource state (HTTP 409)."""
     def __init__(self, message="The request conflicts with existing state.", status_code=409):