diff --git a/docs/04_upgrading/upgrading_to_v3.md b/docs/04_upgrading/upgrading_to_v3.md index 6eae9a21..803db6d8 100644 --- a/docs/04_upgrading/upgrading_to_v3.md +++ b/docs/04_upgrading/upgrading_to_v3.md @@ -69,6 +69,13 @@ Some changes in the related model classes: ## Removed Actor.config property - `Actor.config` property has been removed. Use `Actor.configuration` instead. +## Default storage ids in configuration changed to None +- `Configuration.default_key_value_store_id` changed from `'default'` to `None`. +- `Configuration.default_dataset_id` changed from `'default'` to `None`. +- `Configuration.default_request_queue_id` changed from `'default'` to `None`. + +Previously using the default storage without specifying its `id` in `Configuration` would lead to using specific storage with id `'default'`. Now it will use newly created unnamed storage with `'id'` assigned by the Apify platform, consecutive calls to get the default storage will return the same storage. + ## Actor initialization and ServiceLocator changes `Actor` initialization and global `service_locator` services setup is more strict and predictable. @@ -102,20 +109,51 @@ async def main(): ) ``` -## Removed Actor.config property -- `Actor.config` property has been removed. Use `Actor.configuration` instead. +### Changes in storage clients -## Default storage ids in configuration changed to None -- `Configuration.default_key_value_store_id` changed from `'default'` to `None`. -- `Configuration.default_dataset_id` changed from `'default'` to `None`. -- `Configuration.default_request_queue_id` changed from `'default'` to `None`. +## Explicit control over storage clients used in Actor +- It is now possible to have full control over which storage clients are used by the `Actor`. To make development of Actors convenient, the `Actor` has two storage clients. One that is used when running on Apify platform or when opening storages with `force_cloud=True` and the other client that is used when running outside the Apify platform. The `Actor` has reasonable defaults and for the majority of use-cases there is no need to change it. However, if you need to use a different storage client, you can set it up before entering `Actor` context through `service_locator`. + +**Now (v3.0):** + +```python +from crawlee import service_locator +from apify.storage_clients import ApifyStorageClient, SmartApifyStorageClient, MemoryStorageClient +from apify import Actor + + +async def main(): + service_locator.set_storage_client( + SmartApifyStorageClient( + cloud_storage_client=ApifyStorageClient(request_queue_access="single"), + local_storage_client=MemoryStorageClient() + ) + ) + async with Actor: + rq = await Actor.open_request_queue() +``` -Previously using the default storage without specifying its `id` in `Configuration` would lead to using specific storage with id `'default'`. Now it will use newly created unnamed storage with `'id'` assigned by the Apify platform, consecutive calls to get the default storage will return the same storage. -## Storages +## The default use of optimized ApifyRequestQueueClient - +- The default client for working with Apify platform based `RequestQueue` is now optimized and simplified client which does significantly lower amount of API calls, but does not support multiple consumers working on the same queue. It is cheaper and faster and is suitable for the majority of the use cases. +- The full client is still available, but it has to be explicitly requested via `request_queue_access="shared"` argument when using the `ApifyStorageClient`. -## Storage clients +**Now (v3.0):** + +```python +from crawlee import service_locator +from apify.storage_clients import ApifyStorageClient, SmartApifyStorageClient +from apify import Actor - + +async def main(): + # Full client that supports multiple consumers of the Apify Request Queue + service_locator.set_storage_client( + SmartApifyStorageClient( + cloud_storage_client=ApifyStorageClient(request_queue_access="shared"), + ) + ) + async with Actor: + rq = await Actor.open_request_queue() +``` diff --git a/pyproject.toml b/pyproject.toml index 50e348d2..a4b4fe4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ keywords = [ dependencies = [ "apify-client>=2.0.0,<3.0.0", "apify-shared>=2.0.0,<3.0.0", - "crawlee==0.6.13b42", + "crawlee==0.6.13b46", "cachetools>=5.5.0", "cryptography>=42.0.0", "impit>=0.6.1", diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 133089b2..965725e2 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -25,7 +25,6 @@ EventPersistStateData, EventSystemInfoData, ) -from crawlee.storage_clients import FileSystemStorageClient from apify._charging import ChargeResult, ChargingManager, ChargingManagerImplementation from apify._configuration import Configuration @@ -38,6 +37,7 @@ from apify.log import _configure_logging, logger from apify.storage_clients import ApifyStorageClient from apify.storage_clients._file_system import ApifyFileSystemStorageClient +from apify.storage_clients._smart_apify._storage_client import SmartApifyStorageClient from apify.storages import Dataset, KeyValueStore, RequestQueue if TYPE_CHECKING: @@ -48,7 +48,6 @@ from typing_extensions import Self from crawlee.proxy_configuration import _NewUrlFunction - from crawlee.storage_clients import StorageClient from apify._models import Webhook @@ -131,7 +130,6 @@ def __init__( self._configuration = configuration self._configure_logging = configure_logging self._apify_client: ApifyClientAsync | None = None - self._local_storage_client: StorageClient | None = None self._is_initialized = False @@ -234,45 +232,42 @@ def log(self) -> logging.Logger: """The logging.Logger instance the Actor uses.""" return logger - def _get_local_storage_client(self) -> StorageClient: - """Get the local storage client the Actor instance uses.""" - if self._local_storage_client: - return self._local_storage_client + def _raise_if_not_initialized(self) -> None: + if not self._is_initialized: + raise RuntimeError('The Actor was not initialized!') + + @cached_property + def _storage_client(self) -> SmartApifyStorageClient: + """Storage client used by the actor. + Depending on the initialization of the service locator the client can be created in different ways. + """ try: - # Set implicit default local storage client, unless local storage client was already set. - implicit_storage_client = ApifyFileSystemStorageClient() + # Nothing was set by the user. + implicit_storage_client = SmartApifyStorageClient( + local_storage_client=ApifyFileSystemStorageClient(), cloud_storage_client=ApifyStorageClient() + ) service_locator.set_storage_client(implicit_storage_client) - self._local_storage_client = implicit_storage_client except ServiceConflictError: self.log.debug( 'Storage client in service locator was set explicitly before Actor.init was called.' 'Using the existing storage client as implicit storage client for the Actor.' ) - - self._local_storage_client = service_locator.get_storage_client() - if type(self._local_storage_client) is FileSystemStorageClient: - self.log.warning( - f'Using {FileSystemStorageClient.__module__}.{FileSystemStorageClient.__name__} in Actor context is not' - f' recommended and can lead to problems with reading the input file. Use ' - f'`apify.storage_clients.FileSystemStorageClient` instead.' - ) - - return self._local_storage_client - - def _raise_if_not_initialized(self) -> None: - if not self._is_initialized: - raise RuntimeError('The Actor was not initialized!') - - def _raise_if_cloud_requested_but_not_configured(self, *, force_cloud: bool) -> None: - if not force_cloud: - return - - if not self.is_at_home() and self.configuration.token is None: - raise RuntimeError( - 'In order to use the Apify cloud storage from your computer, ' - 'you need to provide an Apify token using the APIFY_TOKEN environment variable.' - ) + else: + return implicit_storage_client + + # User set something in the service locator. + explicit_storage_client = service_locator.get_storage_client() + if isinstance(explicit_storage_client, SmartApifyStorageClient): + # The client was manually set to the right type in the service locator. This is the explicit way. + return explicit_storage_client + + raise RuntimeError( + 'The storage client in the service locator has to be instance of SmartApifyStorageClient. If you want to ' + 'set the storage client manually you have to call ' + '`service_locator.set_storage_client(SmartApifyStorageClient(...))` before entering Actor context or ' + 'awaiting `Actor.init`.' + ) async def init(self) -> None: """Initialize the Actor instance. @@ -285,6 +280,7 @@ async def init(self) -> None: This method should be called immediately before performing any additional Actor actions, and it should be called only once. """ + self.log.info('Initializing Actor...') if self._configuration: # Set explicitly the configuration in the service locator service_locator.set_configuration(self.configuration) @@ -298,22 +294,13 @@ async def init(self) -> None: if _ActorType._is_any_instance_initialized: self.log.warning('Repeated Actor initialization detected - this is non-standard usage, proceed with care') - # Create an instance of the cloud storage client, the local storage client is obtained - # from the service locator - self._cloud_storage_client = ApifyStorageClient() - # Make sure that the currently initialized instance is also available through the global `Actor` proxy cast('Proxy', Actor).__wrapped__ = self self._is_exiting = False self._was_final_persist_state_emitted = False - # If the Actor is running on the Apify platform, we set the cloud storage client. - if self.is_at_home(): - service_locator.set_storage_client(self._cloud_storage_client) - self._local_storage_client = self._cloud_storage_client - else: - self._get_local_storage_client() + self.log.debug(f'Storage client set to {self._storage_client}') service_locator.set_event_manager(self.event_manager) @@ -321,7 +308,6 @@ async def init(self) -> None: if self._configure_logging: _configure_logging() - self.log.info('Initializing Actor...') self.log.info('System info', extra=get_system_info()) await self.event_manager.__aenter__() @@ -470,16 +456,11 @@ async def open_dataset( An instance of the `Dataset` class for the given ID or name. """ self._raise_if_not_initialized() - self._raise_if_cloud_requested_but_not_configured(force_cloud=force_cloud) - - storage_client = self._cloud_storage_client if force_cloud else self._get_local_storage_client() - return await Dataset.open( id=id, - alias=alias, name=name, - configuration=self.configuration, - storage_client=storage_client, + alias=alias, + storage_client=self._storage_client.get_suitable_storage_client(force_cloud=force_cloud), ) async def open_key_value_store( @@ -509,16 +490,11 @@ async def open_key_value_store( An instance of the `KeyValueStore` class for the given ID or name. """ self._raise_if_not_initialized() - self._raise_if_cloud_requested_but_not_configured(force_cloud=force_cloud) - - storage_client = self._cloud_storage_client if force_cloud else self._get_local_storage_client() - return await KeyValueStore.open( id=id, - alias=alias, name=name, - configuration=self.configuration, - storage_client=storage_client, + alias=alias, + storage_client=self._storage_client.get_suitable_storage_client(force_cloud=force_cloud), ) async def open_request_queue( @@ -550,16 +526,11 @@ async def open_request_queue( An instance of the `RequestQueue` class for the given ID or name. """ self._raise_if_not_initialized() - self._raise_if_cloud_requested_but_not_configured(force_cloud=force_cloud) - - storage_client = self._cloud_storage_client if force_cloud else self._get_local_storage_client() - return await RequestQueue.open( id=id, - alias=alias, name=name, - configuration=self.configuration, - storage_client=storage_client, + alias=alias, + storage_client=self._storage_client.get_suitable_storage_client(force_cloud=force_cloud), ) @overload diff --git a/src/apify/storage_clients/__init__.py b/src/apify/storage_clients/__init__.py index f3e5298c..8a62e3dc 100644 --- a/src/apify/storage_clients/__init__.py +++ b/src/apify/storage_clients/__init__.py @@ -2,9 +2,11 @@ from ._apify import ApifyStorageClient from ._file_system import ApifyFileSystemStorageClient as FileSystemStorageClient +from ._smart_apify import SmartApifyStorageClient __all__ = [ 'ApifyStorageClient', 'FileSystemStorageClient', 'MemoryStorageClient', + 'SmartApifyStorageClient', ] diff --git a/src/apify/storage_clients/_apify/_request_queue_client.py b/src/apify/storage_clients/_apify/_request_queue_client.py index 893f26b9..1928f0ad 100644 --- a/src/apify/storage_clients/_apify/_request_queue_client.py +++ b/src/apify/storage_clients/_apify/_request_queue_client.py @@ -1,15 +1,8 @@ from __future__ import annotations -import asyncio -import re -from base64 import b64encode -from collections import deque -from datetime import datetime, timedelta, timezone -from hashlib import sha256 from logging import getLogger -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal -from cachetools import LRUCache from typing_extensions import override from apify_client import ApifyClientAsync @@ -18,54 +11,24 @@ from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata from crawlee.storages import RequestQueue -from ._models import ( - ApifyRequestQueueMetadata, - CachedRequest, - ProlongRequestLockResponse, - RequestQueueHead, - RequestQueueStats, -) +from ._models import ApifyRequestQueueMetadata, RequestQueueStats +from ._request_queue_shared_client import _ApifyRequestQueueSharedClient +from ._request_queue_single_client import _ApifyRequestQueueSingleClient from ._utils import AliasResolver -from apify import Request if TYPE_CHECKING: from collections.abc import Sequence from apify_client.clients import RequestQueueClientAsync + from crawlee import Request from apify import Configuration logger = getLogger(__name__) -def unique_key_to_request_id(unique_key: str, *, request_id_length: int = 15) -> str: - """Generate a deterministic request ID based on a unique key. - - Args: - unique_key: The unique key to convert into a request ID. - request_id_length: The length of the request ID. - - Returns: - A URL-safe, truncated request ID based on the unique key. - """ - # Encode the unique key and compute its SHA-256 hash - hashed_key = sha256(unique_key.encode('utf-8')).digest() - - # Encode the hash in base64 and decode it to get a string - base64_encoded = b64encode(hashed_key).decode('utf-8') - - # Remove characters that are not URL-safe ('+', '/', or '=') - url_safe_key = re.sub(r'(\+|\/|=)', '', base64_encoded) - - # Truncate the key to the desired length - return url_safe_key[:request_id_length] - - class ApifyRequestQueueClient(RequestQueueClient): - """An Apify platform implementation of the request queue client.""" - - _DEFAULT_LOCK_TIME: Final[timedelta] = timedelta(minutes=3) - """The default lock time for requests in the queue.""" + """Base class for Apify platform implementations of the request queue client.""" _MAX_CACHED_REQUESTS: Final[int] = 1_000_000 """Maximum number of requests that can be cached.""" @@ -75,6 +38,7 @@ def __init__( *, api_client: RequestQueueClientAsync, metadata: RequestQueueMetadata, + access: Literal['single', 'shared'] = 'single', ) -> None: """Initialize a new instance. @@ -83,35 +47,112 @@ def __init__( self._api_client = api_client """The Apify request queue client for API operations.""" - self._metadata = metadata - """Additional data related to the RequestQueue.""" + self._implementation: _ApifyRequestQueueSingleClient | _ApifyRequestQueueSharedClient + """Internal implementation used to communicate with the Apify platform based Request Queue.""" + if access == 'single': + self._implementation = _ApifyRequestQueueSingleClient( + api_client=self._api_client, metadata=metadata, cache_size=self._MAX_CACHED_REQUESTS + ) + elif access == 'shared': + self._implementation = _ApifyRequestQueueSharedClient( + api_client=self._api_client, + metadata=metadata, + cache_size=self._MAX_CACHED_REQUESTS, + metadata_getter=self.get_metadata, + ) + else: + raise RuntimeError(f"Unsupported access type: {access}. Allowed values are 'single' or 'shared'.") + + @property + def _metadata(self) -> RequestQueueMetadata: + return self._implementation.metadata + + @override + async def add_batch_of_requests( + self, + requests: Sequence[Request], + *, + forefront: bool = False, + ) -> AddRequestsResponse: + """Add a batch of requests to the queue. + + Args: + requests: The requests to add. + forefront: Whether to add the requests to the beginning of the queue. + + Returns: + Response containing information about the added requests. + """ + return await self._implementation.add_batch_of_requests(requests, forefront=forefront) + + @override + async def fetch_next_request(self) -> Request | None: + """Return the next request in the queue to be processed. + + Once you successfully finish processing of the request, you need to call `mark_request_as_handled` + to mark the request as handled in the queue. If there was some error in processing the request, call + `reclaim_request` instead, so that the queue will give the request to some other consumer + in another call to the `fetch_next_request` method. - self._queue_head = deque[str]() - """A deque to store request unique keys in the queue head.""" + Returns: + The request or `None` if there are no more pending requests. + """ + return await self._implementation.fetch_next_request() - self._requests_cache: LRUCache[str, CachedRequest] = LRUCache(maxsize=self._MAX_CACHED_REQUESTS) - """A cache to store request objects. Request unique key is used as the cache key.""" + @override + async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None: + """Mark a request as handled after successful processing. + + Handled requests will never again be returned by the `fetch_next_request` method. - self._queue_has_locked_requests: bool | None = None - """Whether the queue has requests locked by another client.""" + Args: + request: The request to mark as handled. - self._should_check_for_forefront_requests = False - """Whether to check for forefront requests in the next list_head call.""" + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + return await self._implementation.mark_request_as_handled(request) - self._fetch_lock = asyncio.Lock() - """Fetch lock to minimize race conditions when communicating with API.""" + @override + async def get_request(self, unique_key: str) -> Request | None: + """Get a request by unique key. - async def _get_metadata_estimate(self) -> RequestQueueMetadata: - """Try to get cached metadata first. If multiple clients, fuse with global metadata. + Args: + unique_key: Unique key of the request to get. - This method is used internally to avoid unnecessary API call unless needed (multiple clients). - Local estimation of metadata is without delay, unlike metadata from API. In situation where there is only one - client, it is the better choice. + Returns: + The request or None if not found. """ - if self._metadata.had_multiple_clients: - return await self.get_metadata() - # Get local estimation (will not include changes done bo another client) - return self._metadata + return await self._implementation.get_request(unique_key) + + @override + async def reclaim_request( + self, + request: Request, + *, + forefront: bool = False, + ) -> ProcessedRequest | None: + """Reclaim a failed request back to the queue. + + The request will be returned for processing later again by another call to `fetch_next_request`. + + Args: + request: The request to return to the queue. + forefront: Whether to add the request to the head or the end of the queue. + + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + return await self._implementation.reclaim_request(request, forefront=forefront) + + @override + async def is_empty(self) -> bool: + """Check if the queue is empty. + + Returns: + True if the queue is empty, False otherwise. + """ + return await self._implementation.is_empty() @override async def get_metadata(self) -> ApifyRequestQueueMetadata: @@ -146,6 +187,7 @@ async def open( name: str | None, alias: str | None, configuration: Configuration, + access: Literal['single', 'shared'] = 'single', ) -> ApifyRequestQueueClient: """Open an Apify request queue client. @@ -163,6 +205,17 @@ async def open( configuration: The configuration object containing API credentials and settings. Must include a valid `token` and `api_base_url`. May also contain a `default_request_queue_id` for fallback when neither `id`, `name`, nor `alias` is provided. + access: Controls the implementation of the request queue client based on expected scenario: + - 'single' is suitable for single consumer scenarios. It makes less API calls, is cheaper and faster. + - 'shared' is suitable for multiple consumers scenarios at the cost of higher API usage. + Detailed constraints for the 'single' access type: + - Only one client is consuming the request queue at the time. + - Multiple producers can put requests to the queue, but their forefront requests are not guaranteed to + be handled so quickly as this client does not aggressively fetch the forefront and relies on local + head estimation. + - Requests are only added to the queue, never deleted by other clients. (Marking as handled is ok.) + - Other producers can add new requests, but not modify existing ones. + (Modifications would not be included in local cache) Returns: An instance for the opened or created storage client. @@ -260,10 +313,7 @@ async def open( metadata_model = RequestQueueMetadata.model_validate(metadata) - return cls( - api_client=apify_rq_client, - metadata=metadata_model, - ) + return cls(api_client=apify_rq_client, metadata=metadata_model, access=access) @override async def purge(self) -> None: @@ -275,540 +325,3 @@ async def purge(self) -> None: @override async def drop(self) -> None: await self._api_client.delete() - - @override - async def add_batch_of_requests( - self, - requests: Sequence[Request], - *, - forefront: bool = False, - ) -> AddRequestsResponse: - """Add a batch of requests to the queue. - - Args: - requests: The requests to add. - forefront: Whether to add the requests to the beginning of the queue. - - Returns: - Response containing information about the added requests. - """ - # Do not try to add previously added requests to avoid pointless expensive calls to API - - new_requests: list[Request] = [] - already_present_requests: list[ProcessedRequest] = [] - - for request in requests: - if self._requests_cache.get(request.unique_key): - # We are not sure if it was already handled at this point, and it is not worth calling API for it. - # It could have been handled by another client in the meantime, so cached information about - # `request.was_already_handled` is not reliable. - already_present_requests.append( - ProcessedRequest.model_validate( - { - 'uniqueKey': request.unique_key, - 'wasAlreadyPresent': True, - 'wasAlreadyHandled': request.was_already_handled, - } - ) - ) - - else: - # Add new request to the cache. - processed_request = ProcessedRequest.model_validate( - { - 'uniqueKey': request.unique_key, - 'wasAlreadyPresent': True, - 'wasAlreadyHandled': request.was_already_handled, - } - ) - self._cache_request( - request.unique_key, - processed_request, - ) - new_requests.append(request) - - if new_requests: - # Prepare requests for API by converting to dictionaries. - requests_dict = [ - request.model_dump( - by_alias=True, - exclude={'id'}, # Exclude ID fields from requests since the API doesn't accept them. - ) - for request in new_requests - ] - - # Send requests to API. - api_response = AddRequestsResponse.model_validate( - await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront) - ) - - # Add the locally known already present processed requests based on the local cache. - api_response.processed_requests.extend(already_present_requests) - - # Remove unprocessed requests from the cache - for unprocessed_request in api_response.unprocessed_requests: - self._requests_cache.pop(unprocessed_request.unique_key, None) - - else: - api_response = AddRequestsResponse.model_validate( - {'unprocessedRequests': [], 'processedRequests': already_present_requests} - ) - - logger.debug( - f'Tried to add new requests: {len(new_requests)}, ' - f'succeeded to add new requests: {len(api_response.processed_requests) - len(already_present_requests)}, ' - f'skipped already present requests: {len(already_present_requests)}' - ) - - # Update assumed total count for newly added requests. - new_request_count = 0 - for processed_request in api_response.processed_requests: - if not processed_request.was_already_present and not processed_request.was_already_handled: - new_request_count += 1 - - self._metadata.total_request_count += new_request_count - - return api_response - - @override - async def get_request(self, unique_key: str) -> Request | None: - """Get a request by unique key. - - Args: - unique_key: Unique key of the request to get. - - Returns: - The request or None if not found. - """ - response = await self._api_client.get_request(unique_key_to_request_id(unique_key)) - - if response is None: - return None - - return Request.model_validate(response) - - @override - async def fetch_next_request(self) -> Request | None: - """Return the next request in the queue to be processed. - - Once you successfully finish processing of the request, you need to call `mark_request_as_handled` - to mark the request as handled in the queue. If there was some error in processing the request, call - `reclaim_request` instead, so that the queue will give the request to some other consumer - in another call to the `fetch_next_request` method. - - Returns: - The request or `None` if there are no more pending requests. - """ - # Ensure the queue head has requests if available. Fetching the head with lock to prevent race conditions. - async with self._fetch_lock: - await self._ensure_head_is_non_empty() - - # If queue head is empty after ensuring, there are no requests - if not self._queue_head: - return None - - # Get the next request ID from the queue head - next_unique_key = self._queue_head.popleft() - - request = await self._get_or_hydrate_request(next_unique_key) - - # Handle potential inconsistency where request might not be in the main table yet - if request is None: - logger.debug( - 'Cannot find a request from the beginning of queue, will be retried later', - extra={'nextRequestUniqueKey': next_unique_key}, - ) - return None - - # If the request was already handled, skip it - if request.handled_at is not None: - logger.debug( - 'Request fetched from the beginning of queue was already handled', - extra={'nextRequestUniqueKey': next_unique_key}, - ) - return None - - # Use get request to ensure we have the full request object. - request = await self.get_request(request.unique_key) - if request is None: - logger.debug( - 'Request fetched from the beginning of queue was not found in the RQ', - extra={'nextRequestUniqueKey': next_unique_key}, - ) - return None - - return request - - @override - async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None: - """Mark a request as handled after successful processing. - - Handled requests will never again be returned by the `fetch_next_request` method. - - Args: - request: The request to mark as handled. - - Returns: - Information about the queue operation. `None` if the given request was not in progress. - """ - # Set the handled_at timestamp if not already set - if request.handled_at is None: - request.handled_at = datetime.now(tz=timezone.utc) - - if cached_request := self._requests_cache[request.unique_key]: - cached_request.was_already_handled = request.was_already_handled - try: - # Update the request in the API - processed_request = await self._update_request(request) - processed_request.unique_key = request.unique_key - - # Update assumed handled count if this wasn't already handled - if not processed_request.was_already_handled: - self._metadata.handled_request_count += 1 - - # Update the cache with the handled request - cache_key = request.unique_key - self._cache_request( - cache_key, - processed_request, - hydrated_request=request, - ) - except Exception as exc: - logger.debug(f'Error marking request {request.unique_key} as handled: {exc!s}') - return None - else: - return processed_request - - @override - async def reclaim_request( - self, - request: Request, - *, - forefront: bool = False, - ) -> ProcessedRequest | None: - """Reclaim a failed request back to the queue. - - The request will be returned for processing later again by another call to `fetch_next_request`. - - Args: - request: The request to return to the queue. - forefront: Whether to add the request to the head or the end of the queue. - - Returns: - Information about the queue operation. `None` if the given request was not in progress. - """ - # Check if the request was marked as handled and clear it. When reclaiming, - # we want to put the request back for processing. - if request.was_already_handled: - request.handled_at = None - - # Reclaim with lock to prevent race conditions that could lead to double processing of the same request. - async with self._fetch_lock: - try: - # Update the request in the API. - processed_request = await self._update_request(request, forefront=forefront) - processed_request.unique_key = request.unique_key - - # If the request was previously handled, decrement our handled count since - # we're putting it back for processing. - if request.was_already_handled and not processed_request.was_already_handled: - self._metadata.handled_request_count -= 1 - - # Update the cache - cache_key = request.unique_key - self._cache_request( - cache_key, - processed_request, - hydrated_request=request, - ) - - # If we're adding to the forefront, we need to check for forefront requests - # in the next list_head call - if forefront: - self._should_check_for_forefront_requests = True - - # Try to release the lock on the request - try: - await self._delete_request_lock(request.unique_key, forefront=forefront) - except Exception as err: - logger.debug(f'Failed to delete request lock for request {request.unique_key}', exc_info=err) - except Exception as exc: - logger.debug(f'Error reclaiming request {request.unique_key}: {exc!s}') - return None - else: - return processed_request - - @override - async def is_empty(self) -> bool: - """Check if the queue is empty. - - Returns: - True if the queue is empty, False otherwise. - """ - # Check _list_head and self._queue_has_locked_requests with lock to make sure they are consistent. - # Without the lock the `is_empty` is prone to falsely report True with some low probability race condition. - async with self._fetch_lock: - head = await self._list_head(limit=1, lock_time=None) - return len(head.items) == 0 and not self._queue_has_locked_requests - - async def _ensure_head_is_non_empty(self) -> None: - """Ensure that the queue head has requests if they are available in the queue.""" - # If queue head has adequate requests, skip fetching more - if len(self._queue_head) > 1 and not self._should_check_for_forefront_requests: - return - - # Fetch requests from the API and populate the queue head - await self._list_head(lock_time=self._DEFAULT_LOCK_TIME) - - async def _get_or_hydrate_request(self, unique_key: str) -> Request | None: - """Get a request by unique key, either from cache or by fetching from API. - - Args: - unique_key: Unique key of the request to get. - - Returns: - The request if found and valid, otherwise None. - """ - # First check if the request is in our cache - cached_entry = self._requests_cache.get(unique_key) - - if cached_entry and cached_entry.hydrated: - # If we have the request hydrated in cache, check if lock is expired - if cached_entry.lock_expires_at and cached_entry.lock_expires_at < datetime.now(tz=timezone.utc): - # Try to prolong the lock if it's expired - try: - lock_secs = int(self._DEFAULT_LOCK_TIME.total_seconds()) - response = await self._prolong_request_lock(unique_key, lock_secs=lock_secs) - cached_entry.lock_expires_at = response.lock_expires_at - except Exception: - # If prolonging the lock fails, we lost the request - logger.debug(f'Failed to prolong lock for request {unique_key}, returning None') - return None - - return cached_entry.hydrated - - # If not in cache or not hydrated, fetch the request - try: - # Try to acquire or prolong the lock - lock_secs = int(self._DEFAULT_LOCK_TIME.total_seconds()) - await self._prolong_request_lock(unique_key, lock_secs=lock_secs) - - # Fetch the request data - request = await self.get_request(unique_key) - - # If request is not found, release lock and return None - if not request: - await self._delete_request_lock(unique_key) - return None - - # Update cache with hydrated request - cache_key = request.unique_key - self._cache_request( - cache_key, - ProcessedRequest( - unique_key=request.unique_key, - was_already_present=True, - was_already_handled=request.handled_at is not None, - ), - hydrated_request=request, - ) - except Exception as exc: - logger.debug(f'Error fetching or locking request {unique_key}: {exc!s}') - return None - else: - return request - - async def _update_request( - self, - request: Request, - *, - forefront: bool = False, - ) -> ProcessedRequest: - """Update a request in the queue. - - Args: - request: The updated request. - forefront: Whether to put the updated request in the beginning or the end of the queue. - - Returns: - The updated request - """ - request_dict = request.model_dump(by_alias=True) - request_dict['id'] = unique_key_to_request_id(request.unique_key) - response = await self._api_client.update_request( - request=request_dict, - forefront=forefront, - ) - - return ProcessedRequest.model_validate( - {'uniqueKey': request.unique_key} | response, - ) - - async def _list_head( - self, - *, - lock_time: timedelta | None = None, - limit: int = 25, - ) -> RequestQueueHead: - """Retrieve requests from the beginning of the queue. - - Args: - lock_time: Duration for which to lock the retrieved requests. - If None, requests will not be locked. - limit: Maximum number of requests to retrieve. - - Returns: - A collection of requests from the beginning of the queue. - """ - # Return from cache if available and we're not checking for new forefront requests - if self._queue_head and not self._should_check_for_forefront_requests: - logger.debug(f'Using cached queue head with {len(self._queue_head)} requests') - # Create a list of requests from the cached queue head - items = [] - for unique_key in list(self._queue_head)[:limit]: - cached_request = self._requests_cache.get(unique_key) - if cached_request and cached_request.hydrated: - items.append(cached_request.hydrated) - - metadata = await self._get_metadata_estimate() - - return RequestQueueHead( - limit=limit, - had_multiple_clients=metadata.had_multiple_clients, - queue_modified_at=metadata.modified_at, - items=items, - queue_has_locked_requests=self._queue_has_locked_requests, - lock_time=lock_time, - ) - leftover_buffer = list[str]() - if self._should_check_for_forefront_requests: - leftover_buffer = list(self._queue_head) - self._queue_head.clear() - self._should_check_for_forefront_requests = False - - # Otherwise fetch from API - lock_time = lock_time or self._DEFAULT_LOCK_TIME - lock_secs = int(lock_time.total_seconds()) - - response = await self._api_client.list_and_lock_head( - lock_secs=lock_secs, - limit=limit, - ) - - # Update the queue head cache - self._queue_has_locked_requests = response.get('queueHasLockedRequests', False) - # Check if there is another client working with the RequestQueue - self._metadata.had_multiple_clients = response.get('hadMultipleClients', False) - - for request_data in response.get('items', []): - request = Request.model_validate(request_data) - - # Skip requests without ID or unique key - if not request.unique_key: - logger.debug( - 'Skipping request from queue head, missing ID or unique key', - extra={ - 'unique_key': request.unique_key, - }, - ) - continue - - # Cache the request - self._cache_request( - request.unique_key, - ProcessedRequest( - unique_key=request.unique_key, - was_already_present=True, - was_already_handled=False, - ), - hydrated_request=request, - ) - self._queue_head.append(request.unique_key) - - for leftover_unique_key in leftover_buffer: - # After adding new requests to the forefront, any existing leftover locked request is kept in the end. - self._queue_head.append(leftover_unique_key) - return RequestQueueHead.model_validate(response) - - async def _prolong_request_lock( - self, - unique_key: str, - *, - lock_secs: int, - ) -> ProlongRequestLockResponse: - """Prolong the lock on a specific request in the queue. - - Args: - unique_key: Unique key of the request whose lock is to be prolonged. - lock_secs: The additional amount of time, in seconds, that the request will remain locked. - - Returns: - A response containing the time at which the lock will expire. - """ - response = await self._api_client.prolong_request_lock( - request_id=unique_key_to_request_id(unique_key), - # All requests reaching this code were the tip of the queue at the moment when they were fetched, - # so if their lock expires, they should be put back to the forefront as their handling is long overdue. - forefront=True, - lock_secs=lock_secs, - ) - - result = ProlongRequestLockResponse( - lock_expires_at=datetime.fromisoformat(response['lockExpiresAt'].replace('Z', '+00:00')) - ) - - # Update the cache with the new lock expiration - for cached_request in self._requests_cache.values(): - if cached_request.unique_key == unique_key: - cached_request.lock_expires_at = result.lock_expires_at - break - - return result - - async def _delete_request_lock( - self, - unique_key: str, - *, - forefront: bool = False, - ) -> None: - """Delete the lock on a specific request in the queue. - - Args: - unique_key: Unique key of the request to delete the lock. - forefront: Whether to put the request in the beginning or the end of the queue after the lock is deleted. - """ - try: - await self._api_client.delete_request_lock( - request_id=unique_key_to_request_id(unique_key), - forefront=forefront, - ) - - # Update the cache to remove the lock - for cached_request in self._requests_cache.values(): - if cached_request.unique_key == unique_key: - cached_request.lock_expires_at = None - break - except Exception as err: - logger.debug(f'Failed to delete request lock for request {unique_key}', exc_info=err) - - def _cache_request( - self, - cache_key: str, - processed_request: ProcessedRequest, - *, - hydrated_request: Request | None = None, - ) -> None: - """Cache a request for future use. - - Args: - cache_key: The key to use for caching the request. It should be request ID. - processed_request: The processed request information. - forefront: Whether the request was added to the forefront of the queue. - hydrated_request: The hydrated request object, if available. - """ - self._requests_cache[cache_key] = CachedRequest( - unique_key=processed_request.unique_key, - was_already_handled=processed_request.was_already_handled, - hydrated=hydrated_request, - lock_expires_at=None, - ) diff --git a/src/apify/storage_clients/_apify/_request_queue_shared_client.py b/src/apify/storage_clients/_apify/_request_queue_shared_client.py new file mode 100644 index 00000000..65ad8daa --- /dev/null +++ b/src/apify/storage_clients/_apify/_request_queue_shared_client.py @@ -0,0 +1,527 @@ +from __future__ import annotations + +import asyncio +from collections import deque +from datetime import datetime, timedelta, timezone +from logging import getLogger +from typing import TYPE_CHECKING, Any, Final + +from cachetools import LRUCache + +from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata + +from ._models import ApifyRequestQueueMetadata, CachedRequest, RequestQueueHead +from ._utils import unique_key_to_request_id +from apify import Request + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine, Sequence + + from apify_client.clients import RequestQueueClientAsync + + +logger = getLogger(__name__) + + +class _ApifyRequestQueueSharedClient: + """An Apify platform implementation of the request queue client. + + This implementation supports multiple producers and multiple consumers scenario. + """ + + _DEFAULT_LOCK_TIME: Final[timedelta] = timedelta(minutes=3) + """The default lock time for requests in the queue.""" + + def __init__( + self, + *, + api_client: RequestQueueClientAsync, + metadata: RequestQueueMetadata, + cache_size: int, + metadata_getter: Callable[[], Coroutine[Any, Any, ApifyRequestQueueMetadata]], + ) -> None: + """Initialize a new instance. + + Preferably use the `ApifyRequestQueueClient.open` class method to create a new instance. + """ + self.metadata = metadata + """Additional data related to the RequestQueue.""" + + self._metadata_getter = metadata_getter + """Async function to get metadata from API.""" + + self._api_client = api_client + """The Apify request queue client for API operations.""" + + self._queue_head = deque[str]() + """A deque to store request unique keys in the queue head.""" + + self._requests_cache: LRUCache[str, CachedRequest] = LRUCache(maxsize=cache_size) + """A cache to store request objects. Request unique key is used as the cache key.""" + + self._queue_has_locked_requests: bool | None = None + """Whether the queue has requests locked by another client.""" + + self._should_check_for_forefront_requests = False + """Whether to check for forefront requests in the next list_head call.""" + + self._fetch_lock = asyncio.Lock() + """Fetch lock to minimize race conditions when communicating with API.""" + + async def _get_metadata_estimate(self) -> RequestQueueMetadata: + """Try to get cached metadata first. If multiple clients, fuse with global metadata. + + This method is used internally to avoid unnecessary API call unless needed (multiple clients). + Local estimation of metadata is without delay, unlike metadata from API. In situation where there is only one + client, it is the better choice. + """ + if self.metadata.had_multiple_clients: + return await self._metadata_getter() + # Get local estimation (will not include changes done bo another client) + return self.metadata + + async def add_batch_of_requests( + self, + requests: Sequence[Request], + *, + forefront: bool = False, + ) -> AddRequestsResponse: + """Add a batch of requests to the queue. + + Args: + requests: The requests to add. + forefront: Whether to add the requests to the beginning of the queue. + + Returns: + Response containing information about the added requests. + """ + # Do not try to add previously added requests to avoid pointless expensive calls to API + + new_requests: list[Request] = [] + already_present_requests: list[ProcessedRequest] = [] + + for request in requests: + if self._requests_cache.get(request.unique_key): + # We are not sure if it was already handled at this point, and it is not worth calling API for it. + # It could have been handled by another client in the meantime, so cached information about + # `request.was_already_handled` is not reliable. + already_present_requests.append( + ProcessedRequest.model_validate( + { + 'uniqueKey': request.unique_key, + 'wasAlreadyPresent': True, + 'wasAlreadyHandled': request.was_already_handled, + } + ) + ) + + else: + # Add new request to the cache. + processed_request = ProcessedRequest.model_validate( + { + 'uniqueKey': request.unique_key, + 'wasAlreadyPresent': True, + 'wasAlreadyHandled': request.was_already_handled, + } + ) + self._cache_request( + request.unique_key, + processed_request, + ) + new_requests.append(request) + + if new_requests: + # Prepare requests for API by converting to dictionaries. + requests_dict = [ + request.model_dump( + by_alias=True, + exclude={'id'}, # Exclude ID fields from requests since the API doesn't accept them. + ) + for request in new_requests + ] + + # Send requests to API. + api_response = AddRequestsResponse.model_validate( + await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront) + ) + + # Add the locally known already present processed requests based on the local cache. + api_response.processed_requests.extend(already_present_requests) + + # Remove unprocessed requests from the cache + for unprocessed_request in api_response.unprocessed_requests: + self._requests_cache.pop(unprocessed_request.unique_key, None) + + else: + api_response = AddRequestsResponse.model_validate( + {'unprocessedRequests': [], 'processedRequests': already_present_requests} + ) + + logger.debug( + f'Tried to add new requests: {len(new_requests)}, ' + f'succeeded to add new requests: {len(api_response.processed_requests) - len(already_present_requests)}, ' + f'skipped already present requests: {len(already_present_requests)}' + ) + + # Update assumed total count for newly added requests. + new_request_count = 0 + for processed_request in api_response.processed_requests: + if not processed_request.was_already_present and not processed_request.was_already_handled: + new_request_count += 1 + + self.metadata.total_request_count += new_request_count + self.metadata.pending_request_count += new_request_count + + return api_response + + async def get_request(self, unique_key: str) -> Request | None: + """Get a request by unique key. + + Args: + unique_key: Unique key of the request to get. + + Returns: + The request or None if not found. + """ + response = await self._api_client.get_request(unique_key_to_request_id(unique_key)) + + if response is None: + return None + + return Request.model_validate(response) + + async def fetch_next_request(self) -> Request | None: + """Return the next request in the queue to be processed. + + Once you successfully finish processing of the request, you need to call `mark_request_as_handled` + to mark the request as handled in the queue. If there was some error in processing the request, call + `reclaim_request` instead, so that the queue will give the request to some other consumer + in another call to the `fetch_next_request` method. + + Returns: + The request or `None` if there are no more pending requests. + """ + # Ensure the queue head has requests if available. Fetching the head with lock to prevent race conditions. + async with self._fetch_lock: + await self._ensure_head_is_non_empty() + + # If queue head is empty after ensuring, there are no requests + if not self._queue_head: + return None + + # Get the next request ID from the queue head + next_unique_key = self._queue_head.popleft() + + request = await self._get_or_hydrate_request(next_unique_key) + + # Handle potential inconsistency where request might not be in the main table yet + if request is None: + logger.debug( + 'Cannot find a request from the beginning of queue, will be retried later', + extra={'nextRequestUniqueKey': next_unique_key}, + ) + return None + + # If the request was already handled, skip it + if request.handled_at is not None: + logger.debug( + 'Request fetched from the beginning of queue was already handled', + extra={'nextRequestUniqueKey': next_unique_key}, + ) + return None + + # Use get request to ensure we have the full request object. + request = await self.get_request(request.unique_key) + if request is None: + logger.debug( + 'Request fetched from the beginning of queue was not found in the RQ', + extra={'nextRequestUniqueKey': next_unique_key}, + ) + return None + + return request + + async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None: + """Mark a request as handled after successful processing. + + Handled requests will never again be returned by the `fetch_next_request` method. + + Args: + request: The request to mark as handled. + + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + # Set the handled_at timestamp if not already set + if request.handled_at is None: + request.handled_at = datetime.now(tz=timezone.utc) + + if cached_request := self._requests_cache[request.unique_key]: + cached_request.was_already_handled = request.was_already_handled + try: + # Update the request in the API + processed_request = await self._update_request(request) + processed_request.unique_key = request.unique_key + + # Update assumed handled count if this wasn't already handled + if not processed_request.was_already_handled: + self.metadata.handled_request_count += 1 + self.metadata.pending_request_count -= 1 + + # Update the cache with the handled request + cache_key = request.unique_key + self._cache_request( + cache_key, + processed_request, + hydrated_request=request, + ) + except Exception as exc: + logger.debug(f'Error marking request {request.unique_key} as handled: {exc!s}') + return None + else: + return processed_request + + async def reclaim_request( + self, + request: Request, + *, + forefront: bool = False, + ) -> ProcessedRequest | None: + """Reclaim a failed request back to the queue. + + The request will be returned for processing later again by another call to `fetch_next_request`. + + Args: + request: The request to return to the queue. + forefront: Whether to add the request to the head or the end of the queue. + + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + # Check if the request was marked as handled and clear it. When reclaiming, + # we want to put the request back for processing. + if request.was_already_handled: + request.handled_at = None + + # Reclaim with lock to prevent race conditions that could lead to double processing of the same request. + async with self._fetch_lock: + try: + # Update the request in the API. + processed_request = await self._update_request(request, forefront=forefront) + processed_request.unique_key = request.unique_key + + # If the request was previously handled, decrement our handled count since + # we're putting it back for processing. + if request.was_already_handled and not processed_request.was_already_handled: + self.metadata.handled_request_count -= 1 + self.metadata.pending_request_count += 1 + + # Update the cache + cache_key = request.unique_key + self._cache_request( + cache_key, + processed_request, + hydrated_request=request, + ) + + # If we're adding to the forefront, we need to check for forefront requests + # in the next list_head call + if forefront: + self._should_check_for_forefront_requests = True + + except Exception as exc: + logger.debug(f'Error reclaiming request {request.unique_key}: {exc!s}') + return None + else: + return processed_request + + async def is_empty(self) -> bool: + """Check if the queue is empty. + + Returns: + True if the queue is empty, False otherwise. + """ + # Check _list_head. + # Without the lock the `is_empty` is prone to falsely report True with some low probability race condition. + async with self._fetch_lock: + head = await self._list_head(limit=1) + return len(head.items) == 0 and not self._queue_has_locked_requests + + async def _ensure_head_is_non_empty(self) -> None: + """Ensure that the queue head has requests if they are available in the queue.""" + # If queue head has adequate requests, skip fetching more + if len(self._queue_head) > 1 and not self._should_check_for_forefront_requests: + return + + # Fetch requests from the API and populate the queue head + await self._list_head() + + async def _get_or_hydrate_request(self, unique_key: str) -> Request | None: + """Get a request by unique key, either from cache or by fetching from API. + + Args: + unique_key: Unique key of the request to get. + + Returns: + The request if found and valid, otherwise None. + """ + # First check if the request is in our cache + cached_entry = self._requests_cache.get(unique_key) + + if cached_entry and cached_entry.hydrated: + # If we have the request hydrated in cache, return it + return cached_entry.hydrated + + # If not in cache or not hydrated, fetch the request + try: + # Fetch the request data + request = await self.get_request(unique_key) + + # If request is not found and return None + if not request: + return None + + # Update cache with hydrated request + cache_key = request.unique_key + self._cache_request( + cache_key, + ProcessedRequest( + unique_key=request.unique_key, + was_already_present=True, + was_already_handled=request.handled_at is not None, + ), + hydrated_request=request, + ) + except Exception as exc: + logger.debug(f'Error fetching request {unique_key}: {exc!s}') + return None + else: + return request + + async def _update_request( + self, + request: Request, + *, + forefront: bool = False, + ) -> ProcessedRequest: + """Update a request in the queue. + + Args: + request: The updated request. + forefront: Whether to put the updated request in the beginning or the end of the queue. + + Returns: + The updated request + """ + request_dict = request.model_dump(by_alias=True) + request_dict['id'] = unique_key_to_request_id(request.unique_key) + response = await self._api_client.update_request( + request=request_dict, + forefront=forefront, + ) + + return ProcessedRequest.model_validate( + {'uniqueKey': request.unique_key} | response, + ) + + async def _list_head( + self, + *, + limit: int = 25, + ) -> RequestQueueHead: + """Retrieve requests from the beginning of the queue. + + Args: + limit: Maximum number of requests to retrieve. + + Returns: + A collection of requests from the beginning of the queue. + """ + # Return from cache if available and we're not checking for new forefront requests + if self._queue_head and not self._should_check_for_forefront_requests: + logger.debug(f'Using cached queue head with {len(self._queue_head)} requests') + # Create a list of requests from the cached queue head + items = [] + for unique_key in list(self._queue_head)[:limit]: + cached_request = self._requests_cache.get(unique_key) + if cached_request and cached_request.hydrated: + items.append(cached_request.hydrated) + + metadata = await self._get_metadata_estimate() + + return RequestQueueHead( + limit=limit, + had_multiple_clients=metadata.had_multiple_clients, + queue_modified_at=metadata.modified_at, + items=items, + lock_time=None, + queue_has_locked_requests=self._queue_has_locked_requests, + ) + leftover_buffer = list[str]() + if self._should_check_for_forefront_requests: + leftover_buffer = list(self._queue_head) + self._queue_head.clear() + self._should_check_for_forefront_requests = False + + # Otherwise fetch from API + response = await self._api_client.list_and_lock_head( + lock_secs=int(self._DEFAULT_LOCK_TIME.total_seconds()), + limit=limit, + ) + + # Update the queue head cache + self._queue_has_locked_requests = response.get('queueHasLockedRequests', False) + # Check if there is another client working with the RequestQueue + self.metadata.had_multiple_clients = response.get('hadMultipleClients', False) + + for request_data in response.get('items', []): + request = Request.model_validate(request_data) + + # Skip requests without ID or unique key + if not request.unique_key: + logger.debug( + 'Skipping request from queue head, missing unique key', + extra={ + 'unique_key': request.unique_key, + }, + ) + continue + + # Cache the request + self._cache_request( + request.unique_key, + ProcessedRequest( + unique_key=request.unique_key, + was_already_present=True, + was_already_handled=False, + ), + hydrated_request=request, + ) + self._queue_head.append(request.unique_key) + + for leftover_unique_key in leftover_buffer: + # After adding new requests to the forefront, any existing leftover locked request is kept in the end. + self._queue_head.append(leftover_unique_key) + return RequestQueueHead.model_validate(response) + + def _cache_request( + self, + cache_key: str, + processed_request: ProcessedRequest, + *, + hydrated_request: Request | None = None, + ) -> None: + """Cache a request for future use. + + Args: + cache_key: The key to use for caching the request. It should be request ID. + processed_request: The processed request information. + forefront: Whether the request was added to the forefront of the queue. + hydrated_request: The hydrated request object, if available. + """ + self._requests_cache[cache_key] = CachedRequest( + unique_key=processed_request.unique_key, + was_already_handled=processed_request.was_already_handled, + hydrated=hydrated_request, + lock_expires_at=None, + ) diff --git a/src/apify/storage_clients/_apify/_request_queue_single_client.py b/src/apify/storage_clients/_apify/_request_queue_single_client.py new file mode 100644 index 00000000..1b9502b0 --- /dev/null +++ b/src/apify/storage_clients/_apify/_request_queue_single_client.py @@ -0,0 +1,399 @@ +from __future__ import annotations + +from collections import deque +from datetime import datetime, timezone +from logging import getLogger +from typing import TYPE_CHECKING, Final + +from cachetools import LRUCache + +from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata + +from apify import Request +from apify.storage_clients._apify._utils import unique_key_to_request_id + +if TYPE_CHECKING: + from collections.abc import Sequence + + from apify_client.clients import RequestQueueClientAsync + + +logger = getLogger(__name__) + + +class _ApifyRequestQueueSingleClient: + """An Apify platform implementation of the request queue client with limited capability. + + This client is designed to use as little resources as possible, but has to be used in constrained context. + Constraints: + - Only one client is consuming the request queue at the time. + - Multiple producers can put requests to the queue, but their forefront requests are not guaranteed to be handled + so quickly as this client does not aggressively fetch the forefront and relies on local head estimation. + - Requests are only added to the queue, never deleted. (Marking as handled is ok.) + - Other producers can add new requests, but not modify existing ones (otherwise caching can miss the updates) + + If the constraints are not met, the client might work in an unpredictable way. + """ + + _MAX_HEAD_ITEMS: Final[int] = 1000 + """The maximum head items read count limited by API.""" + + def __init__( + self, + *, + api_client: RequestQueueClientAsync, + metadata: RequestQueueMetadata, + cache_size: int, + ) -> None: + """Initialize a new instance. + + Preferably use the `ApifyRequestQueueClient.open` class method to create a new instance. + """ + self.metadata = metadata + """Additional data related to the RequestQueue.""" + + self._api_client = api_client + """The Apify request queue client for API operations.""" + + self._requests_cache: LRUCache[str, Request] = LRUCache(maxsize=cache_size) + """A cache to store request objects. Request unique key is used as the cache key.""" + + self._head_requests: deque[str] = deque() + """Ordered unique keys of requests that represent queue head.""" + + self._requests_already_handled: set[str] = set() + """Local estimation of requests unique keys that are already present and handled on the platform. + + - To enhance local deduplication. + - To reduce the _requests_cache size. Already handled requests are most likely not going to be needed again, + so no need to cache more than their unique_key. + """ + + self._requests_in_progress: set[str] = set() + """Set of requests unique keys that are being processed locally. + + - To help decide if the RQ is finished or not. This is the only consumer, so it can be tracked locally. + """ + + self._initialized_caches = False + """This flag indicates whether the local caches were already initialized. + + Initialization is done lazily only if deduplication is needed (When calling add_batch_of_requests). + """ + + async def add_batch_of_requests( + self, + requests: Sequence[Request], + *, + forefront: bool = False, + ) -> AddRequestsResponse: + """Add a batch of requests to the queue. + + Args: + requests: The requests to add. + forefront: Whether to add the requests to the beginning of the queue. + + Returns: + Response containing information about the added requests. + """ + if not self._initialized_caches: + # One time process to initialize local caches for existing request queues. + await self._init_caches() + self._initialized_caches = True + + new_requests: list[Request] = [] + already_present_requests: list[ProcessedRequest] = [] + + for request in requests: + # Check if request is known to be already handled (it has to be present as well.) + if request.unique_key in self._requests_already_handled: + already_present_requests.append( + ProcessedRequest.model_validate( + { + 'uniqueKey': request.unique_key, + 'wasAlreadyPresent': True, + 'wasAlreadyHandled': True, + } + ) + ) + # Check if request is known to be already present, but unhandled + elif self._requests_cache.get(request.unique_key): + already_present_requests.append( + ProcessedRequest.model_validate( + { + 'uniqueKey': request.unique_key, + 'wasAlreadyPresent': True, + 'wasAlreadyHandled': request.was_already_handled, + } + ) + ) + else: + # Push the request to the platform. Probably not there, or we are not aware of it + new_requests.append(request) + + # Update local caches + self._requests_cache[request.unique_key] = request + if forefront: + self._head_requests.append(request.unique_key) + else: + self._head_requests.appendleft(request.unique_key) + + if new_requests: + # Prepare requests for API by converting to dictionaries. + requests_dict = [ + request.model_dump( + by_alias=True, + ) + for request in new_requests + ] + + # Send requests to API. + api_response = AddRequestsResponse.model_validate( + await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront) + ) + # Add the locally known already present processed requests based on the local cache. + api_response.processed_requests.extend(already_present_requests) + # Remove unprocessed requests from the cache + for unprocessed_request in api_response.unprocessed_requests: + self._requests_cache.pop(unprocessed_request.unique_key, None) + + else: + api_response = AddRequestsResponse.model_validate( + {'unprocessedRequests': [], 'processedRequests': already_present_requests} + ) + + # Update assumed total count for newly added requests. + new_request_count = 0 + for processed_request in api_response.processed_requests: + if not processed_request.was_already_present and not processed_request.was_already_handled: + new_request_count += 1 + self.metadata.total_request_count += new_request_count + self.metadata.pending_request_count += new_request_count + + return api_response + + async def get_request(self, unique_key: str) -> Request | None: + """Get a request by unique key. + + Args: + unique_key: Unique key of the request to get. + + Returns: + The request or None if not found. + """ + if unique_key in self._requests_cache: + return self._requests_cache[unique_key] + + response = await self._api_client.get_request(unique_key_to_request_id(unique_key)) + + if response is None: + return None + + return Request.model_validate(response) + + async def fetch_next_request(self) -> Request | None: + """Return the next request in the queue to be processed. + + Once you successfully finish processing of the request, you need to call `mark_request_as_handled` + to mark the request as handled in the queue. If there was some error in processing the request, call + `reclaim_request` instead, so that the queue will give the request to some other consumer + in another call to the `fetch_next_request` method. + + Returns: + The request or `None` if there are no more pending requests. + """ + await self._ensure_head_is_non_empty() + + while self._head_requests: + request_unique_key = self._head_requests.pop() + if ( + request_unique_key not in self._requests_in_progress + and request_unique_key not in self._requests_already_handled + ): + self._requests_in_progress.add(request_unique_key) + return await self.get_request(request_unique_key) + # No request locally and the ones returned from the platform are already in progress. + return None + + async def _ensure_head_is_non_empty(self) -> None: + """Ensure that the queue head has requests if they are available in the queue.""" + if len(self._head_requests) <= 1: + await self._list_head() + + async def _list_head(self) -> None: + desired_new_head_items = 200 + # The head will contain in progress requests as well, so we need to fetch more, to get some new ones. + requested_head_items = max(self._MAX_HEAD_ITEMS, desired_new_head_items + len(self._requests_in_progress)) + response = await self._api_client.list_head(limit=requested_head_items) + + # Update metadata + # Check if there is another client working with the RequestQueue + self.metadata.had_multiple_clients = response.get('hadMultipleClients', False) + # Should warn once? This might be outside expected context if the other consumers consumes at the same time + + if modified_at := response.get('queueModifiedAt'): + self.metadata.modified_at = max(self.metadata.modified_at, modified_at) + + # Update the cached data + for request_data in response.get('items', []): + request = Request.model_validate(request_data) + + if request.unique_key in self._requests_in_progress: + # Ignore requests that are already in progress, we will not process them again. + continue + if request.was_already_handled: + # Do not cache fully handled requests, we do not need them. Just cache their unique_key. + self._requests_already_handled.add(request.unique_key) + else: + # Only fetch the request if we do not know it yet. + if request.unique_key not in self._requests_cache: + request = Request.model_validate( + await self._api_client.get_request(unique_key_to_request_id(request.unique_key)) + ) + self._requests_cache[request.unique_key] = request + + # Add new requests to the end of the head, unless already present in head + if request.unique_key not in self._head_requests: + self._head_requests.appendleft(request.unique_key) + + async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None: + """Mark a request as handled after successful processing. + + Handled requests will never again be returned by the `fetch_next_request` method. + + Args: + request: The request to mark as handled. + + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + # Set the handled_at timestamp if not already set + + if request.handled_at is None: + request.handled_at = datetime.now(tz=timezone.utc) + self.metadata.handled_request_count += 1 + self.metadata.pending_request_count -= 1 + + if cached_request := self._requests_cache.get(request.unique_key): + cached_request.handled_at = request.handled_at + + try: + # Update the request in the API + # Works as upsert - adds the request if it does not exist yet. (Local request that was handled before + # adding to the queue.) + processed_request = await self._update_request(request) + # Remember that we handled this request, to optimize local deduplication. + self._requests_already_handled.add(request.unique_key) + # Remove request from cache. It will most likely not be needed. + self._requests_cache.pop(request.unique_key) + self._requests_in_progress.discard(request.unique_key) + + except Exception as exc: + logger.debug(f'Error marking request {request.unique_key} as handled: {exc!s}') + return None + else: + return processed_request + + async def reclaim_request( + self, + request: Request, + *, + forefront: bool = False, + ) -> ProcessedRequest | None: + """Reclaim a failed request back to the queue. + + The request will be returned for processing later again by another call to `fetch_next_request`. + + Args: + request: The request to return to the queue. + forefront: Whether to add the request to the head or the end of the queue. + + Returns: + Information about the queue operation. `None` if the given request was not in progress. + """ + # Check if the request was marked as handled and clear it. When reclaiming, + # we want to put the request back for processing. + if request.was_already_handled: + request.handled_at = None + + try: + # Make sure request is in the local cache. We might need it. + self._requests_cache[request.unique_key] = request + + # No longer in progress + self._requests_in_progress.discard(request.unique_key) + # No longer handled + self._requests_already_handled.discard(request.unique_key) + + if forefront: + # Append to top of the local head estimation + self._head_requests.append(request.unique_key) + + processed_request = await self._update_request(request, forefront=forefront) + processed_request.unique_key = request.unique_key + # If the request was previously handled, decrement our handled count since + # we're putting it back for processing. + if request.was_already_handled and not processed_request.was_already_handled: + self.metadata.handled_request_count -= 1 + self.metadata.pending_request_count += 1 + + except Exception as exc: + logger.debug(f'Error reclaiming request {request.unique_key}: {exc!s}') + return None + else: + return processed_request + + async def is_empty(self) -> bool: + """Check if the queue is empty. + + Returns: + True if the queue is empty, False otherwise. + """ + # Without the lock the `is_empty` is prone to falsely report True with some low probability race condition. + await self._ensure_head_is_non_empty() + return not self._head_requests and not self._requests_in_progress + + async def _update_request( + self, + request: Request, + *, + forefront: bool = False, + ) -> ProcessedRequest: + """Update a request in the queue. + + Args: + request: The updated request. + forefront: Whether to put the updated request in the beginning or the end of the queue. + + Returns: + The updated request + """ + request_dict = request.model_dump(by_alias=True) + request_dict['id'] = unique_key_to_request_id(request.unique_key) + response = await self._api_client.update_request( + request=request_dict, + forefront=forefront, + ) + + return ProcessedRequest.model_validate( + {'uniqueKey': request.unique_key} | response, + ) + + async def _init_caches(self) -> None: + """Initialize the local caches by getting requests from the existing queue. + + This is mainly done to improve local deduplication capability. List request can return up to 10k requests, but + their order is implementation detail and does not respect head order or insertion order. + + Deduplication on platform is expensive, it takes 1 API call per request and 1 write operation per request. + Local deduplication is cheaper, it takes 1 API call for whole cache and 1 read operation per request. + """ + response = await self._api_client.list_requests(limit=10_000) + for request_data in response.get('items', []): + request = Request.model_validate(request_data) + if request.was_already_handled: + # Cache just unique_key for deduplication + self._requests_already_handled.add(request.unique_key) + else: + # Cache full request + self._requests_cache[request.unique_key] = request diff --git a/src/apify/storage_clients/_apify/_storage_client.py b/src/apify/storage_clients/_apify/_storage_client.py index 2ff3fad9..2bee6527 100644 --- a/src/apify/storage_clients/_apify/_storage_client.py +++ b/src/apify/storage_clients/_apify/_storage_client.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from typing_extensions import override @@ -23,15 +23,35 @@ class ApifyStorageClient(StorageClient): """Apify storage client.""" + def __init__(self, *, request_queue_access: Literal['single', 'shared'] = 'single') -> None: + """Initialize the Apify storage client. + + Args: + request_queue_access: Controls the implementation of the request queue client based on expected scenario: + - 'single' is suitable for single consumer scenarios. It makes less API calls, is cheaper and faster. + - 'shared' is suitable for multiple consumers scenarios at the cost of higher API usage. + Detailed constraints for the 'single' access type: + - Only one client is consuming the request queue at the time. + - Multiple producers can put requests to the queue, but their forefront requests are not guaranteed to + be handled so quickly as this client does not aggressively fetch the forefront and relies on local + head estimation. + - Requests are only added to the queue, never deleted by other clients. (Marking as handled is ok.) + - Other producers can add new requests, but not modify existing ones. + (Modifications would not be included in local cache) + """ + self._request_queue_access = request_queue_access + # This class breaches Liskov Substitution Principle. It requires specialized Configuration compared to its parent. _lsp_violation_error_message_template = ( 'Expected "configuration" to be an instance of "apify.Configuration", but got {} instead.' ) @override - def get_additional_cache_key(self, configuration: CrawleeConfiguration) -> Hashable: + def get_storage_client_cache_key(self, configuration: CrawleeConfiguration) -> Hashable: if isinstance(configuration, ApifyConfiguration): - return hash_api_base_url_and_token(configuration) + # It is not supported to open exactly same queue with 'single' and 'shared' client at the same time. + # Whichever client variation gets used first, wins. + return super().get_storage_client_cache_key(configuration), hash_api_base_url_and_token(configuration) config_class = type(configuration) raise TypeError( @@ -79,6 +99,8 @@ async def create_rq_client( ) -> ApifyRequestQueueClient: configuration = configuration or ApifyConfiguration.get_global_configuration() if isinstance(configuration, ApifyConfiguration): - return await ApifyRequestQueueClient.open(id=id, name=name, alias=alias, configuration=configuration) + return await ApifyRequestQueueClient.open( + id=id, name=name, alias=alias, configuration=configuration, access=self._request_queue_access + ) raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__)) diff --git a/src/apify/storage_clients/_apify/_utils.py b/src/apify/storage_clients/_apify/_utils.py index ebae80f7..eee87367 100644 --- a/src/apify/storage_clients/_apify/_utils.py +++ b/src/apify/storage_clients/_apify/_utils.py @@ -1,7 +1,10 @@ from __future__ import annotations import logging +import re from asyncio import Lock +from base64 import b64encode +from hashlib import sha256 from logging import getLogger from typing import TYPE_CHECKING, ClassVar @@ -107,7 +110,7 @@ async def store_mapping(self, storage_id: str) -> None: # Update in-memory mapping (await self._get_alias_map())[self._storage_key] = storage_id if not Configuration.get_global_configuration().is_at_home: - logging.getLogger(__name__).warning( + logging.getLogger(__name__).debug( 'AliasResolver storage limited retention is only supported on Apify platform. Storage is not exported.' ) return @@ -166,3 +169,26 @@ def hash_api_base_url_and_token(configuration: Configuration) -> str: if configuration.api_public_base_url is None or configuration.token is None: raise ValueError("'Configuration.api_public_base_url' and 'Configuration.token' must be set.") return compute_short_hash(f'{configuration.api_public_base_url}{configuration.token}'.encode()) + + +def unique_key_to_request_id(unique_key: str, *, request_id_length: int = 15) -> str: + """Generate a deterministic request ID based on a unique key. + + Args: + unique_key: The unique key to convert into a request ID. + request_id_length: The length of the request ID. + + Returns: + A URL-safe, truncated request ID based on the unique key. + """ + # Encode the unique key and compute its SHA-256 hash + hashed_key = sha256(unique_key.encode('utf-8')).digest() + + # Encode the hash in base64 and decode it to get a string + base64_encoded = b64encode(hashed_key).decode('utf-8') + + # Remove characters that are not URL-safe ('+', '/', or '=') + url_safe_key = re.sub(r'(\+|\/|=)', '', base64_encoded) + + # Truncate the key to the desired length + return url_safe_key[:request_id_length] diff --git a/src/apify/storage_clients/_smart_apify/__init__.py b/src/apify/storage_clients/_smart_apify/__init__.py new file mode 100644 index 00000000..605be630 --- /dev/null +++ b/src/apify/storage_clients/_smart_apify/__init__.py @@ -0,0 +1 @@ +from ._storage_client import SmartApifyStorageClient diff --git a/src/apify/storage_clients/_smart_apify/_storage_client.py b/src/apify/storage_clients/_smart_apify/_storage_client.py new file mode 100644 index 00000000..db1b8b5d --- /dev/null +++ b/src/apify/storage_clients/_smart_apify/_storage_client.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from typing_extensions import override + +from crawlee.storage_clients._base import DatasetClient, KeyValueStoreClient, RequestQueueClient, StorageClient + +from apify._configuration import Configuration as ApifyConfiguration +from apify._utils import docs_group +from apify.storage_clients import ApifyStorageClient +from apify.storage_clients._file_system import ApifyFileSystemStorageClient + +if TYPE_CHECKING: + from collections.abc import Hashable + + from crawlee.configuration import Configuration as CrawleeConfiguration + + +@docs_group('Storage clients') +class SmartApifyStorageClient(StorageClient): + """SmartApifyStorageClient that delegates to cloud_storage_client or local_storage_client. + + When running on Apify platform use cloud_storage_client, else use local_storage_client. This storage client is + designed to work specifically in Actor context. + """ + + def __init__( + self, + *, + cloud_storage_client: ApifyStorageClient | None = None, + local_storage_client: StorageClient | None = None, + ) -> None: + """Initialize the Apify storage client. + + Args: + cloud_storage_client: Client used to communicate with the Apify platform storage. Either through + `force_cloud` argument when opening storages or automatically when running on the Apify platform. + local_storage_client: Client used to communicate with the storage when not running on the Apify + platform and not using `force_cloud` argument when opening storages. + """ + self._cloud_storage_client = cloud_storage_client or ApifyStorageClient(request_queue_access='single') + self._local_storage_client = local_storage_client or ApifyFileSystemStorageClient() + + def __str__(self) -> str: + return ( + f'{self.__class__.__name__}(cloud_storage_client={self._cloud_storage_client.__class__.__name__},' + f' local_storage_client={self._local_storage_client.__class__.__name__})' + ) + + def get_suitable_storage_client(self, *, force_cloud: bool = False) -> StorageClient: + """Get a suitable storage client based on the global configuration and the value of the force_cloud flag. + + Args: + force_cloud: If True, return `cloud_storage_client`. + """ + if ApifyConfiguration.get_global_configuration().is_at_home: + return self._cloud_storage_client + + configuration = ApifyConfiguration.get_global_configuration() + if force_cloud: + if configuration.token is None: + raise RuntimeError( + 'In order to use the Apify cloud storage from your computer, ' + 'you need to provide an Apify token using the APIFY_TOKEN environment variable.' + ) + return self._cloud_storage_client + + return self._local_storage_client + + @override + def get_storage_client_cache_key(self, configuration: CrawleeConfiguration) -> Hashable: + if ApifyConfiguration.get_global_configuration().is_at_home: + if isinstance(configuration, ApifyConfiguration): + return self._cloud_storage_client.get_storage_client_cache_key(configuration) + raise TypeError('Expecting ApifyConfiguration') + + return self._local_storage_client.get_storage_client_cache_key(configuration) + + @override + async def create_dataset_client( + self, + *, + id: str | None = None, + name: str | None = None, + alias: str | None = None, + configuration: CrawleeConfiguration | None = None, + ) -> DatasetClient: + return await self.get_suitable_storage_client().create_dataset_client( + id=id, name=id, alias=alias, configuration=configuration + ) + + @override + async def create_kvs_client( + self, + *, + id: str | None = None, + name: str | None = None, + alias: str | None = None, + configuration: CrawleeConfiguration | None = None, + ) -> KeyValueStoreClient: + return await self.get_suitable_storage_client().create_kvs_client( + id=id, name=id, alias=alias, configuration=configuration + ) + + @override + async def create_rq_client( + self, + *, + id: str | None = None, + name: str | None = None, + alias: str | None = None, + configuration: CrawleeConfiguration | None = None, + ) -> RequestQueueClient: + return await self.get_suitable_storage_client().create_rq_client( + id=id, name=id, alias=alias, configuration=configuration + ) diff --git a/tests/integration/actor_source_base/requirements.txt b/tests/integration/actor_source_base/requirements.txt index e9bc6dad..7f4f8246 100644 --- a/tests/integration/actor_source_base/requirements.txt +++ b/tests/integration/actor_source_base/requirements.txt @@ -1,4 +1,4 @@ # The test fixture will put the Apify SDK wheel path on the next line APIFY_SDK_WHEEL_PLACEHOLDER uvicorn[standard] -crawlee[parsel]==0.6.13b42 +crawlee[parsel]==0.6.13b46 diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 52bee53a..aea770db 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -20,14 +20,15 @@ from ._utils import generate_unique_resource_name from apify import Actor from apify._models import ActorRun +from apify.storage_clients import ApifyStorageClient from apify.storage_clients._apify._utils import AliasResolver +from apify.storages import RequestQueue if TYPE_CHECKING: from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Iterator, Mapping from decimal import Decimal from apify_client.clients.resource_clients import ActorClientAsync - from crawlee.storages import RequestQueue _TOKEN_ENV_VAR = 'APIFY_TEST_USER_API_TOKEN' _API_URL_ENV_VAR = 'APIFY_INTEGRATION_TESTS_API_URL' @@ -50,6 +51,9 @@ def prepare_test_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Callabl """ def _prepare_test_env() -> None: + # Reset the Actor class state. + apify._actor.Actor.__wrapped__.__class__._is_any_instance_initialized = False # type: ignore[attr-defined] + apify._actor.Actor.__wrapped__.__class__._is_rebooting = False # type: ignore[attr-defined] delattr(apify._actor.Actor, '__wrapped__') # Set the environment variable for the local storage directory to the temporary path. @@ -103,14 +107,15 @@ def apify_client_async(apify_token: str) -> ApifyClientAsync: return ApifyClientAsync(apify_token, api_url=api_url) -@pytest.fixture -async def request_queue_force_cloud(apify_token: str, monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[RequestQueue]: +@pytest.fixture(params=['single', 'shared']) +async def request_queue_apify( + apify_token: str, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> AsyncGenerator[RequestQueue]: """Create an instance of the Apify request queue on the platform and drop it when the test is finished.""" - request_queue_name = generate_unique_resource_name('request_queue') monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_token) async with Actor: - rq = await Actor.open_request_queue(name=request_queue_name, force_cloud=True) + rq = await RequestQueue.open(storage_client=ApifyStorageClient(request_queue_access=request.param)) yield rq await rq.drop() diff --git a/tests/integration/test_actor_request_queue.py b/tests/integration/test_actor_request_queue.py index 9521234b..3a9053c7 100644 --- a/tests/integration/test_actor_request_queue.py +++ b/tests/integration/test_actor_request_queue.py @@ -12,26 +12,29 @@ from ._utils import generate_unique_resource_name from apify import Actor, Request from apify._models import ActorRun +from apify.storage_clients import ApifyStorageClient +from apify.storages import RequestQueue if TYPE_CHECKING: from collections.abc import AsyncGenerator from apify_client import ApifyClientAsync - from crawlee.storages import RequestQueue from .conftest import MakeActorFunction, RunActorFunction -@pytest.fixture +@pytest.fixture(params=['single', 'shared']) async def apify_named_rq( - apify_client_async: ApifyClientAsync, monkeypatch: pytest.MonkeyPatch + apify_client_async: ApifyClientAsync, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest ) -> AsyncGenerator[RequestQueue]: assert apify_client_async.token monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_client_async.token) request_queue_name = generate_unique_resource_name('request_queue') async with Actor: - request_queue = await Actor.open_request_queue(name=request_queue_name, force_cloud=True) + request_queue = await RequestQueue.open( + name=request_queue_name, storage_client=ApifyStorageClient(request_queue_access=request.param) + ) yield request_queue await request_queue.drop() diff --git a/tests/integration/test_apify_storages.py b/tests/integration/test_apify_storages.py index 83ad7ebd..32cb5061 100644 --- a/tests/integration/test_apify_storages.py +++ b/tests/integration/test_apify_storages.py @@ -5,8 +5,9 @@ from crawlee import service_locator from crawlee.storages import Dataset, KeyValueStore, RequestQueue +from .conftest import MakeActorFunction, RunActorFunction from apify import Actor, Configuration -from apify.storage_clients import ApifyStorageClient +from apify.storage_clients import ApifyStorageClient, MemoryStorageClient, SmartApifyStorageClient @pytest.mark.parametrize( @@ -66,10 +67,86 @@ async def test_aliases_not_stored_on_platform_when_local( ) -> None: """Test that default Apify storage used locally is not persisting aliases to Apify based default KVS.""" service_locator.set_configuration(Configuration(token=apify_token)) - service_locator.set_storage_client(ApifyStorageClient()) async with Actor(configure_logging=False): await storage_type.open(alias='test') default_kvs = await Actor.open_key_value_store(force_cloud=True) # The default KVS should be empty assert len(await default_kvs.list_keys()) == 0 + + +async def test_actor_full_explicit_storage_init(apify_token: str) -> None: + service_locator.set_configuration(Configuration(token=apify_token)) + service_locator.set_storage_client( + SmartApifyStorageClient( + local_storage_client=MemoryStorageClient(), + cloud_storage_client=ApifyStorageClient(request_queue_access='shared'), + ) + ) + async with Actor: + # If service locator was already set with SmartApifyStorageClient, the actor will use it. + # Storages should be different when force_cloud is used outside the Apify platform + assert await Actor.open_dataset() is not await Actor.open_dataset(force_cloud=True) + assert await Actor.open_key_value_store() is not await Actor.open_key_value_store(force_cloud=True) + assert await Actor.open_request_queue() is not await Actor.open_request_queue(force_cloud=True) + + +async def test_actor_full_explicit_storage_init_same_client(apify_token: str) -> None: + service_locator.set_configuration(Configuration(token=apify_token)) + service_locator.set_storage_client( + SmartApifyStorageClient( + local_storage_client=ApifyStorageClient(request_queue_access='shared'), + cloud_storage_client=ApifyStorageClient(request_queue_access='shared'), + ) + ) + async with Actor: + # If service locator was already set with SmartApifyStorageClient, the actor will use it. + # Storages should be same as the equivalent storage client is for both local and cloud storage client + assert await Actor.open_dataset() is await Actor.open_dataset(force_cloud=True) + assert await Actor.open_key_value_store() is await Actor.open_key_value_store(force_cloud=True) + assert await Actor.open_request_queue() is await Actor.open_request_queue(force_cloud=True) + + +async def test_actor_partial_explicit_cloud_storage_init(apify_token: str) -> None: + service_locator.set_configuration(Configuration(token=apify_token)) + service_locator.set_storage_client(ApifyStorageClient(request_queue_access='shared')) + with pytest.raises( + RuntimeError, match=r'^The storage client in the service locator has to be instance of SmartApifyStorageClient' + ): + async with Actor: + # If service locator was explicitly set to something different than SmartApifyStorageClient, raise an error. + ... + + +async def test_actor_implicit_storage_init(apify_token: str) -> None: + service_locator.set_configuration(Configuration(token=apify_token)) + async with Actor: + assert await Actor.open_dataset() is not await Actor.open_dataset(force_cloud=True) + assert await Actor.open_key_value_store() is not await Actor.open_key_value_store(force_cloud=True) + assert await Actor.open_request_queue() is not await Actor.open_request_queue(force_cloud=True) + + +async def test_actor_full_explicit_storage_init_on_platform( + make_actor: MakeActorFunction, run_actor: RunActorFunction +) -> None: + async def main() -> None: + from crawlee import service_locator + + from apify.storage_clients import ApifyStorageClient, MemoryStorageClient, SmartApifyStorageClient + + service_locator.set_storage_client( + SmartApifyStorageClient( + local_storage_client=MemoryStorageClient(), + cloud_storage_client=ApifyStorageClient(request_queue_access='shared'), + ) + ) + async with Actor: + # Storages should be same as the cloud client is used on the platform + assert await Actor.open_dataset() is await Actor.open_dataset(force_cloud=True) + assert await Actor.open_key_value_store() is await Actor.open_key_value_store(force_cloud=True) + assert await Actor.open_request_queue() is await Actor.open_request_queue(force_cloud=True) + + actor = await make_actor(label='explicit_storage_init', main_func=main) + run_result = await run_actor(actor) + + assert run_result.status == 'SUCCEEDED' diff --git a/tests/integration/test_request_queue.py b/tests/integration/test_request_queue.py index 17ef094f..fbbdfb74 100644 --- a/tests/integration/test_request_queue.py +++ b/tests/integration/test_request_queue.py @@ -1,751 +1,601 @@ from __future__ import annotations import asyncio -from typing import TYPE_CHECKING, cast +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Literal, cast import pytest -from crawlee import Request +from apify_shared.consts import ApifyEnvVars +from crawlee import Request, service_locator +from crawlee.crawlers import BasicCrawler +from ._utils import generate_unique_resource_name from apify import Actor +from apify.storage_clients import ApifyStorageClient +from apify.storages import RequestQueue if TYPE_CHECKING: from apify_client import ApifyClientAsync - from crawlee.storages import RequestQueue + from crawlee._types import BasicCrawlingContext from .conftest import MakeActorFunction, RunActorFunction from apify.storage_clients._apify._models import ApifyRequestQueueMetadata -async def test_add_and_fetch_requests( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_add_and_fetch_requests(request_queue_apify: RequestQueue) -> None: """Test basic functionality of adding and fetching requests.""" - async def main() -> None: - async with Actor: - desired_request_count = 100 - Actor.log.info('Opening request queue...') - rq = await Actor.open_request_queue() - - # Add some requests - for i in range(desired_request_count): - Actor.log.info(f'Adding request {i}...') - await rq.add_request(f'https://example.com/{i}') - - handled_request_count = 0 - while next_request := await rq.fetch_next_request(): - Actor.log.info('Fetching next request...') - queue_operation_info = await rq.mark_request_as_handled(next_request) - assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' - assert queue_operation_info.was_already_handled is False, ( - f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' - ) - handled_request_count += 1 - - assert handled_request_count == desired_request_count, ( - f'handled_request_count={handled_request_count}', - f'desired_request_count={desired_request_count}', - ) - Actor.log.info('Waiting for queue to be finished...') - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' - - actor = await make_actor(label='rq-simple-test', main_func=main) - run_result = await run_actor(actor) - - assert run_result.status == 'SUCCEEDED' + desired_request_count = 100 + Actor.log.info('Opening request queue...') + rq = request_queue_apify + + # Add some requests + for i in range(desired_request_count): + Actor.log.info(f'Adding request {i}...') + await rq.add_request(f'https://example.com/{i}') + + handled_request_count = 0 + while next_request := await rq.fetch_next_request(): + Actor.log.info('Fetching next request...') + queue_operation_info = await rq.mark_request_as_handled(next_request) + assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' + assert queue_operation_info.was_already_handled is False, ( + f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' + ) + handled_request_count += 1 + + assert handled_request_count == desired_request_count, ( + f'handled_request_count={handled_request_count}', + f'desired_request_count={desired_request_count}', + ) + Actor.log.info('Waiting for queue to be finished...') + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' -async def test_add_requests_in_batches( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_add_requests_in_batches(request_queue_apify: RequestQueue) -> None: """Test adding multiple requests in a single batch operation.""" - async def main() -> None: - async with Actor: - desired_request_count = 100 - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add some requests - await rq.add_requests([f'https://example.com/{i}' for i in range(desired_request_count)]) - total_count = await rq.get_total_count() - Actor.log.info(f'Added {desired_request_count} requests in batch, total in queue: {total_count}') - - handled_request_count = 0 - while next_request := await rq.fetch_next_request(): - if handled_request_count % 20 == 0: - Actor.log.info(f'Processing request {handled_request_count + 1}...') - queue_operation_info = await rq.mark_request_as_handled(next_request) - assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' - assert queue_operation_info.was_already_handled is False, ( - f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' - ) - handled_request_count += 1 - - assert handled_request_count == desired_request_count, ( - f'handled_request_count={handled_request_count}', - f'desired_request_count={desired_request_count}', - ) - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' - - actor = await make_actor(label='rq-batch-test', main_func=main) - run_result = await run_actor(actor) - - assert run_result.status == 'SUCCEEDED' + desired_request_count = 100 + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Add some requests + await rq.add_requests([f'https://example.com/{i}' for i in range(desired_request_count)]) + total_count = await rq.get_total_count() + Actor.log.info(f'Added {desired_request_count} requests in batch, total in queue: {total_count}') + + handled_request_count = 0 + while next_request := await rq.fetch_next_request(): + if handled_request_count % 20 == 0: + Actor.log.info(f'Processing request {handled_request_count + 1}...') + queue_operation_info = await rq.mark_request_as_handled(next_request) + assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' + assert queue_operation_info.was_already_handled is False, ( + f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' + ) + handled_request_count += 1 + + assert handled_request_count == desired_request_count, ( + f'handled_request_count={handled_request_count}', + f'desired_request_count={desired_request_count}', + ) + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' -async def test_add_non_unique_requests_in_batch( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_add_non_unique_requests_in_batch(request_queue_apify: RequestQueue) -> None: """Test adding requests with duplicate unique keys in batch.""" - async def main() -> None: - from apify import Request - - async with Actor: - desired_request_count = 100 - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add some requests - requests_to_add = [ - Request.from_url(f'https://example.com/{i}', unique_key=str(i - 1 if i % 4 == 1 else i)) - for i in range(desired_request_count) - ] - await rq.add_requests(requests_to_add) - total_count = await rq.get_total_count() - Actor.log.info( - f'Added {desired_request_count} requests with duplicate unique keys, total in queue: {total_count}' - ) - - handled_request_count = 0 - while next_request := await rq.fetch_next_request(): - if handled_request_count % 20 == 0: - Actor.log.info(f'Processing request {handled_request_count + 1}: {next_request.url}') - queue_operation_info = await rq.mark_request_as_handled(next_request) - assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' - assert queue_operation_info.was_already_handled is False, ( - f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' - ) - handled_request_count += 1 - - expected_count = int(desired_request_count * 3 / 4) - assert handled_request_count == expected_count, ( - f'handled_request_count={handled_request_count}', - f'expected_count={expected_count}', - ) - is_finished = await rq.is_finished() - Actor.log.info(f'Processed {handled_request_count}/{expected_count} requests, finished: {is_finished}') - assert is_finished is True, f'is_finished={is_finished}' - - actor = await make_actor(label='rq-batch-test', main_func=main) - run_result = await run_actor(actor) - - assert run_result.status == 'SUCCEEDED' + desired_request_count = 100 + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Add some requests + requests_to_add = [ + Request.from_url(f'https://example.com/{i}', unique_key=str(i - 1 if i % 4 == 1 else i)) + for i in range(desired_request_count) + ] + await rq.add_requests(requests_to_add) + total_count = await rq.get_total_count() + Actor.log.info(f'Added {desired_request_count} requests with duplicate unique keys, total in queue: {total_count}') + + handled_request_count = 0 + while next_request := await rq.fetch_next_request(): + if handled_request_count % 20 == 0: + Actor.log.info(f'Processing request {handled_request_count + 1}: {next_request.url}') + queue_operation_info = await rq.mark_request_as_handled(next_request) + assert queue_operation_info is not None, f'queue_operation_info={queue_operation_info}' + assert queue_operation_info.was_already_handled is False, ( + f'queue_operation_info.was_already_handled={queue_operation_info.was_already_handled}' + ) + handled_request_count += 1 + + expected_count = int(desired_request_count * 3 / 4) + assert handled_request_count == expected_count, ( + f'handled_request_count={handled_request_count}', + f'expected_count={expected_count}', + ) + is_finished = await rq.is_finished() + Actor.log.info(f'Processed {handled_request_count}/{expected_count} requests, finished: {is_finished}') + assert is_finished is True, f'is_finished={is_finished}' -async def test_forefront_requests_ordering( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_forefront_requests_ordering(request_queue_apify: RequestQueue) -> None: """Test that forefront requests are processed before regular requests.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add regular requests - await rq.add_request('https://example.com/1') - await rq.add_request('https://example.com/2') - await rq.add_request('https://example.com/3') - Actor.log.info('Added 3 regular requests') - - # Add forefront requests - await rq.add_request('https://example.com/priority1', forefront=True) - await rq.add_request('https://example.com/priority2', forefront=True) - total_count = await rq.get_total_count() - Actor.log.info(f'Added 2 forefront requests, total in queue: {total_count}') - - # Fetch requests and verify order - fetched_urls = [] - while next_request := await rq.fetch_next_request(): - Actor.log.info(f'Fetched request: {next_request.url}') - fetched_urls.append(next_request.url) - await rq.mark_request_as_handled(next_request) - - # Forefront requests should come first (in reverse order of addition) - expected_order = [ - 'https://example.com/priority2', - 'https://example.com/priority1', - 'https://example.com/1', - 'https://example.com/2', - 'https://example.com/3', - ] - assert fetched_urls == expected_order, ( - f'fetched_urls={fetched_urls}', - f'expected_order={expected_order}', - ) - - actor = await make_actor(label='rq-forefront-order-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Add regular requests + await rq.add_request('https://example.com/1') + await rq.add_request('https://example.com/2') + await rq.add_request('https://example.com/3') + Actor.log.info('Added 3 regular requests') + + # Add forefront requests + await rq.add_request('https://example.com/priority1', forefront=True) + await rq.add_request('https://example.com/priority2', forefront=True) + total_count = await rq.get_total_count() + Actor.log.info(f'Added 2 forefront requests, total in queue: {total_count}') + + # Fetch requests and verify order + fetched_urls = [] + while next_request := await rq.fetch_next_request(): + Actor.log.info(f'Fetched request: {next_request.url}') + fetched_urls.append(next_request.url) + await rq.mark_request_as_handled(next_request) + + # Forefront requests should come first (in reverse order of addition) + expected_order = [ + 'https://example.com/priority2', + 'https://example.com/priority1', + 'https://example.com/1', + 'https://example.com/2', + 'https://example.com/3', + ] + assert fetched_urls == expected_order, ( + f'fetched_urls={fetched_urls}', + f'expected_order={expected_order}', + ) -async def test_request_unique_key_behavior( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_request_unique_key_behavior(request_queue_apify: RequestQueue) -> None: """Test behavior of custom unique keys.""" - async def main() -> None: - from apify import Request - - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add requests with custom unique keys - req1 = Request.from_url('https://example.com/page1', unique_key='custom-key-1') - req2 = Request.from_url('https://example.com/page2', unique_key='custom-key-1') # Same key - req3 = Request.from_url('https://example.com/page3', unique_key='custom-key-2') # Different key + rq = request_queue_apify + Actor.log.info('Request queue opened') - result1 = await rq.add_request(req1) - result2 = await rq.add_request(req2) - result3 = await rq.add_request(req3) + # Add requests with custom unique keys + req1 = Request.from_url('https://example.com/page1', unique_key='custom-key-1') + req2 = Request.from_url('https://example.com/page2', unique_key='custom-key-1') # Same key + req3 = Request.from_url('https://example.com/page3', unique_key='custom-key-2') # Different key - Actor.log.info( - f'Added requests - was_already_present: [{result1.was_already_present}, ' - f'{result2.was_already_present}, {result3.was_already_present}]' - ) + result1 = await rq.add_request(req1) + result2 = await rq.add_request(req2) + result3 = await rq.add_request(req3) - # Second request should be marked as already present - assert result1.was_already_present is False, f'result1.was_already_present={result1.was_already_present}' - assert result2.was_already_present is True, f'result2.was_already_present={result2.was_already_present}' - assert result3.was_already_present is False, f'result3.was_already_present={result3.was_already_present}' - - # Only 2 requests should be fetchable - fetched_count = 0 - fetched_requests = [] - while next_request := await rq.fetch_next_request(): - fetched_count += 1 - fetched_requests.append(next_request) - await rq.mark_request_as_handled(next_request) - - assert fetched_count == 2, f'fetched_count={fetched_count}' - - # Verify the fetched requests have the correct unique keys - unique_keys = {req.unique_key for req in fetched_requests} - expected_keys = {'custom-key-1', 'custom-key-2'} - assert unique_keys == expected_keys, ( - f'unique_keys={unique_keys}', - f'expected_keys={expected_keys}', - ) + Actor.log.info( + f'Added requests - was_already_present: [{result1.was_already_present}, ' + f'{result2.was_already_present}, {result3.was_already_present}]' + ) - actor = await make_actor(label='rq-unique-key-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + # Second request should be marked as already present + assert result1.was_already_present is False, f'result1.was_already_present={result1.was_already_present}' + assert result2.was_already_present is True, f'result2.was_already_present={result2.was_already_present}' + assert result3.was_already_present is False, f'result3.was_already_present={result3.was_already_present}' + + # Only 2 requests should be fetchable + fetched_count = 0 + fetched_requests = [] + while next_request := await rq.fetch_next_request(): + fetched_count += 1 + fetched_requests.append(next_request) + await rq.mark_request_as_handled(next_request) + + assert fetched_count == 2, f'fetched_count={fetched_count}' + + # Verify the fetched requests have the correct unique keys + unique_keys = {req.unique_key for req in fetched_requests} + expected_keys = {'custom-key-1', 'custom-key-2'} + assert unique_keys == expected_keys, ( + f'unique_keys={unique_keys}', + f'expected_keys={expected_keys}', + ) -async def test_request_reclaim_functionality( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_request_reclaim_functionality(request_queue_apify: RequestQueue) -> None: """Test request reclaiming for failed processing.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Add a test request - await rq.add_request('https://example.com/test') - Actor.log.info('Added test request') + # Add a test request + await rq.add_request('https://example.com/test') + Actor.log.info('Added test request') - # Fetch and reclaim the request - request = await rq.fetch_next_request() - assert request is not None, f'request={request}' - Actor.log.info(f'Fetched request: {request.url}') - - # Reclaim the request (simulate failed processing) - reclaim_result = await rq.reclaim_request(request) - assert reclaim_result is not None, f'reclaim_result={reclaim_result}' - assert reclaim_result.was_already_handled is False, ( - f'reclaim_result.was_already_handled={reclaim_result.was_already_handled}' - ) - Actor.log.info('Request reclaimed successfully') - - # Should be able to fetch the same request again - request2 = await rq.fetch_next_request() - assert request2 is not None, f'request2={request2}' - assert request2.url == request.url, ( - f'request2.url={request2.url}', - f'request.url={request.url}', - ) - Actor.log.info(f'Successfully fetched reclaimed request: {request2.url}') + # Fetch and reclaim the request + request = await rq.fetch_next_request() + assert request is not None, f'request={request}' + Actor.log.info(f'Fetched request: {request.url}') - # Mark as handled this time - await rq.mark_request_as_handled(request2) - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' + # Reclaim the request (simulate failed processing) + reclaim_result = await rq.reclaim_request(request) + assert reclaim_result is not None, f'reclaim_result={reclaim_result}' + assert reclaim_result.was_already_handled is False, ( + f'reclaim_result.was_already_handled={reclaim_result.was_already_handled}' + ) + Actor.log.info('Request reclaimed successfully') + + # Should be able to fetch the same request again + request2 = await rq.fetch_next_request() + assert request2 is not None, f'request2={request2}' + assert request2.url == request.url, ( + f'request2.url={request2.url}', + f'request.url={request.url}', + ) + Actor.log.info(f'Successfully fetched reclaimed request: {request2.url}') - actor = await make_actor(label='rq-reclaim-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + # Mark as handled this time + await rq.mark_request_as_handled(request2) + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' -async def test_request_reclaim_with_forefront( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_request_reclaim_with_forefront(request_queue_apify: RequestQueue) -> None: """Test reclaiming requests to the front of the queue.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add multiple requests - await rq.add_request('https://example.com/1') - await rq.add_request('https://example.com/2') - await rq.add_request('https://example.com/3') - Actor.log.info('Added 3 requests') - - # Fetch first request - first_request = await rq.fetch_next_request() - assert first_request is not None, f'first_request={first_request}' - Actor.log.info(f'Fetched first request: {first_request.url}') - - # Reclaim to forefront - await rq.reclaim_request(first_request, forefront=True) - Actor.log.info('Request reclaimed to forefront') - - # The reclaimed request should be fetched first again - next_request = await rq.fetch_next_request() - assert next_request is not None, f'next_request={next_request}' - assert next_request.url == first_request.url, ( - f'next_request.url={next_request.url}', - f'first_request.url={first_request.url}', - ) - Actor.log.info(f'Confirmed reclaimed request came first: {next_request.url}') - - # Clean up - await rq.mark_request_as_handled(next_request) - remaining_count = 0 + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Add multiple requests + await rq.add_request('https://example.com/1') + await rq.add_request('https://example.com/2') + await rq.add_request('https://example.com/3') + Actor.log.info('Added 3 requests') + + # Fetch first request + first_request = await rq.fetch_next_request() + assert first_request is not None, f'first_request={first_request}' + Actor.log.info(f'Fetched first request: {first_request.url}') + + # Reclaim to forefront + await rq.reclaim_request(first_request, forefront=True) + Actor.log.info('Request reclaimed to forefront') + + # The reclaimed request should be fetched first again + next_request = await rq.fetch_next_request() + assert next_request is not None, f'next_request={next_request}' + assert next_request.url == first_request.url, ( + f'next_request.url={next_request.url}', + f'first_request.url={first_request.url}', + ) + Actor.log.info(f'Confirmed reclaimed request came first: {next_request.url}') - while next_request := await rq.fetch_next_request(): - remaining_count += 1 - await rq.mark_request_as_handled(next_request) + # Clean up + await rq.mark_request_as_handled(next_request) + remaining_count = 0 - Actor.log.info(f'Test completed - processed {remaining_count} additional requests') + while next_request := await rq.fetch_next_request(): + remaining_count += 1 + await rq.mark_request_as_handled(next_request) - actor = await make_actor(label='rq-reclaim-forefront-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + Actor.log.info(f'Test completed - processed {remaining_count} additional requests') -async def test_complex_request_objects( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_complex_request_objects(request_queue_apify: RequestQueue) -> None: """Test handling complex Request objects with various properties.""" - async def main() -> None: - from apify import Request - - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Create request with various properties - request = Request.from_url( - 'https://example.com/api/data', - method='POST', - headers={'Authorization': 'Bearer token123', 'Content-Type': 'application/json'}, - user_data={'category': 'api', 'priority': 'high'}, - unique_key='api-request-1', - ) - await rq.add_request(request) - Actor.log.info(f'Added complex request: {request.url} with method {request.method}') - - # Fetch and verify all properties are preserved - fetched_request = await rq.fetch_next_request() - assert fetched_request is not None, f'fetched_request={fetched_request}' - Actor.log.info(f'Fetched request: {fetched_request.url}') - - assert fetched_request.url == 'https://example.com/api/data', f'fetched_request.url={fetched_request.url}' - assert fetched_request.method == 'POST', f'fetched_request.method={fetched_request.method}' - assert fetched_request.headers['Authorization'] == 'Bearer token123', ( - f'fetched_request.headers["Authorization"]={fetched_request.headers["Authorization"]}' - ) - assert fetched_request.headers['Content-Type'] == 'application/json', ( - f'fetched_request.headers["Content-Type"]={fetched_request.headers["Content-Type"]}' - ) - assert fetched_request.user_data['category'] == 'api', ( - f'fetched_request.user_data["category"]={fetched_request.user_data["category"]}' - ) - assert fetched_request.user_data['priority'] == 'high', ( - f'fetched_request.user_data["priority"]={fetched_request.user_data["priority"]}' - ) - assert fetched_request.unique_key == 'api-request-1', ( - f'fetched_request.unique_key={fetched_request.unique_key}' - ) - Actor.log.info('All properties verified successfully') + # Create request with various properties + request = Request.from_url( + 'https://example.com/api/data', + method='POST', + headers={'Authorization': 'Bearer token123', 'Content-Type': 'application/json'}, + user_data={'category': 'api', 'priority': 'high'}, + unique_key='api-request-1', + ) + await rq.add_request(request) + Actor.log.info(f'Added complex request: {request.url} with method {request.method}') + + # Fetch and verify all properties are preserved + fetched_request = await rq.fetch_next_request() + assert fetched_request is not None, f'fetched_request={fetched_request}' + Actor.log.info(f'Fetched request: {fetched_request.url}') + + assert fetched_request.url == 'https://example.com/api/data', f'fetched_request.url={fetched_request.url}' + assert fetched_request.method == 'POST', f'fetched_request.method={fetched_request.method}' + assert fetched_request.headers['Authorization'] == 'Bearer token123', ( + f'fetched_request.headers["Authorization"]={fetched_request.headers["Authorization"]}' + ) + assert fetched_request.headers['Content-Type'] == 'application/json', ( + f'fetched_request.headers["Content-Type"]={fetched_request.headers["Content-Type"]}' + ) + assert fetched_request.user_data['category'] == 'api', ( + f'fetched_request.user_data["category"]={fetched_request.user_data["category"]}' + ) + assert fetched_request.user_data['priority'] == 'high', ( + f'fetched_request.user_data["priority"]={fetched_request.user_data["priority"]}' + ) + assert fetched_request.unique_key == 'api-request-1', f'fetched_request.unique_key={fetched_request.unique_key}' + Actor.log.info('All properties verified successfully') - await rq.mark_request_as_handled(fetched_request) - Actor.log.info('Complex request test completed') - - actor = await make_actor(label='rq-complex-request-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + await rq.mark_request_as_handled(fetched_request) + Actor.log.info('Complex request test completed') -async def test_get_request_by_unique_key( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_get_request_by_unique_key(request_queue_apify: RequestQueue) -> None: """Test retrieving specific requests by their unique_key.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add a request and get its unique_key - add_result = await rq.add_request('https://example.com/test') - request_unique_key = add_result.unique_key - Actor.log.info(f'Request added with unique_key: {request_unique_key}') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Retrieve the request by unique_key - retrieved_request = await rq.get_request(request_unique_key) - assert retrieved_request is not None, f'retrieved_request={retrieved_request}' - assert retrieved_request.url == 'https://example.com/test', f'retrieved_request.url={retrieved_request.url}' - assert retrieved_request.unique_key == request_unique_key, (f'{request_unique_key=}',) - Actor.log.info('Request retrieved successfully by unique_key') + # Add a request and get its unique_key + add_result = await rq.add_request('https://example.com/test') + request_unique_key = add_result.unique_key + Actor.log.info(f'Request added with unique_key: {request_unique_key}') - # Test with non-existent unique_key - non_existent_request = await rq.get_request('non-existent-unique_key') - assert non_existent_request is None, f'non_existent_request={non_existent_request}' - Actor.log.info('Non-existent unique_key correctly returned None') + # Retrieve the request by unique_key + retrieved_request = await rq.get_request(request_unique_key) + assert retrieved_request is not None, f'retrieved_request={retrieved_request}' + assert retrieved_request.url == 'https://example.com/test', f'retrieved_request.url={retrieved_request.url}' + assert retrieved_request.unique_key == request_unique_key, (f'{request_unique_key=}',) + Actor.log.info('Request retrieved successfully by unique_key') - actor = await make_actor(label='rq-get-by-unique-key-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + # Test with non-existent unique_key + non_existent_request = await rq.get_request('non-existent-unique_key') + assert non_existent_request is None, f'non_existent_request={non_existent_request}' + Actor.log.info('Non-existent unique_key correctly returned None') -async def test_metadata_tracking( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_metadata_tracking(request_queue_apify: RequestQueue) -> None: """Test request queue metadata and counts.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Check initial state - initial_total = await rq.get_total_count() - initial_handled = await rq.get_handled_count() - Actor.log.info(f'Initial state - Total: {initial_total}, Handled: {initial_handled}') - assert initial_total == 0, f'initial_total={initial_total}' - assert initial_handled == 0, f'initial_handled={initial_handled}' - - # Add requests - await rq.add_requests([f'https://example.com/{i}' for i in range(5)]) - Actor.log.info('Added 5 requests in batch') - - # Check counts after adding - total_after_add = await rq.get_total_count() - handled_after_add = await rq.get_handled_count() - Actor.log.info(f'After adding - Total: {total_after_add}, Handled: {handled_after_add}') - assert total_after_add == 5, f'total_after_add={total_after_add}' - assert handled_after_add == 0, f'handled_after_add={handled_after_add}' - - # Process some requests - for _ in range(3): - request = await rq.fetch_next_request() - if request: - await rq.mark_request_as_handled(request) + # Check initial state + initial_total = await rq.get_total_count() + initial_handled = await rq.get_handled_count() + Actor.log.info(f'Initial state - Total: {initial_total}, Handled: {initial_handled}') + assert initial_total == 0, f'initial_total={initial_total}' + assert initial_handled == 0, f'initial_handled={initial_handled}' - Actor.log.info('Processed 3 requests') + # Add requests + await rq.add_requests([f'https://example.com/{i}' for i in range(5)]) + Actor.log.info('Added 5 requests in batch') - # Check counts after processing - final_total = await rq.get_total_count() - final_handled = await rq.get_handled_count() - Actor.log.info(f'Final state - Total: {final_total}, Handled: {final_handled}') - assert final_total == 5, f'final_total={final_total}' - assert final_handled == 3, f'final_handled={final_handled}' + # Check counts after adding + total_after_add = await rq.get_total_count() + handled_after_add = await rq.get_handled_count() + Actor.log.info(f'After adding - Total: {total_after_add}, Handled: {handled_after_add}') + assert total_after_add == 5, f'total_after_add={total_after_add}' + assert handled_after_add == 0, f'handled_after_add={handled_after_add}' - actor = await make_actor(label='rq-metadata-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + # Process some requests + for _ in range(3): + request = await rq.fetch_next_request() + if request: + await rq.mark_request_as_handled(request) + Actor.log.info('Processed 3 requests') -async def test_batch_operations_performance( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: + # Check counts after processing + final_total = await rq.get_total_count() + final_handled = await rq.get_handled_count() + Actor.log.info(f'Final state - Total: {final_total}, Handled: {final_handled}') + assert final_total == 5, f'final_total={final_total}' + assert final_handled == 3, f'final_handled={final_handled}' + + +async def test_batch_operations_performance(request_queue_apify: RequestQueue) -> None: """Test batch operations vs individual operations.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Test batch add vs individual adds - batch_requests = [f'https://example.com/batch/{i}' for i in range(50)] - Actor.log.info(f'Prepared {len(batch_requests)} requests for batch add') + # Test batch add vs individual adds + batch_requests = [f'https://example.com/batch/{i}' for i in range(50)] + Actor.log.info(f'Prepared {len(batch_requests)} requests for batch add') - # Add in batch - await rq.add_requests(batch_requests) - Actor.log.info('Batch add completed') + # Add in batch + await rq.add_requests(batch_requests) + Actor.log.info('Batch add completed') - # Verify all requests were added - total_count = await rq.get_total_count() - handled_count = await rq.get_handled_count() - Actor.log.info(f'After batch add - Total: {total_count}, Handled: {handled_count}') - assert total_count == 50, f'total_count={total_count}' - assert handled_count == 0, f'handled_count={handled_count}' - - # Process all requests - processed_count = 0 - while next_request := await rq.fetch_next_request(): - processed_count += 1 - await rq.mark_request_as_handled(next_request) - if processed_count >= 50: # Safety break - break + # Verify all requests were added + total_count = await rq.get_total_count() + handled_count = await rq.get_handled_count() + Actor.log.info(f'After batch add - Total: {total_count}, Handled: {handled_count}') + assert total_count == 50, f'total_count={total_count}' + assert handled_count == 0, f'handled_count={handled_count}' - Actor.log.info(f'Processing completed. Total processed: {processed_count}') - assert processed_count == 50, f'processed_count={processed_count}' + # Process all requests + processed_count = 0 + while next_request := await rq.fetch_next_request(): + processed_count += 1 + await rq.mark_request_as_handled(next_request) + if processed_count >= 50: # Safety break + break - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' + Actor.log.info(f'Processing completed. Total processed: {processed_count}') + assert processed_count == 50, f'processed_count={processed_count}' - actor = await make_actor(label='rq-batch-performance-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' -async def test_state_consistency( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_state_consistency(request_queue_apify: RequestQueue) -> None: """Test queue state consistency during concurrent operations.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Add initial requests - for i in range(10): - await rq.add_request(f'https://example.com/{i}') - Actor.log.info('Added 10 initial requests') + # Add initial requests + for i in range(10): + await rq.add_request(f'https://example.com/{i}') + Actor.log.info('Added 10 initial requests') - initial_total = await rq.get_total_count() - Actor.log.info(f'Initial total count: {initial_total}') + initial_total = await rq.get_total_count() + Actor.log.info(f'Initial total count: {initial_total}') - # Simulate some requests being processed and others being reclaimed - processed_requests = [] - reclaimed_requests = [] + # Simulate some requests being processed and others being reclaimed + processed_requests = [] + reclaimed_requests = [] - for i in range(5): - request = await rq.fetch_next_request() - if request: - if i % 2 == 0: # Process even indices - await rq.mark_request_as_handled(request) - processed_requests.append(request) - else: # Reclaim odd indices - await rq.reclaim_request(request) - reclaimed_requests.append(request) - - Actor.log.info(f'Processed {len(processed_requests)} requests, reclaimed {len(reclaimed_requests)}') + for i in range(5): + request = await rq.fetch_next_request() + if request: + if i % 2 == 0: # Process even indices + await rq.mark_request_as_handled(request) + processed_requests.append(request) + else: # Reclaim odd indices + await rq.reclaim_request(request) + reclaimed_requests.append(request) - # Verify queue state - expected_handled = len(processed_requests) - current_handled = await rq.get_handled_count() - current_total = await rq.get_total_count() + Actor.log.info(f'Processed {len(processed_requests)} requests, reclaimed {len(reclaimed_requests)}') - Actor.log.info(f'Expected handled: {expected_handled}, Actual handled: {current_handled}') - Actor.log.info(f'Current total: {current_total}') + # Verify queue state + expected_handled = len(processed_requests) + current_handled = await rq.get_handled_count() + current_total = await rq.get_total_count() - assert current_handled == expected_handled, ( - f'current_handled={current_handled}', - f'expected_handled={expected_handled}', - ) - assert current_total == 10, f'current_total={current_total}' + Actor.log.info(f'Expected handled: {expected_handled}, Actual handled: {current_handled}') + Actor.log.info(f'Current total: {current_total}') - # Process remaining requests - remaining_count = 0 - while next_request := await rq.fetch_next_request(): - remaining_count += 1 - await rq.mark_request_as_handled(next_request) + assert current_handled == expected_handled, ( + f'current_handled={current_handled}', + f'expected_handled={expected_handled}', + ) + assert current_total == 10, f'current_total={current_total}' - Actor.log.info(f'Processed {remaining_count} remaining requests') - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' + # Process remaining requests + remaining_count = 0 + while next_request := await rq.fetch_next_request(): + remaining_count += 1 + await rq.mark_request_as_handled(next_request) - actor = await make_actor(label='rq-state-consistency-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + Actor.log.info(f'Processed {remaining_count} remaining requests') + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' -async def test_empty_rq_behavior( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_empty_rq_behavior(request_queue_apify: RequestQueue) -> None: """Test behavior with empty queues.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Test empty queue operations - is_empty = await rq.is_empty() - is_finished = await rq.is_finished() - Actor.log.info(f'Empty queue - is_empty: {is_empty}, is_finished: {is_finished}') - assert is_empty is True, f'is_empty={is_empty}' - assert is_finished is True, f'is_finished={is_finished}' - - # Fetch from empty queue - request = await rq.fetch_next_request() - Actor.log.info(f'Fetch result from empty queue: {request}') - assert request is None, f'request={request}' - - # Check metadata for empty queue - metadata = await rq.get_metadata() - assert metadata is not None, f'metadata={metadata}' - Actor.log.info( - f'Empty queue metadata - Total: {metadata.total_request_count}, ' - f'Handled: {metadata.handled_request_count}, ' - f'Pending: {metadata.pending_request_count}' - ) - assert metadata.total_request_count == 0, f'metadata.total_request_count={metadata.total_request_count}' - assert metadata.handled_request_count == 0, ( - f'metadata.handled_request_count={metadata.handled_request_count}' - ) - assert metadata.pending_request_count == 0, ( - f'metadata.pending_request_count={metadata.pending_request_count}' - ) - - actor = await make_actor(label='rq-empty-queue-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Test empty queue operations + is_empty = await rq.is_empty() + is_finished = await rq.is_finished() + Actor.log.info(f'Empty queue - is_empty: {is_empty}, is_finished: {is_finished}') + assert is_empty is True, f'is_empty={is_empty}' + assert is_finished is True, f'is_finished={is_finished}' + + # Fetch from empty queue + request = await rq.fetch_next_request() + Actor.log.info(f'Fetch result from empty queue: {request}') + assert request is None, f'request={request}' + + # Check metadata for empty queue + metadata = await rq.get_metadata() + assert metadata is not None, f'metadata={metadata}' + Actor.log.info( + f'Empty queue metadata - Total: {metadata.total_request_count}, ' + f'Handled: {metadata.handled_request_count}, ' + f'Pending: {metadata.pending_request_count}' + ) + assert metadata.total_request_count == 0, f'metadata.total_request_count={metadata.total_request_count}' + assert metadata.handled_request_count == 0, f'metadata.handled_request_count={metadata.handled_request_count}' + assert metadata.pending_request_count == 0, f'metadata.pending_request_count={metadata.pending_request_count}' -async def test_large_batch_operations( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_large_batch_operations(request_queue_apify: RequestQueue) -> None: """Test handling large batches of requests.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Create a large batch of requests - large_batch = [f'https://example.com/large/{i}' for i in range(500)] - Actor.log.info(f'Created batch of {len(large_batch)} requests') + # Create a large batch of requests + large_batch = [f'https://example.com/large/{i}' for i in range(500)] + Actor.log.info(f'Created batch of {len(large_batch)} requests') - # Add in batch - await rq.add_requests(large_batch, batch_size=100, wait_for_all_requests_to_be_added=True) - Actor.log.info('Large batch add completed') + # Add in batch + await rq.add_requests(large_batch, batch_size=100, wait_for_all_requests_to_be_added=True) + Actor.log.info('Large batch add completed') - # Verify all requests were added - total_count = await rq.get_total_count() - assert total_count == 500, f'total_count={total_count}' + # Verify all requests were added + total_count = await rq.get_total_count() + assert total_count == 500, f'total_count={total_count}' - # Process all in chunks to test performance - processed_count = 0 + # Process all in chunks to test performance + processed_count = 0 - while not await rq.is_empty(): - request = await rq.fetch_next_request() + while not await rq.is_empty(): + request = await rq.fetch_next_request() - # The RQ is_empty should ensure we don't get None - assert request is not None, f'request={request}' + # The RQ is_empty should ensure we don't get None + assert request is not None, f'request={request}' - await rq.mark_request_as_handled(request) - processed_count += 1 + await rq.mark_request_as_handled(request) + processed_count += 1 - Actor.log.info(f'Processing completed. Total processed: {processed_count}') - assert processed_count == 500, f'processed_count={processed_count}' + Actor.log.info(f'Processing completed. Total processed: {processed_count}') + assert processed_count == 500, f'processed_count={processed_count}' - is_finished = await rq.is_finished() - assert is_finished is True, f'is_finished={is_finished}' + is_finished = await rq.is_finished() + assert is_finished is True, f'is_finished={is_finished}' - actor = await make_actor(label='rq-large-batch-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' - -async def test_mixed_string_and_request_objects( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_mixed_string_and_request_objects(request_queue_apify: RequestQueue) -> None: """Test adding both string URLs and Request objects.""" - async def main() -> None: - from apify import Request + rq = request_queue_apify + Actor.log.info('Request queue opened') - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') + # Add string URLs + await rq.add_request('https://example.com/string1') + await rq.add_request('https://example.com/string2') + Actor.log.info('Added string URL requests') - # Add string URLs - await rq.add_request('https://example.com/string1') - await rq.add_request('https://example.com/string2') - Actor.log.info('Added string URL requests') - - # Add Request objects - request_obj = Request.from_url('https://example.com/object1', user_data={'type': 'request_object'}) - await rq.add_request(request_obj) - Actor.log.info('Added Request object with user_data') - - # Add mixed batch - mixed_batch: list[str | Request] = [ - 'https://example.com/mixed1', - Request.from_url('https://example.com/mixed2', method='POST'), - 'https://example.com/mixed3', - ] - await rq.add_requests(mixed_batch) - Actor.log.info('Added mixed batch of strings and Request objects') + # Add Request objects + request_obj = Request.from_url('https://example.com/object1', user_data={'type': 'request_object'}) + await rq.add_request(request_obj) + Actor.log.info('Added Request object with user_data') - total_count = await rq.get_total_count() - Actor.log.info(f'Total requests in queue: {total_count}') + # Add mixed batch + mixed_batch: list[str | Request] = [ + 'https://example.com/mixed1', + Request.from_url('https://example.com/mixed2', method='POST'), + 'https://example.com/mixed3', + ] + await rq.add_requests(mixed_batch) + Actor.log.info('Added mixed batch of strings and Request objects') - # Fetch and verify all types work - fetched_requests = [] - while next_request := await rq.fetch_next_request(): - fetched_requests.append(next_request) - await rq.mark_request_as_handled(next_request) + total_count = await rq.get_total_count() + Actor.log.info(f'Total requests in queue: {total_count}') - assert len(fetched_requests) == 6, f'len(fetched_requests)={len(fetched_requests)}' + # Fetch and verify all types work + fetched_requests = [] + while next_request := await rq.fetch_next_request(): + fetched_requests.append(next_request) + await rq.mark_request_as_handled(next_request) - # Find the request object we added - request_obj_found = None - for req in fetched_requests: - if req.user_data and req.user_data.get('type') == 'request_object': - request_obj_found = req - break + assert len(fetched_requests) == 6, f'len(fetched_requests)={len(fetched_requests)}' - assert request_obj_found is not None, f'request_obj_found={request_obj_found}' - assert request_obj_found.url == 'https://example.com/object1', ( - f'request_obj_found.url={request_obj_found.url}' - ) - Actor.log.info('Mixed types verified - found request object with user_data') + # Find the request object we added + request_obj_found = None + for req in fetched_requests: + if req.user_data and req.user_data.get('type') == 'request_object': + request_obj_found = req + break - actor = await make_actor(label='rq-mixed-types-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + assert request_obj_found is not None, f'request_obj_found={request_obj_found}' + assert request_obj_found.url == 'https://example.com/object1', f'request_obj_found.url={request_obj_found.url}' + Actor.log.info('Mixed types verified - found request object with user_data') @pytest.mark.skip( @@ -841,193 +691,160 @@ async def worker() -> int: assert run_result.status == 'SUCCEEDED' -async def test_persistence_across_operations( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_persistence_across_operations(request_queue_apify: RequestQueue) -> None: """Test that queue state persists across different operations.""" - async def main() -> None: - async with Actor: - # Open queue and add some requests - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add initial batch - initial_requests = [f'https://example.com/persist/{i}' for i in range(10)] - await rq.add_requests(initial_requests, wait_for_all_requests_to_be_added=True) - Actor.log.info(f'Added initial batch of {len(initial_requests)} requests') - - initial_total = await rq.get_total_count() - Actor.log.info(f'Total count after initial batch: {initial_total}') - - # Process some requests - processed_count = 0 - for _ in range(5): - request = await rq.fetch_next_request() - if request: - await rq.mark_request_as_handled(request) - processed_count += 1 - - Actor.log.info(f'Processed {processed_count} requests from initial batch') - handled_after_first_batch = await rq.get_handled_count() - Actor.log.info(f'Handled count after processing: {handled_after_first_batch}') - - # Add more requests - additional_requests = [f'https://example.com/additional/{i}' for i in range(5)] - await rq.add_requests(additional_requests, wait_for_all_requests_to_be_added=True) - Actor.log.info(f'Added additional batch of {len(additional_requests)} requests') - - # Check final state - total_after_additional = await rq.get_total_count() - handled_after_additional = await rq.get_handled_count() - Actor.log.info( - f'After adding additional batch - Total: {total_after_additional}, Handled: {handled_after_additional}' - ) - assert total_after_additional == 15, f'total_after_additional={total_after_additional}' - assert handled_after_additional == 5, f'handled_after_additional={handled_after_additional}' - - # Process remaining - remaining_processed = 0 - while not await rq.is_finished(): - request = await rq.fetch_next_request() - if request: - remaining_processed += 1 - await rq.mark_request_as_handled(request) - else: - break - - Actor.log.info(f'Processed {remaining_processed} remaining requests') - is_finished = await rq.is_finished() - final_total = await rq.get_total_count() - final_handled = await rq.get_handled_count() - - Actor.log.info(f'Final state - Finished: {is_finished}, Total: {final_total}, Handled: {final_handled}') - assert is_finished is True, f'is_finished={is_finished}' - assert final_total == 15, f'final_total={final_total}' - assert final_handled == 15, f'final_handled={final_handled}' - - actor = await make_actor(label='rq-persistence-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' - - -async def test_request_deduplication_edge_cases( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: + # Open queue and add some requests + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Add initial batch + initial_requests = [f'https://example.com/persist/{i}' for i in range(10)] + await rq.add_requests(initial_requests, wait_for_all_requests_to_be_added=True) + Actor.log.info(f'Added initial batch of {len(initial_requests)} requests') + + initial_total = await rq.get_total_count() + Actor.log.info(f'Total count after initial batch: {initial_total}') + + # Process some requests + processed_count = 0 + for _ in range(5): + request = await rq.fetch_next_request() + if request: + await rq.mark_request_as_handled(request) + processed_count += 1 + + Actor.log.info(f'Processed {processed_count} requests from initial batch') + handled_after_first_batch = await rq.get_handled_count() + Actor.log.info(f'Handled count after processing: {handled_after_first_batch}') + + # Add more requests + additional_requests = [f'https://example.com/additional/{i}' for i in range(5)] + await rq.add_requests(additional_requests, wait_for_all_requests_to_be_added=True) + Actor.log.info(f'Added additional batch of {len(additional_requests)} requests') + + # Check final state + total_after_additional = await rq.get_total_count() + handled_after_additional = await rq.get_handled_count() + Actor.log.info( + f'After adding additional batch - Total: {total_after_additional}, Handled: {handled_after_additional}' + ) + assert total_after_additional == 15, f'total_after_additional={total_after_additional}' + assert handled_after_additional == 5, f'handled_after_additional={handled_after_additional}' + + # Process remaining + remaining_processed = 0 + while not await rq.is_finished(): + request = await rq.fetch_next_request() + if request: + remaining_processed += 1 + await rq.mark_request_as_handled(request) + else: + break + + Actor.log.info(f'Processed {remaining_processed} remaining requests') + is_finished = await rq.is_finished() + final_total = await rq.get_total_count() + final_handled = await rq.get_handled_count() + + Actor.log.info(f'Final state - Finished: {is_finished}, Total: {final_total}, Handled: {final_handled}') + assert is_finished is True, f'is_finished={is_finished}' + assert final_total == 15, f'final_total={final_total}' + assert final_handled == 15, f'final_handled={final_handled}' + + +async def test_request_deduplication_edge_cases(request_queue_apify: RequestQueue) -> None: """Test edge cases in request deduplication.""" - - async def main() -> None: - from apify import Request - - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Test URL normalization and deduplication with expected results - urls_and_deduplication_expectations: list[tuple[str | Request, bool]] = [ - ('https://example.com/page', False), - ('https://example.com/page/', True), # Should be deduplicated (same as first) - ('https://example.com/page?', True), # Should be deduplicated (same as first) - ( - Request.from_url('https://example.com/page#fragment', use_extended_unique_key=True), - False, - ), # Different extended unique key - ('https://example.com/page?param=1', False), # Different unique key - ] - Actor.log.info(f'Testing deduplication with {len(urls_and_deduplication_expectations)} URLs') - - results = list[bool]() - for url, expected_duplicate in urls_and_deduplication_expectations: - result = await rq.add_request(url) - results.append(result.was_already_present) - assert result.was_already_present == expected_duplicate, ( - f'url={url}', - f'expected_duplicate={expected_duplicate}', - f'actual_was_already_present={result.was_already_present}', - ) - - Actor.log.info(f'was_already_present results: {results}') - - # Calculate expected unique count - expected_unique_count = sum( - 1 for _, is_duplicate in urls_and_deduplication_expectations if not is_duplicate - ) - Actor.log.info(f'Expected {expected_unique_count} unique requests') - - # Fetch all unique requests - fetched_urls = list[str]() - while next_request := await rq.fetch_next_request(): - fetched_urls.append(next_request.url) - await rq.mark_request_as_handled(next_request) - - # Assert exact expected count - assert len(fetched_urls) == expected_unique_count, ( - f'len(fetched_urls)={len(fetched_urls)}', - f'expected_unique_count={expected_unique_count}', - ) - Actor.log.info( - f'Added {len(urls_and_deduplication_expectations)} URLs, ' - f'got {len(fetched_urls)} unique requests as expected' - ) - - actor = await make_actor(label='rq-deduplication-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Test URL normalization and deduplication with expected results + urls_and_deduplication_expectations: list[tuple[str | Request, bool]] = [ + ('https://example.com/page', False), + ('https://example.com/page/', True), # Should be deduplicated (same as first) + ('https://example.com/page?', True), # Should be deduplicated (same as first) + ( + Request.from_url('https://example.com/page#fragment', use_extended_unique_key=True), + False, + ), # Different extended unique key + ('https://example.com/page?param=1', False), # Different unique key + ] + Actor.log.info(f'Testing deduplication with {len(urls_and_deduplication_expectations)} URLs') + + results = list[bool]() + for url, expected_duplicate in urls_and_deduplication_expectations: + result = await rq.add_request(url) + results.append(result.was_already_present) + assert result.was_already_present == expected_duplicate, ( + f'url={url}', + f'expected_duplicate={expected_duplicate}', + f'actual_was_already_present={result.was_already_present}', + ) + + Actor.log.info(f'was_already_present results: {results}') + + # Calculate expected unique count + expected_unique_count = sum(1 for _, is_duplicate in urls_and_deduplication_expectations if not is_duplicate) + Actor.log.info(f'Expected {expected_unique_count} unique requests') + + # Fetch all unique requests + fetched_urls = list[str]() + while next_request := await rq.fetch_next_request(): + fetched_urls.append(next_request.url) + await rq.mark_request_as_handled(next_request) + + # Assert exact expected count + assert len(fetched_urls) == expected_unique_count, ( + f'len(fetched_urls)={len(fetched_urls)}', + f'expected_unique_count={expected_unique_count}', + ) + Actor.log.info( + f'Added {len(urls_and_deduplication_expectations)} URLs, got {len(fetched_urls)} unique requests as expected' + ) -async def test_request_ordering_with_mixed_operations( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_request_ordering_with_mixed_operations(request_queue_apify: RequestQueue) -> None: """Test request ordering with mixed add/reclaim operations.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Add initial requests - await rq.add_request('https://example.com/1') - await rq.add_request('https://example.com/2') - Actor.log.info('Added initial requests') + rq = request_queue_apify + Actor.log.info('Request queue opened') - # Fetch one and reclaim to forefront - request1 = await rq.fetch_next_request() - assert request1 is not None, f'request1={request1}' - assert request1.url == 'https://example.com/1', f'request1.url={request1.url}' - Actor.log.info(f'Fetched request: {request1.url}') + # Add initial requests + await rq.add_request('https://example.com/1') + await rq.add_request('https://example.com/2') + Actor.log.info('Added initial requests') - await rq.reclaim_request(request1, forefront=True) - Actor.log.info('Reclaimed request to forefront') + # Fetch one and reclaim to forefront + request1 = await rq.fetch_next_request() + assert request1 is not None, f'request1={request1}' + assert request1.url == 'https://example.com/1', f'request1.url={request1.url}' + Actor.log.info(f'Fetched request: {request1.url}') - # Add forefront request - await rq.add_request('https://example.com/priority', forefront=True) - Actor.log.info('Added new forefront request') + await rq.reclaim_request(request1, forefront=True) + Actor.log.info('Reclaimed request to forefront') - # Fetch all requests and verify forefront behavior - urls_ordered = list[str]() - while next_request := await rq.fetch_next_request(): - urls_ordered.append(next_request.url) - await rq.mark_request_as_handled(next_request) + # Add forefront request + await rq.add_request('https://example.com/priority', forefront=True) + Actor.log.info('Added new forefront request') - Actor.log.info(f'Final order of fetched URLs: {urls_ordered}') + # Fetch all requests and verify forefront behavior + urls_ordered = list[str]() + while next_request := await rq.fetch_next_request(): + urls_ordered.append(next_request.url) + await rq.mark_request_as_handled(next_request) - # Verify that we got all 3 requests - assert len(urls_ordered) == 3, f'len(urls_ordered)={len(urls_ordered)}' + Actor.log.info(f'Final order of fetched URLs: {urls_ordered}') - assert urls_ordered[0] == 'https://example.com/priority', f'urls_ordered[0]={urls_ordered[0]}' - assert urls_ordered[1] == request1.url, ( - f'urls_ordered[1]={urls_ordered[1]}', - f'request1.url={request1.url}', - ) - assert urls_ordered[2] == 'https://example.com/2', f'urls_ordered[2]={urls_ordered[2]}' - Actor.log.info('Request ordering verified successfully') + # Verify that we got all 3 requests + assert len(urls_ordered) == 3, f'len(urls_ordered)={len(urls_ordered)}' - actor = await make_actor(label='rq-mixed-ordering-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + assert urls_ordered[0] == 'https://example.com/priority', f'urls_ordered[0]={urls_ordered[0]}' + assert urls_ordered[1] == request1.url, ( + f'urls_ordered[1]={urls_ordered[1]}', + f'request1.url={request1.url}', + ) + assert urls_ordered[2] == 'https://example.com/2', f'urls_ordered[2]={urls_ordered[2]}' + Actor.log.info('Request ordering verified successfully') async def test_rq_isolation( @@ -1081,180 +898,161 @@ async def main() -> None: assert run_result.status == 'SUCCEEDED' -async def test_finished_state_accuracy( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: +async def test_finished_state_accuracy(request_queue_apify: RequestQueue) -> None: """Test accuracy of is_finished() method in various scenarios.""" - async def main() -> None: - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Initially should be finished - initial_finished = await rq.is_finished() - Actor.log.info(f'Initial finished state: {initial_finished}') - assert initial_finished is True, f'initial_finished={initial_finished}' - - # Add requests - should not be finished - await rq.add_request('https://example.com/test1') - await rq.add_request('https://example.com/test2') - after_add_finished = await rq.is_finished() - Actor.log.info(f'Finished state after adding requests: {after_add_finished}') - assert after_add_finished is False, f'after_add_finished={after_add_finished}' - - # Fetch but don't handle - should not be finished - request1 = await rq.fetch_next_request() - assert request1 is not None, f'request1={request1}' - after_fetch_finished = await rq.is_finished() - Actor.log.info(f'Finished state after fetch (not handled): {after_fetch_finished}') - assert after_fetch_finished is False, f'after_fetch_finished={after_fetch_finished}' - - # Reclaim request - should still not be finished - await rq.reclaim_request(request1) - after_reclaim_finished = await rq.is_finished() - Actor.log.info(f'Finished state after reclaim: {after_reclaim_finished}') - assert after_reclaim_finished is False, f'after_reclaim_finished={after_reclaim_finished}' - - # Handle all requests - should be finished - processed_count = 0 - while next_request := await rq.fetch_next_request(): - processed_count += 1 - await rq.mark_request_as_handled(next_request) - - Actor.log.info(f'Processed {processed_count} requests') - final_finished = await rq.is_finished() - assert final_finished is True, f'final_finished={final_finished}' - - actor = await make_actor(label='rq-finished-state-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' - - -async def test_operations_performance_pattern( - make_actor: MakeActorFunction, - run_actor: RunActorFunction, -) -> None: + rq = request_queue_apify + Actor.log.info('Request queue opened') + + # Initially should be finished + initial_finished = await rq.is_finished() + Actor.log.info(f'Initial finished state: {initial_finished}') + assert initial_finished is True, f'initial_finished={initial_finished}' + + # Add requests - should not be finished + await rq.add_request('https://example.com/test1') + await rq.add_request('https://example.com/test2') + after_add_finished = await rq.is_finished() + Actor.log.info(f'Finished state after adding requests: {after_add_finished}') + assert after_add_finished is False, f'after_add_finished={after_add_finished}' + + # Fetch but don't handle - should not be finished + request1 = await rq.fetch_next_request() + assert request1 is not None, f'request1={request1}' + after_fetch_finished = await rq.is_finished() + Actor.log.info(f'Finished state after fetch (not handled): {after_fetch_finished}') + assert after_fetch_finished is False, f'after_fetch_finished={after_fetch_finished}' + + # Reclaim request - should still not be finished + await rq.reclaim_request(request1) + after_reclaim_finished = await rq.is_finished() + Actor.log.info(f'Finished state after reclaim: {after_reclaim_finished}') + assert after_reclaim_finished is False, f'after_reclaim_finished={after_reclaim_finished}' + + # Handle all requests - should be finished + processed_count = 0 + while next_request := await rq.fetch_next_request(): + processed_count += 1 + await rq.mark_request_as_handled(next_request) + + Actor.log.info(f'Processed {processed_count} requests') + final_finished = await rq.is_finished() + assert final_finished is True, f'final_finished={final_finished}' + + +async def test_operations_performance_pattern(request_queue_apify: RequestQueue) -> None: """Test a common performance pattern: producer-consumer.""" + Actor.log.info('Request queue opened') + rq = request_queue_apify + + # Producer: Add requests in background + async def producer() -> None: + for i in range(20): + await rq.add_request(f'https://example.com/item/{i}') + if i % 5 == 0: # Add some delay to simulate real production + await asyncio.sleep(0.01) + Actor.log.info('Producer finished adding all 20 requests') + + # Consumer: Process requests as they become available + async def consumer() -> int: + processed = 0 + consecutive_empty = 0 + max_empty_attempts = 5 + + while consecutive_empty < max_empty_attempts: + request = await rq.fetch_next_request() + if request is None: + consecutive_empty += 1 + await asyncio.sleep(0.01) # Brief wait for more requests + continue - async def main() -> None: - import asyncio - - async with Actor: - rq = await Actor.open_request_queue() - Actor.log.info('Request queue opened') - - # Producer: Add requests in background - async def producer() -> None: - for i in range(20): - await rq.add_request(f'https://example.com/item/{i}') - if i % 5 == 0: # Add some delay to simulate real production - await asyncio.sleep(0.01) - Actor.log.info('Producer finished adding all 20 requests') - - # Consumer: Process requests as they become available - async def consumer() -> int: - processed = 0 - consecutive_empty = 0 - max_empty_attempts = 5 - - while consecutive_empty < max_empty_attempts: - request = await rq.fetch_next_request() - if request is None: - consecutive_empty += 1 - await asyncio.sleep(0.01) # Brief wait for more requests - continue + consecutive_empty = 0 + await rq.mark_request_as_handled(request) + processed += 1 - consecutive_empty = 0 - await rq.mark_request_as_handled(request) - processed += 1 + Actor.log.info(f'Consumer finished initial processing, processed {processed} requests') + return processed - Actor.log.info(f'Consumer finished initial processing, processed {processed} requests') - return processed + # Run producer and consumer concurrently + producer_task = asyncio.create_task(producer()) + consumer_task = asyncio.create_task(consumer()) - # Run producer and consumer concurrently - producer_task = asyncio.create_task(producer()) - consumer_task = asyncio.create_task(consumer()) + # Wait for both to complete + await producer_task + processed_count = await consumer_task + Actor.log.info(f'Concurrent phase completed, processed {processed_count} requests') - # Wait for both to complete - await producer_task - processed_count = await consumer_task - Actor.log.info(f'Concurrent phase completed, processed {processed_count} requests') + # Process any remaining requests + remaining_count = 0 + while next_request := await rq.fetch_next_request(): + await rq.mark_request_as_handled(next_request) + processed_count += 1 + remaining_count += 1 - # Process any remaining requests - remaining_count = 0 - while next_request := await rq.fetch_next_request(): - await rq.mark_request_as_handled(next_request) - processed_count += 1 - remaining_count += 1 + Actor.log.info(f'Processed {remaining_count} remaining requests') + Actor.log.info(f'Total processed: {processed_count} requests') + assert processed_count == 20, f'processed_count={processed_count}' - Actor.log.info(f'Processed {remaining_count} remaining requests') - Actor.log.info(f'Total processed: {processed_count} requests') - assert processed_count == 20, f'processed_count={processed_count}' - - final_finished = await rq.is_finished() - assert final_finished is True, f'final_finished={final_finished}' - - actor = await make_actor(label='rq-performance-pattern-test', main_func=main) - run_result = await run_actor(actor) - assert run_result.status == 'SUCCEEDED' + final_finished = await rq.is_finished() + assert final_finished is True, f'final_finished={final_finished}' async def test_request_queue_enhanced_metadata( - request_queue_force_cloud: RequestQueue, + request_queue_apify: RequestQueue, apify_client_async: ApifyClientAsync, ) -> None: """Test metadata tracking. Multiple clients scenarios are not guaranteed to give correct results without delay. But at least multiple clients, single producer, should be reliable on the producer side.""" - + rq = request_queue_apify for i in range(1, 10): - await request_queue_force_cloud.add_request(Request.from_url(f'http://example.com/{i}')) + await rq.add_request(Request.from_url(f'http://example.com/{i}')) # Reliable information as the API response is enhanced with local metadata estimation. - assert (await request_queue_force_cloud.get_metadata()).total_request_count == i + assert (await rq.get_metadata()).total_request_count == i # Accessed with client created explicitly with `client_key=None` should appear as distinct client - api_client = apify_client_async.request_queue(request_queue_id=request_queue_force_cloud.id, client_key=None) + api_client = apify_client_async.request_queue(request_queue_id=rq.id, client_key=None) await api_client.list_head() # The presence of another non-producing client should not affect the metadata for i in range(10, 20): - await request_queue_force_cloud.add_request(Request.from_url(f'http://example.com/{i}')) + await rq.add_request(Request.from_url(f'http://example.com/{i}')) # Reliable information as the API response is enhanced with local metadata estimation. - assert (await request_queue_force_cloud.get_metadata()).total_request_count == i + assert (await rq.get_metadata()).total_request_count == i async def test_request_queue_metadata_another_client( - request_queue_force_cloud: RequestQueue, + request_queue_apify: RequestQueue, apify_client_async: ApifyClientAsync, ) -> None: """Test metadata tracking. The delayed metadata should be reliable even when changed by another client.""" - api_client = apify_client_async.request_queue(request_queue_id=request_queue_force_cloud.id, client_key=None) + rq = request_queue_apify + api_client = apify_client_async.request_queue(request_queue_id=rq.id, client_key=None) await api_client.add_request(Request.from_url('http://example.com/1').model_dump(by_alias=True, exclude={'id'})) # Wait to be sure that the API has updated the global metadata await asyncio.sleep(10) - assert (await request_queue_force_cloud.get_metadata()).total_request_count == 1 + assert (await rq.get_metadata()).total_request_count == 1 async def test_request_queue_had_multiple_clients( - request_queue_force_cloud: RequestQueue, + request_queue_apify: RequestQueue, apify_client_async: ApifyClientAsync, ) -> None: """Test that `RequestQueue` correctly detects multiple clients. Clients created with different `client_key` should appear as distinct clients.""" - await request_queue_force_cloud.fetch_next_request() + rq = request_queue_apify + await rq.fetch_next_request() # Accessed with client created explicitly with `client_key=None` should appear as distinct client - api_client = apify_client_async.request_queue(request_queue_id=request_queue_force_cloud.id, client_key=None) + api_client = apify_client_async.request_queue(request_queue_id=request_queue_apify.id, client_key=None) await api_client.list_head() # Check that it is correctly in the RequestQueueClient metadata - assert (await request_queue_force_cloud.get_metadata()).had_multiple_clients is True + assert (await rq.get_metadata()).had_multiple_clients is True # Check that it is correctly in the API api_response = await api_client.get() @@ -1263,42 +1061,131 @@ async def test_request_queue_had_multiple_clients( async def test_request_queue_not_had_multiple_clients( - request_queue_force_cloud: RequestQueue, apify_client_async: ApifyClientAsync + request_queue_apify: RequestQueue, apify_client_async: ApifyClientAsync ) -> None: """Test that same `RequestQueue` created from Actor does not act as multiple clients.""" - + rq = request_queue_apify # Two calls to API to create situation where different `client_key` can set `had_multiple_clients` to True - await request_queue_force_cloud.fetch_next_request() - await request_queue_force_cloud.fetch_next_request() + await rq.fetch_next_request() + await rq.fetch_next_request() # Check that it is correctly in the RequestQueueClient metadata - assert (await request_queue_force_cloud.get_metadata()).had_multiple_clients is False + assert (await rq.get_metadata()).had_multiple_clients is False # Check that it is correctly in the API - api_client = apify_client_async.request_queue(request_queue_id=request_queue_force_cloud.id) + api_client = apify_client_async.request_queue(request_queue_id=rq.id) api_response = await api_client.get() assert api_response assert api_response['hadMultipleClients'] is False -async def test_request_queue_has_stats(request_queue_force_cloud: RequestQueue) -> None: - """Test that Apify based request queue has stats in metadata.""" +async def test_request_queue_simple_and_full_at_the_same_time( + apify_token: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Test using two variants of the ApifyStorageClient on the same queue resolves to the first client used.""" + monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_token) + + async with Actor: + rq_simple = await RequestQueue.open(storage_client=ApifyStorageClient(request_queue_access='single')) + rq_full = await RequestQueue.open(storage_client=ApifyStorageClient(request_queue_access='shared')) + # Opening same queue again with different ApifyStorageClient will resolve to the first client used. + assert rq_simple is rq_full + await rq_simple.drop() + +@pytest.mark.parametrize( + ('access', 'expected_write_count_per_request'), + [pytest.param('single', 2, id='Simple rq client'), pytest.param('shared', 3, id='Full rq client')], +) +async def test_crawler_run_request_queue_variant_stats( + *, + apify_token: str, + monkeypatch: pytest.MonkeyPatch, + access: Literal['single', 'shared'], + expected_write_count_per_request: int, +) -> None: + """Check the main difference in the simple vs full request queue client - writeCount per request. + + The simple client also has lower readCount, but the costs of read are order of magnitude cheaper than writes, so we + do test that. + """ + monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_token) + async with Actor: + requests = 5 + rq = await RequestQueue.open(storage_client=ApifyStorageClient(request_queue_access=access)) + crawler = BasicCrawler(request_manager=rq) + + @crawler.router.default_handler + async def default_handler(context: BasicCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url} ...') + + await crawler.run([Request.from_url(f'https://example.com/{i}') for i in range(requests)]) + + # Make sure all requests were handled. + assert crawler.statistics.state.requests_finished == requests + + # Check the request queue stats + await asyncio.sleep(10) # Wait to be sure that metadata are updated + + metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata()) + Actor.log.info(f'{metadata.stats=}') + assert metadata.stats.write_count == requests * expected_write_count_per_request + await rq.drop() + + +async def test_cache_initialization(apify_token: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Test that Apify based simple `RequestQueue` initializes cache correctly to reduce unnecessary API calls.""" + + # Create an instance of the Apify request queue on the platform and drop it when the test is finished. + request_queue_name = generate_unique_resource_name('request_queue') + monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_token) + + requests = [Request.from_url(f'http://example.com/{i}', handled_at=datetime.now(timezone.utc)) for i in range(10)] + + async with Actor: + rq = await Actor.open_request_queue(name=request_queue_name, force_cloud=True) + try: + await rq.add_requests(requests) + + # Check that it is correctly in the API + await asyncio.sleep(10) # Wait to be sure that metadata are updated + metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata()) + stats_before = metadata.stats + Actor.log.info(stats_before) + + # Clear service locator cache to simulate creating RQ instance from scratch + service_locator.storage_instance_manager.clear_cache() + + # Try to enqueue same requests again. It should be deduplicated from local cache created on initialization + rq = await Actor.open_request_queue(name=request_queue_name, force_cloud=True) + await rq.add_requests(requests) + + await asyncio.sleep(10) # Wait to be sure that metadata are updated + metadata = cast('ApifyRequestQueueMetadata', await rq.get_metadata()) + stats_after = metadata.stats + Actor.log.info(stats_after) + + # Cache was actually initialized, readCount increased + assert (stats_after.read_count - stats_before.read_count) == len(requests) + # Deduplication happened locally, writeCount should be the same + assert stats_after.write_count == stats_before.write_count + + finally: + await rq.drop() + + +async def test_request_queue_has_stats(request_queue_apify: RequestQueue) -> None: + """Test that Apify based request queue has stats in metadata.""" + rq = request_queue_apify add_request_count = 3 - read_request_count = 2 - await request_queue_force_cloud.add_requests( - [Request.from_url(f'http://example.com/{i}') for i in range(add_request_count)] - ) - for _ in range(read_request_count): - await request_queue_force_cloud.get_request(Request.from_url('http://example.com/1').unique_key) + await rq.add_requests([Request.from_url(f'http://example.com/{i}') for i in range(add_request_count)]) # Wait for stats to become stable await asyncio.sleep(10) - metadata = await request_queue_force_cloud.get_metadata() + metadata = await rq.get_metadata() assert hasattr(metadata, 'stats') apify_metadata = cast('ApifyRequestQueueMetadata', metadata) - assert apify_metadata.stats.read_count == read_request_count assert apify_metadata.stats.write_count == add_request_count diff --git a/tests/unit/actor/test_actor_log.py b/tests/unit/actor/test_actor_log.py index 356f8bb3..ecb90ab6 100644 --- a/tests/unit/actor/test_actor_log.py +++ b/tests/unit/actor/test_actor_log.py @@ -37,7 +37,7 @@ async def test_actor_logs_messages_correctly(caplog: pytest.LogCaptureFixture) - raise RuntimeError('Dummy RuntimeError') # Updated expected number of log records (an extra record is now captured) - assert len(caplog.records) == 14 + assert len(caplog.records) == 15 # Record 0: Extra Pytest context log assert caplog.records[0].levelno == logging.DEBUG @@ -51,54 +51,58 @@ async def test_actor_logs_messages_correctly(caplog: pytest.LogCaptureFixture) - assert caplog.records[2].levelno == logging.INFO assert caplog.records[2].message == 'Initializing Actor...' + # Record 2: Initializing Actor... + assert caplog.records[3].levelno == logging.DEBUG + assert caplog.records[3].message.startswith('Storage client set to') + # Record 3: System info - assert caplog.records[3].levelno == logging.INFO - assert caplog.records[3].message == 'System info' + assert caplog.records[4].levelno == logging.INFO + assert caplog.records[4].message == 'System info' # Record 4: Event manager initialized - assert caplog.records[4].levelno == logging.DEBUG - assert caplog.records[4].message == 'Event manager initialized' + assert caplog.records[5].levelno == logging.DEBUG + assert caplog.records[5].message == 'Event manager initialized' # Record 5: Charging manager initialized - assert caplog.records[5].levelno == logging.DEBUG - assert caplog.records[5].message == 'Charging manager initialized' + assert caplog.records[6].levelno == logging.DEBUG + assert caplog.records[6].message == 'Charging manager initialized' # Record 6: Debug message - assert caplog.records[6].levelno == logging.DEBUG - assert caplog.records[6].message == 'Debug message' + assert caplog.records[7].levelno == logging.DEBUG + assert caplog.records[7].message == 'Debug message' # Record 7: Info message - assert caplog.records[7].levelno == logging.INFO - assert caplog.records[7].message == 'Info message' + assert caplog.records[8].levelno == logging.INFO + assert caplog.records[8].message == 'Info message' # Record 8: Warning message - assert caplog.records[8].levelno == logging.WARNING - assert caplog.records[8].message == 'Warning message' + assert caplog.records[9].levelno == logging.WARNING + assert caplog.records[9].message == 'Warning message' # Record 9: Error message - assert caplog.records[9].levelno == logging.ERROR - assert caplog.records[9].message == 'Error message' + assert caplog.records[10].levelno == logging.ERROR + assert caplog.records[10].message == 'Error message' # Record 10: Exception message with traceback (ValueError) - assert caplog.records[10].levelno == logging.ERROR - assert caplog.records[10].message == 'Exception message' - assert caplog.records[10].exc_info is not None - assert caplog.records[10].exc_info[0] is ValueError - assert isinstance(caplog.records[10].exc_info[1], ValueError) - assert str(caplog.records[10].exc_info[1]) == 'Dummy ValueError' + assert caplog.records[11].levelno == logging.ERROR + assert caplog.records[11].message == 'Exception message' + assert caplog.records[11].exc_info is not None + assert caplog.records[11].exc_info[0] is ValueError + assert isinstance(caplog.records[11].exc_info[1], ValueError) + assert str(caplog.records[11].exc_info[1]) == 'Dummy ValueError' # Record 11: Multiline log message - assert caplog.records[11].levelno == logging.INFO - assert caplog.records[11].message == 'Multi\nline\nlog\nmessage' + assert caplog.records[12].levelno == logging.INFO + assert caplog.records[12].message == 'Multi\nline\nlog\nmessage' # Record 12: Actor failed with an exception (RuntimeError) - assert caplog.records[12].levelno == logging.ERROR - assert caplog.records[12].message == 'Actor failed with an exception' - assert caplog.records[12].exc_info is not None - assert caplog.records[12].exc_info[0] is RuntimeError - assert isinstance(caplog.records[12].exc_info[1], RuntimeError) - assert str(caplog.records[12].exc_info[1]) == 'Dummy RuntimeError' + assert caplog.records[13].levelno == logging.ERROR + assert caplog.records[13].message == 'Actor failed with an exception' + assert caplog.records[13].exc_info is not None + assert caplog.records[13].exc_info[0] is RuntimeError + assert isinstance(caplog.records[13].exc_info[1], RuntimeError) + assert str(caplog.records[13].exc_info[1]) == 'Dummy RuntimeError' # Record 13: Exiting Actor - assert caplog.records[13].levelno == logging.INFO - assert caplog.records[13].message == 'Exiting Actor' + assert caplog.records[14].levelno == logging.INFO + assert caplog.records[14].message == 'Exiting Actor' diff --git a/tests/unit/actor/test_configuration.py b/tests/unit/actor/test_configuration.py index 7212f5e3..97500eab 100644 --- a/tests/unit/actor/test_configuration.py +++ b/tests/unit/actor/test_configuration.py @@ -7,11 +7,10 @@ from crawlee.configuration import Configuration as CrawleeConfiguration from crawlee.crawlers import BasicCrawler from crawlee.errors import ServiceConflictError -from crawlee.storage_clients import FileSystemStorageClient from apify import Actor from apify import Configuration as ApifyConfiguration -from apify.storage_clients import FileSystemStorageClient as ApifyFileSystemStorageClient +from apify.storage_clients._smart_apify._storage_client import SmartApifyStorageClient @pytest.mark.parametrize( @@ -111,8 +110,8 @@ async def test_crawler_implicit_local_storage() -> None: async with Actor(): crawler = BasicCrawler() - assert isinstance(service_locator.get_storage_client(), ApifyFileSystemStorageClient) - assert isinstance(crawler._service_locator.get_storage_client(), ApifyFileSystemStorageClient) + assert isinstance(service_locator.get_storage_client(), SmartApifyStorageClient) + assert isinstance(crawler._service_locator.get_storage_client(), SmartApifyStorageClient) async def test_crawlers_own_configuration(tmp_path: Path) -> None: @@ -242,16 +241,3 @@ def test_apify_configuration_is_always_used(caplog: pytest.LogCaptureFixture) -> 'It is recommended to set `apify.Configuration` explicitly as early as possible by using ' 'service_locator.set_configuration' ) in caplog.messages - - -async def test_file_system_storage_client_warning(caplog: pytest.LogCaptureFixture) -> None: - service_locator.set_storage_client(FileSystemStorageClient()) - caplog.set_level('WARNING') - async with Actor(): - ... - - assert ( - 'Using crawlee.storage_clients._file_system._storage_client.FileSystemStorageClient in Actor context is not ' - 'recommended and can lead to problems with reading the input file. Use ' - '`apify.storage_clients.FileSystemStorageClient` instead.' - ) in caplog.messages diff --git a/tests/unit/storage_clients/test_apify_request_queue_client.py b/tests/unit/storage_clients/test_apify_request_queue_client.py index 019b2e0b..f00b2d3a 100644 --- a/tests/unit/storage_clients/test_apify_request_queue_client.py +++ b/tests/unit/storage_clients/test_apify_request_queue_client.py @@ -1,6 +1,6 @@ import pytest -from apify.storage_clients._apify._request_queue_client import unique_key_to_request_id +from apify.storage_clients._apify._utils import unique_key_to_request_id def test_unique_key_to_request_id_length() -> None: diff --git a/uv.lock b/uv.lock index 6a71bfa1..5f826e54 100644 --- a/uv.lock +++ b/uv.lock @@ -76,7 +76,7 @@ requires-dist = [ { name = "apify-client", specifier = ">=2.0.0,<3.0.0" }, { name = "apify-shared", specifier = ">=2.0.0,<3.0.0" }, { name = "cachetools", specifier = ">=5.5.0" }, - { name = "crawlee", specifier = "==0.6.13b42" }, + { name = "crawlee", specifier = "==0.6.13b46" }, { name = "cryptography", specifier = ">=42.0.0" }, { name = "impit", specifier = ">=0.6.1" }, { name = "lazy-object-proxy", specifier = ">=1.11.0" }, @@ -516,7 +516,7 @@ toml = [ [[package]] name = "crawlee" -version = "0.6.13b42" +version = "0.6.13b46" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -532,9 +532,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "yarl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/98/8e/8c5bf3cd84335aeb157f95ecaadc5cb61b9bb0f1ffa28a50f9a2485c38a6/crawlee-0.6.13b42.tar.gz", hash = "sha256:5a8c7bcf6abf77c6b7be3323e3cfa017a9717f0b5e5275bbb7ad8de589c851af", size = 24842767, upload-time = "2025-09-17T15:19:26.706Z" } +sdist = { url = "https://files.pythonhosted.org/packages/94/5d/f42c7684b0120eaaf0fd0b7667e6222ab0e0bed2c197a348ad6b534061e8/crawlee-0.6.13b46.tar.gz", hash = "sha256:a1ba1fd649c5673801b85c7b3035c8288a8f783a25e6980961d87d08d55701d4", size = 24846309, upload-time = "2025-09-25T06:40:30.654Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/eb/6a048e5916a78c30ea1b550452a6ede24facf5cafd564bbb1bc5e8ba6fea/crawlee-0.6.13b42-py3-none-any.whl", hash = "sha256:e9c258d49c8d4269d41a1dd9babfc262d241c62c9549d4dd54d1cad0ddbf9569", size = 279764, upload-time = "2025-09-17T15:19:23.817Z" }, + { url = "https://files.pythonhosted.org/packages/ae/a8/1210d96728108c0ff1d1cd935fb41998768265a9e1a03265f2476c2d734f/crawlee-0.6.13b46-py3-none-any.whl", hash = "sha256:48788749a861024fa21eba114c1c209e34a321bc5475e0ce38c493ada333f785", size = 280069, upload-time = "2025-09-25T06:40:28.272Z" }, ] [package.optional-dependencies]