Refactor utils

Author: vdusek

src/apify_client/_http_client.py

@@ -1,32 +1,36 @@
 from __future__ import annotations
 
+import asyncio
 import gzip
 import json as jsonlib
 import logging
 import os
+import random
 import sys
+import time
 from datetime import datetime, timezone
 from http import HTTPStatus
 from importlib import metadata
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 from urllib.parse import urlencode
 
 import impit
 
 from apify_client._logging import log_context, logger_name
 from apify_client._statistics import ClientStatistics
-from apify_client._utils import is_retryable_error, retry_with_exp_backoff, retry_with_exp_backoff_async
-from apify_client.errors import ApifyApiError
+from apify_client.errors import ApifyApiError, InvalidResponseBodyError
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
     from apify_client._config import ClientConfig
     from apify_client._consts import JsonSerializable
 
 DEFAULT_BACKOFF_EXPONENTIAL_FACTOR = 2
 DEFAULT_BACKOFF_RANDOM_FACTOR = 1
 
+T = TypeVar('T')
+
 logger = logging.getLogger(logger_name)
 
 
@@ -99,6 +103,26 @@ def _parse_params(params: dict | None) -> dict | None:
 
         return parsed_params
 
+    @staticmethod
+    def _is_retryable_error(exc: Exception) -> bool:
+        """Check if an exception should be retried.
+
+        Args:
+            exc: The exception to check.
+
+        Returns:
+            True if the exception is retryable (network errors, timeouts, etc.).
+        """
+        return isinstance(
+            exc,
+            (
+                InvalidResponseBodyError,
+                impit.NetworkError,
+                impit.TimeoutException,
+                impit.RemoteProtocolError,
+            ),
+        )
+
     def _prepare_request_call(
         self,
         headers: dict | None = None,
@@ -201,7 +225,7 @@ def _make_request(stop_retrying: Callable, attempt: int) -> impit.Response:
 
             except Exception as exc:
                 logger.debug('Request threw exception', exc_info=exc)
-                if not is_retryable_error(exc):
+                if not self._is_retryable_error(exc):
                     logger.debug('Exception is not retryable', exc_info=exc)
                     stop_retrying()
                 raise
@@ -217,14 +241,59 @@ def _make_request(stop_retrying: Callable, attempt: int) -> impit.Response:
             response.read()
             raise ApifyApiError(response, attempt, method=method)
 
-        return retry_with_exp_backoff(
+        return self._retry_with_exp_backoff(
             _make_request,
             max_retries=self._config.max_retries,
             backoff_base_millis=self._config.min_delay_between_retries_millis,
             backoff_factor=DEFAULT_BACKOFF_EXPONENTIAL_FACTOR,
             random_factor=DEFAULT_BACKOFF_RANDOM_FACTOR,
         )
 
+    @staticmethod
+    def _retry_with_exp_backoff(
+        func: Callable[[Callable[[], None], int], T],
+        *,
+        max_retries: int = 8,
+        backoff_base_millis: int = 500,
+        backoff_factor: float = 2,
+        random_factor: float = 1,
+    ) -> T:
+        """Retry a function with exponential backoff.
+
+        Args:
+            func: Function to retry. Receives a stop_retrying callback and attempt number.
+            max_retries: Maximum number of retry attempts.
+            backoff_base_millis: Base backoff delay in milliseconds.
+            backoff_factor: Exponential backoff multiplier (1-10).
+            random_factor: Random jitter factor (0-1).
+
+        Returns:
+            The return value of the function.
+        """
+        random_factor = min(max(0, random_factor), 1)
+        backoff_factor = min(max(1, backoff_factor), 10)
+        swallow = True
+
+        def stop_retrying() -> None:
+            nonlocal swallow
+            swallow = False
+
+        for attempt in range(1, max_retries + 1):
+            try:
+                return func(stop_retrying, attempt)
+            except Exception:
+                if not swallow:
+                    raise
+
+            random_sleep_factor = random.uniform(1, 1 + random_factor)
+            backoff_base_secs = backoff_base_millis / 1000
+            backoff_exp_factor = backoff_factor ** (attempt - 1)
+
+            sleep_time_secs = random_sleep_factor * backoff_base_secs * backoff_exp_factor
+            time.sleep(sleep_time_secs)
+
+        return func(stop_retrying, max_retries + 1)
+
 
 class HttpClientAsync(_BaseHttpClient):
     async def call(
@@ -279,7 +348,7 @@ async def _make_request(stop_retrying: Callable, attempt: int) -> impit.Response
 
             except Exception as exc:
                 logger.debug('Request threw exception', exc_info=exc)
-                if not is_retryable_error(exc):
+                if not self._is_retryable_error(exc):
                     logger.debug('Exception is not retryable', exc_info=exc)
                     stop_retrying()
                 raise
@@ -295,10 +364,55 @@ async def _make_request(stop_retrying: Callable, attempt: int) -> impit.Response
             await response.aread()
             raise ApifyApiError(response, attempt, method=method)
 
-        return await retry_with_exp_backoff_async(
+        return await self._retry_with_exp_backoff(
             _make_request,
             max_retries=self._config.max_retries,
             backoff_base_millis=self._config.min_delay_between_retries_millis,
             backoff_factor=DEFAULT_BACKOFF_EXPONENTIAL_FACTOR,
             random_factor=DEFAULT_BACKOFF_RANDOM_FACTOR,
         )
+
+    @staticmethod
+    async def _retry_with_exp_backoff(
+        func: Callable[[Callable[[], None], int], Awaitable[T]],
+        *,
+        max_retries: int = 8,
+        backoff_base_millis: int = 500,
+        backoff_factor: float = 2,
+        random_factor: float = 1,
+    ) -> T:
+        """Retry a function with exponential backoff.
+
+        Args:
+            func: Function to retry. Receives a stop_retrying callback and attempt number.
+            max_retries: Maximum number of retry attempts.
+            backoff_base_millis: Base backoff delay in milliseconds.
+            backoff_factor: Exponential backoff multiplier (1-10).
+            random_factor: Random jitter factor (0-1).
+
+        Returns:
+            The return value of the function.
+        """
+        random_factor = min(max(0, random_factor), 1)
+        backoff_factor = min(max(1, backoff_factor), 10)
+        swallow = True
+
+        def stop_retrying() -> None:
+            nonlocal swallow
+            swallow = False
+
+        for attempt in range(1, max_retries + 1):
+            try:
+                return await func(stop_retrying, attempt)
+            except Exception:
+                if not swallow:
+                    raise
+
+            random_sleep_factor = random.uniform(1, 1 + random_factor)
+            backoff_base_secs = backoff_base_millis / 1000
+            backoff_exp_factor = backoff_factor ** (attempt - 1)
+
+            sleep_time_secs = random_sleep_factor * backoff_base_secs * backoff_exp_factor
+            await asyncio.sleep(sleep_time_secs)
+
+        return await func(stop_retrying, max_retries + 1)

src/apify_client/_resource_clients/_resource_client.py

@@ -1,10 +1,16 @@
 from __future__ import annotations
 
+import asyncio
+import math
+import time
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlencode
 
+from apify_client._consts import DEFAULT_WAIT_FOR_FINISH_SEC, DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC, ActorJobStatus
 from apify_client._logging import WithLogDetailsClient
-from apify_client._utils import to_safe_id
+from apify_client._utils import catch_not_found_or_throw, response_to_dict, to_safe_id
+from apify_client.errors import ApifyApiError, ApifyClientError
 
 if TYPE_CHECKING:
     from apify_client._client_classes import ClientRegistry, ClientRegistryAsync
@@ -117,6 +123,77 @@ def _build_params(self, **kwargs: Any) -> dict:
         merged = {**self._default_params, **kwargs}
         return {k: v for k, v in merged.items() if v is not None}
 
+    def _wait_for_finish(
+        self,
+        url: str,
+        params: dict,
+        wait_secs: int | None = None,
+    ) -> dict | None:
+        """Wait synchronously for an Actor job (run or build) to finish.
+
+        Polls the job status until it reaches a terminal state or timeout.
+        Handles 404 errors gracefully (job might not exist yet in replicas).
+
+        Args:
+            url: Full URL to the job endpoint.
+            params: Base query parameters to include in each request.
+            wait_secs: Maximum seconds to wait (None = indefinite).
+
+        Returns:
+            Job data dict when finished, or None if job doesn't exist after
+            DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC seconds.
+
+        Raises:
+            ApifyApiError: If API returns errors other than 404.
+        """
+        started_at = datetime.now(timezone.utc)
+        should_repeat = True
+        job: dict | None = None
+        seconds_elapsed = 0
+
+        while should_repeat:
+            wait_for_finish = DEFAULT_WAIT_FOR_FINISH_SEC
+            if wait_secs is not None:
+                wait_for_finish = wait_secs - seconds_elapsed
+
+            try:
+                response = self._http_client.call(
+                    url=url,
+                    method='GET',
+                    params={**params, 'waitForFinish': wait_for_finish},
+                )
+                job_response = response_to_dict(response)
+                job = job_response.get('data') if isinstance(job_response, dict) else job_response
+                seconds_elapsed = math.floor((datetime.now(timezone.utc) - started_at).total_seconds())
+
+                if not isinstance(job, dict):
+                    raise ApifyClientError(
+                        f'Unexpected response format received from the API. '
+                        f'Expected dict with "status" field, got: {type(job).__name__}'
+                    )
+
+                is_terminal = ActorJobStatus(job['status']).is_terminal
+                is_timed_out = wait_secs is not None and seconds_elapsed >= wait_secs
+                if is_terminal or is_timed_out:
+                    should_repeat = False
+
+                if not should_repeat:
+                    # Early return here so that we avoid the sleep below if not needed
+                    return job
+
+            except ApifyApiError as exc:
+                catch_not_found_or_throw(exc)
+
+                # If there are still not found errors after DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC, we give up
+                # and return None. In such case, the requested record probably really doesn't exist.
+                if seconds_elapsed > DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC:
+                    return None
+
+            # It might take some time for database replicas to get up-to-date so sleep a bit before retrying
+            time.sleep(0.25)
+
+        return job
+
 
 class ResourceClientAsync(metaclass=WithLogDetailsClient):
     """Base class for asynchronous resource clients.
@@ -223,3 +300,74 @@ def _build_params(self, **kwargs: Any) -> dict:
         """
         merged = {**self._default_params, **kwargs}
         return {k: v for k, v in merged.items() if v is not None}
+
+    async def _wait_for_finish(
+        self,
+        url: str,
+        params: dict,
+        wait_secs: int | None = None,
+    ) -> dict | None:
+        """Wait synchronously for an Actor job (run or build) to finish.
+
+        Polls the job status until it reaches a terminal state or timeout.
+        Handles 404 errors gracefully (job might not exist yet in replicas).
+
+        Args:
+            url: Full URL to the job endpoint.
+            params: Base query parameters to include in each request.
+            wait_secs: Maximum seconds to wait (None = indefinite).
+
+        Returns:
+            Job data dict when finished, or None if job doesn't exist after
+            DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC seconds.
+
+        Raises:
+            ApifyApiError: If API returns errors other than 404.
+        """
+        started_at = datetime.now(timezone.utc)
+        should_repeat = True
+        job: dict | None = None
+        seconds_elapsed = 0
+
+        while should_repeat:
+            wait_for_finish = DEFAULT_WAIT_FOR_FINISH_SEC
+            if wait_secs is not None:
+                wait_for_finish = wait_secs - seconds_elapsed
+
+            try:
+                response = await self._http_client.call(
+                    url=url,
+                    method='GET',
+                    params={**params, 'waitForFinish': wait_for_finish},
+                )
+                job_response = response_to_dict(response)
+                job = job_response.get('data') if isinstance(job_response, dict) else job_response
+
+                if not isinstance(job, dict):
+                    raise ApifyClientError(
+                        f'Unexpected response format received from the API. '
+                        f'Expected dict with "status" field, got: {type(job).__name__}'
+                    )
+
+                seconds_elapsed = math.floor((datetime.now(timezone.utc) - started_at).total_seconds())
+                is_terminal = ActorJobStatus(job['status']).is_terminal
+                is_timed_out = wait_secs is not None and seconds_elapsed >= wait_secs
+                if is_terminal or is_timed_out:
+                    should_repeat = False
+
+                if not should_repeat:
+                    # Early return here so that we avoid the sleep below if not needed
+                    return job
+
+            except ApifyApiError as exc:
+                catch_not_found_or_throw(exc)
+
+                # If there are still not found errors after DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC, we give up
+                # and return None. In such case, the requested record probably really doesn't exist.
+                if seconds_elapsed > DEFAULT_WAIT_WHEN_JOB_NOT_EXIST_SEC:
+                    return None
+
+            # It might take some time for database replicas to get up-to-date so sleep a bit before retrying
+            await asyncio.sleep(0.25)
+
+        return job