Use factory in Rq and Dataset

Author: Pijukatel

src/apify/storage_clients/_apify/_dataset_client.py

@@ -7,18 +7,19 @@
 
 from typing_extensions import override
 
+from apify_client.clients import DatasetClientAsync
 from crawlee._utils.byte_size import ByteSize
 from crawlee._utils.file import json_dumps
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 from crawlee.storages import Dataset
 
-from ._utils import AliasResolver, create_apify_client
+from ._utils import ApiClientFactory
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
-    from apify_client.clients import DatasetClientAsync
+    from apify_client.clients import DatasetCollectionClientAsync
     from crawlee._types import JsonSerializable
 
     from apify import Configuration
@@ -101,66 +102,12 @@ async def open(
                 `id`, `name`, or `alias` is provided, or if none are provided and no default storage ID is available
                 in the configuration.
         """
-        if sum(1 for param in [id, name, alias] if param is not None) > 1:
-            raise ValueError('Only one of "id", "name", or "alias" can be specified, not multiple.')
-
-        apify_client_async = create_apify_client(configuration)
-        apify_datasets_client = apify_client_async.datasets()
-
-        # Normalize unnamed default storage in cases where not defined in `configuration.default_dataset_id` to unnamed
-        # storage aliased as `__default__`
-        if not any([alias, name, id, configuration.default_dataset_id]):
-            alias = '__default__'
-
-        if alias:
-            # Check if there is pre-existing alias mapping in the default KVS.
-            async with AliasResolver(storage_type=Dataset, alias=alias, configuration=configuration) as _alias:
-                id = await _alias.resolve_id()
-
-                # There was no pre-existing alias in the mapping.
-                # Create a new unnamed storage and store the mapping.
-                if id is None:
-                    new_storage_metadata = DatasetMetadata.model_validate(
-                        await apify_datasets_client.get_or_create(),
-                    )
-                    id = new_storage_metadata.id
-                    await _alias.store_mapping(storage_id=id)
-
-        # If name is provided, get or create the storage by name.
-        elif name:
-            id = DatasetMetadata.model_validate(
-                await apify_datasets_client.get_or_create(name=name),
-            ).id
-
-        # If none are provided, try to get the default storage ID from environment variables.
-        elif id is None:
-            id = configuration.default_dataset_id
-            if not id:
-                raise ValueError(
-                    'Dataset "id", "name", or "alias" must be specified, '
-                    'or a default dataset ID must be set in the configuration.'
-                )
-
-        # Now create the client for the determined ID
-        apify_dataset_client = apify_client_async.dataset(dataset_id=id)
-
-        # Fetch its metadata.
-        metadata = await apify_dataset_client.get()
-
-        # If metadata is None, it means the storage does not exist, so we create it.
-        if metadata is None:
-            id = DatasetMetadata.model_validate(
-                await apify_datasets_client.get_or_create(),
-            ).id
-            apify_dataset_client = apify_client_async.dataset(dataset_id=id)
-
-        # Verify that the storage exists by fetching its metadata again.
-        metadata = await apify_dataset_client.get()
-        if metadata is None:
-            raise ValueError(f'Opening dataset with id={id}, name={name}, and alias={alias} failed.')
+        api_client, _ = await DatasetApiClientFactory(
+            configuration=configuration, alias=alias, name=name, id=id
+        ).get_client_with_metadata()
 
         return cls(
-            api_client=apify_dataset_client,
+            api_client=api_client,
             api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
             lock=asyncio.Lock(),
         )
@@ -309,3 +256,24 @@ async def _chunk_by_size(self, items: AsyncIterator[str]) -> AsyncIterator[str]:
                 last_chunk_size = payload_size + ByteSize(2)  # Add 2 bytes for [] wrapper.
 
         yield f'[{",".join(current_chunk)}]'
+
+
+class DatasetApiClientFactory(ApiClientFactory[DatasetClientAsync, DatasetMetadata]):
+    @property
+    def _collection_client(self) -> DatasetCollectionClientAsync:
+        return self._api_client.datasets()
+
+    def _get_resource_client(self, id: str) -> DatasetClientAsync:
+        return self._api_client.dataset(dataset_id=id)
+
+    @property
+    def _default_id(self) -> str | None:
+        return self._configuration.default_dataset_id
+
+    @property
+    def _storage_type(self) -> type[Dataset]:
+        return Dataset
+
+    @staticmethod
+    def _get_metadata(raw_metadata: dict | None) -> DatasetMetadata:
+        return DatasetMetadata.model_validate(raw_metadata)

src/apify/storage_clients/_apify/_key_value_store_client.py

@@ -29,27 +29,6 @@
 logger = getLogger(__name__)
 
 
-class KvsApiClientFactory(ApiClientFactory[KeyValueStoreClientAsync, ApifyKeyValueStoreMetadata]):
-    @property
-    def _collection_client(self) -> KeyValueStoreCollectionClientAsync:
-        return self._api_client.key_value_stores()
-
-    def _get_resource_client(self, id: str) -> KeyValueStoreClientAsync:
-        return self._api_client.key_value_store(key_value_store_id=id)
-
-    @property
-    def _default_id(self) -> str | None:
-        return self._configuration.default_key_value_store_id
-
-    @property
-    def _storage_type(self) -> type[KeyValueStore]:
-        return KeyValueStore
-
-    @staticmethod
-    def _get_metadata(raw_metadata: dict | None) -> ApifyKeyValueStoreMetadata:
-        return ApifyKeyValueStoreMetadata.model_validate(raw_metadata)
-
-
 class ApifyKeyValueStoreClient(KeyValueStoreClient):
     """An Apify platform implementation of the key-value store client."""
 
@@ -116,9 +95,9 @@ async def open(
                 `id`, `name`, or `alias` is provided, or if none are provided and no default storage ID is available
                 in the configuration.
         """
-        api_client, _ =await KvsApiClientFactory(
-                configuration=configuration, alias=alias, name=name, id=id
-            ).get_client_with_metadata()
+        api_client, _ = await KvsApiClientFactory(
+            configuration=configuration, alias=alias, name=name, id=id
+        ).get_client_with_metadata()
         return cls(
             api_client=api_client,
             api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
@@ -203,3 +182,24 @@ async def get_public_url(self, key: str) -> str:
             A public URL that can be used to access the value of the given key in the KVS.
         """
         return await self._api_client.get_record_public_url(key=key)
+
+
+class KvsApiClientFactory(ApiClientFactory[KeyValueStoreClientAsync, ApifyKeyValueStoreMetadata]):
+    @property
+    def _collection_client(self) -> KeyValueStoreCollectionClientAsync:
+        return self._api_client.key_value_stores()
+
+    def _get_resource_client(self, id: str) -> KeyValueStoreClientAsync:
+        return self._api_client.key_value_store(key_value_store_id=id)
+
+    @property
+    def _default_id(self) -> str | None:
+        return self._configuration.default_key_value_store_id
+
+    @property
+    def _storage_type(self) -> type[KeyValueStore]:
+        return KeyValueStore
+
+    @staticmethod
+    def _get_metadata(raw_metadata: dict | None) -> ApifyKeyValueStoreMetadata:
+        return ApifyKeyValueStoreMetadata.model_validate(raw_metadata)

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -5,16 +5,19 @@
 
 from typing_extensions import override
 
+from apify_client.clients import RequestQueueClientAsync
 from crawlee.storage_clients._base import RequestQueueClient
+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 ApiClientFactory
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from apify_client.clients import RequestQueueClientAsync
+    from apify_client.clients import RequestQueueCollectionClientAsync
     from crawlee import Request
     from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
@@ -221,13 +224,11 @@ async def open(
                 `id`, `name`, or `alias` is provided, or if none are provided and no default storage ID is available
                 in the configuration.
         """
-        _api_client, metadata =await RqApiClientFactory(
-                configuration=configuration, alias=alias, name=name, id=id
-            ).get_client_with_metadata()
+        _api_client, metadata = await RqApiClientFactory(
+            configuration=configuration, alias=alias, name=name, id=id
+        ).get_client_with_metadata()
         return cls(
-            api_client=await RqApiClientFactory(
-                configuration=configuration, alias=alias, name=name, id=id
-            ).create_api_client(),
+            api_client=_api_client,
             metadata=metadata,
             access=access,
         )
@@ -242,3 +243,24 @@ async def purge(self) -> None:
     @override
     async def drop(self) -> None:
         await self._api_client.delete()
+
+
+class RqApiClientFactory(ApiClientFactory[RequestQueueClientAsync, ApifyRequestQueueMetadata]):
+    @property
+    def _collection_client(self) -> RequestQueueCollectionClientAsync:
+        return self._api_client.request_queues()
+
+    def _get_resource_client(self, id: str) -> RequestQueueClientAsync:
+        return self._api_client.request_queue(request_queue_id=id)
+
+    @property
+    def _default_id(self) -> str | None:
+        return self._configuration.default_request_queue_id
+
+    @property
+    def _storage_type(self) -> type[RequestQueue]:
+        return RequestQueue
+
+    @staticmethod
+    def _get_metadata(raw_metadata: dict | None) -> ApifyRequestQueueMetadata:
+        return ApifyRequestQueueMetadata.model_validate(raw_metadata)

src/apify/storage_clients/_apify/_utils.py

@@ -42,9 +42,8 @@
 TResourceClient = TypeVar(
     'TResourceClient', bound=KeyValueStoreClientAsync | RequestQueueClientAsync | DatasetClientAsync
 )
-TStorageMetadata = TypeVar(
-    'TStorageMetadata', bound=KeyValueStoreMetadata | RequestQueueMetadata | DatasetMetadata
-)
+TStorageMetadata = TypeVar('TStorageMetadata', bound=KeyValueStoreMetadata | RequestQueueMetadata | DatasetMetadata)
+
 
 class AliasResolver:
     """Class for handling aliases.
@@ -245,26 +244,24 @@ def __init__(self, configuration: Configuration, alias: str | None, name: str |
         if sum(1 for param in [id, name, alias] if param is not None) > 1:
             raise ValueError('Only one of "id", "name", or "alias" can be specified, not multiple.')
 
-        self.alias = alias
-        self.name = name
-        self.id = id
+        self._alias = alias
+        self._name = name
+        self._id = id
         self._configuration = configuration
         self._api_client = create_apify_client(configuration)
 
     async def get_client_with_metadata(self) -> tuple[TResourceClient, TStorageMetadata]:
-        match (self.alias, self.name, self.id, self._default_id):
+        match (self._alias, self._name, self._id, self._default_id):
             case (None, None, None, None):
                 # Normalize unnamed default storage in cases where not defined in `self._default_id` to
                 # unnamed storage aliased as `__default__`. Used only when running locally.
                 return await self._open_by_alias('__default__')
 
             case (str(), None, None, _):
-                return await self._open_by_alias(self.alias)
+                return await self._open_by_alias(self._alias)
 
             case (None, None, None, str()):
-                # Now create the client for the determined ID
                 resource_client = self._get_resource_client(id=self._default_id)
-                # Fetch its metadata.
                 raw_metadata = await resource_client.get()
                 metadata = self._get_metadata(raw_metadata)
                 if not raw_metadata:
@@ -276,22 +273,20 @@ async def get_client_with_metadata(self) -> tuple[TResourceClient, TStorageMetad
                 return resource_client, metadata
 
             case (None, str(), None, _):
-                metadata = self._get_metadata(
-                    await self._collection_client.get_or_create(name=self.name))
+                metadata = self._get_metadata(await self._collection_client.get_or_create(name=self._name))
                 # Freshly fetched named storage. No need to fetch it again.
                 return self._get_resource_client(id=metadata.id), metadata
 
             case (None, None, str(), _):
-                # Now create the client for the determined ID
-                resource_client = self._get_resource_client(id=self.id)
+                resource_client = self._get_resource_client(id=self._id)
                 # Fetch its metadata.
                 raw_metadata = await resource_client.get()
                 # If metadata is None, it means the storage does not exist.
                 if raw_metadata is None:
-                    raise ValueError(f'Opening key-value store with id={self.id} failed.')
+                    raise ValueError(f'Opening {self._storage_type} with id={self._id} failed.')
                 return resource_client, self._get_metadata(raw_metadata)
 
-        raise RuntimeError('Will never happen')
+        raise RuntimeError('Unreachable code')
 
     @property
     @abstractmethod