Use composition instead of inheritance

Author: Pijukatel

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -1,53 +1,32 @@
 from __future__ import annotations
 
-import re
-from base64 import b64encode
-from hashlib import sha256
 from logging import getLogger
-from typing import TYPE_CHECKING, Final
+from typing import TYPE_CHECKING, Final, Literal
 
 from typing_extensions import override
 
 from apify_client import ApifyClientAsync
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee.storage_clients._base import RequestQueueClient
-from crawlee.storage_clients.models import RequestQueueMetadata
+from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 from crawlee.storages import RequestQueue
 
 from ._models import ApifyRequestQueueMetadata, RequestQueueStats
+from ._request_queue_shared_client import _ApifyRequestQueueSharedClient
+from ._request_queue_single_client import _ApifyRequestQueueSingleClient
 from ._utils import AliasResolver
 
 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):
     """Base class for Apify platform implementations of the request queue client."""
 
@@ -59,6 +38,7 @@ def __init__(
         *,
         api_client: RequestQueueClientAsync,
         metadata: RequestQueueMetadata,
+        access: Literal['single', 'shared'] = 'single',
     ) -> None:
         """Initialize a new instance.
 
@@ -67,8 +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.
+
+        Returns:
+            The request or `None` if there are no more pending requests.
+        """
+        return await self._implementation.fetch_next_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.
+        """
+        return await self._implementation.mark_request_as_handled(request)
+
+    @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.
+        """
+        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:
@@ -103,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.
 
@@ -120,6 +205,18 @@ 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.
@@ -217,10 +314,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:

src/apify/storage_clients/_apify/_request_queue_shared_client.py

@@ -4,28 +4,26 @@
 from collections import deque
 from datetime import datetime, timedelta, timezone
 from logging import getLogger
-from typing import TYPE_CHECKING, Final
+from typing import TYPE_CHECKING, Any, Final
 
 from cachetools import LRUCache
-from typing_extensions import override
 
 from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
-from . import ApifyRequestQueueClient
-from ._models import CachedRequest, RequestQueueHead
-from ._request_queue_client import unique_key_to_request_id
+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 Sequence
+    from collections.abc import Callable, Coroutine, Sequence
 
     from apify_client.clients import RequestQueueClientAsync
 
 
 logger = getLogger(__name__)
 
 
-class ApifyRequestQueueSharedClient(ApifyRequestQueueClient):
+class _ApifyRequestQueueSharedClient:
     """An Apify platform implementation of the request queue client.
 
     This implementation supports multiple producers and multiple consumers scenario.
@@ -39,21 +37,26 @@ def __init__(
         *,
         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._metadata = metadata
-        """Additional data related to the RequestQueue."""
-
         self._queue_head = deque[str]()
         """A deque to store request unique keys in the queue head."""
 
-        self._requests_cache: LRUCache[str, CachedRequest] = LRUCache(maxsize=self._MAX_CACHED_REQUESTS)
+        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
@@ -72,12 +75,11 @@ async def _get_metadata_estimate(self) -> RequestQueueMetadata:
         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.get_metadata()
+        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
+        return self.metadata
 
-    @override
     async def add_batch_of_requests(
         self,
         requests: Sequence[Request],
@@ -167,11 +169,10 @@ async def add_batch_of_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.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.
 
@@ -188,7 +189,6 @@ async def get_request(self, unique_key: str) -> Request | 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.
 
@@ -240,7 +240,6 @@ async def fetch_next_request(self) -> Request | None:
 
         return request
 
-    @override
     async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
         """Mark a request as handled after successful processing.
 
@@ -265,7 +264,7 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
 
             # 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.handled_request_count += 1
 
             # Update the cache with the handled request
             cache_key = request.unique_key
@@ -280,7 +279,6 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
         else:
             return processed_request
 
-    @override
     async def reclaim_request(
         self,
         request: Request,
@@ -313,7 +311,7 @@ async def reclaim_request(
                 # 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.handled_request_count -= 1
 
                 # Update the cache
                 cache_key = request.unique_key
@@ -334,7 +332,6 @@ async def reclaim_request(
             else:
                 return processed_request
 
-    @override
     async def is_empty(self) -> bool:
         """Check if the queue is empty.
 
@@ -472,7 +469,7 @@ async def _list_head(
         # 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)
+        self.metadata.had_multiple_clients = response.get('hadMultipleClients', False)
 
         for request_data in response.get('items', []):
             request = Request.model_validate(request_data)