Finished KVS, TODO: RQ and Dataset

Author: Pijukatel

src/apify/storage_clients/_apify/_key_value_store_client.py

@@ -7,23 +7,49 @@
 
 from typing_extensions import override
 
+from apify_client.clients import (
+    KeyValueStoreClientAsync,
+)
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreRecord, KeyValueStoreRecordMetadata
 from crawlee.storages import KeyValueStore
 
 from ._models import ApifyKeyValueStoreMetadata, KeyValueStoreListKeysPage
-from ._utils import AliasResolver, create_apify_client
+from ._utils import ApiClientFactory
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
-    from apify_client.clients import KeyValueStoreClientAsync
+    from apify_client.clients import (
+        KeyValueStoreCollectionClientAsync,
+    )
 
     from apify import Configuration
 
 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."""
 
@@ -90,91 +116,14 @@ 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_kvss_client = apify_client_async.key_value_stores()
-
-        # Normalize unnamed default storage in cases where not defined in `configuration.default_key_value_store_id` to
-        # unnamed storage aliased as `__default__`
-        if not any([alias, name, id, configuration.default_key_value_store_id]):
-            alias = '__default__'
-
-        if alias:
-            # Check if there is pre-existing alias mapping in the default KVS.
-            async with AliasResolver(storage_type=KeyValueStore, alias=alias, configuration=configuration) as _alias:
-                id = await _alias.resolve_id()
-
-                if id:
-                    # There was id, storage has to exist, fetch metadata to confirm it.
-                    apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=id)
-                    raw_metadata = await apify_kvs_client.get()
-                    if raw_metadata:
-                        return cls(
-                            api_client=apify_kvs_client,
-                            api_public_base_url='',
-                            # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
-                            lock=asyncio.Lock(),
-                        )
-
-                # There was no pre-existing alias in the mapping or the id did not point to existing storage.
-                # Create a new unnamed storage and store the alias mapping.
-                metadata = ApifyKeyValueStoreMetadata.model_validate(
-                    await apify_kvss_client.get_or_create(),
-                )
-                await _alias.store_mapping(storage_id=metadata.id)
-
-                # Return the client for the newly created storage directly.
-                # It was just created, no need to refetch it.
-                return cls(
-                    api_client=apify_client_async.key_value_store(key_value_store_id=metadata.id),
-                    api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
-                    lock=asyncio.Lock(),
-                )
-
-        # If name is provided, get or create the storage by name.
-        elif name:
-            metadata = ApifyKeyValueStoreMetadata.model_validate(await apify_kvss_client.get_or_create(name=name))
-
-            # Freshly fetched named storage. No need to fetch it again.
-            return cls(
-                api_client=apify_client_async.key_value_store(key_value_store_id=metadata.id),
-                api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
-                lock=asyncio.Lock(),
-            )
-        # If id is provided, then storage has to exists.
-        elif id:
-            # Now create the client for the determined ID
-            apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=id)
-            # Fetch its metadata.
-            raw_metadata = await apify_kvs_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={id} failed.')
-            return cls(
-                api_client=apify_kvs_client,
-                api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
-                lock=asyncio.Lock(),
-            )
-        # Default key-value store ID from configuration
-        elif configuration.default_key_value_store_id:
-            # Now create the client for the determined ID
-            apify_kvs_client = apify_client_async.key_value_store(
-                key_value_store_id=configuration.default_key_value_store_id
-            )
-            # Fetch its metadata.
-            raw_metadata = await apify_kvs_client.get()
-            if not raw_metadata:
-                metadata = ApifyKeyValueStoreMetadata.model_validate(await apify_kvss_client.get_or_create(name=name))
-                apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=metadata.id)
-
-            return cls(
-                api_client=apify_kvs_client,
-                api_public_base_url='',  # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
-                lock=asyncio.Lock(),
-            )
-        raise RuntimeError('Will never happen')
+        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
+            lock=asyncio.Lock(),
+        )
 
     @override
     async def purge(self) -> None:

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -5,21 +5,18 @@
 
 from typing_extensions import override
 
-from crawlee._utils.crypto import crypto_random_object_id
 from crawlee.storage_clients._base import RequestQueueClient
-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, create_apify_client
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from apify_client.clients import RequestQueueClientAsync
     from crawlee import Request
+    from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
     from apify import Configuration
 
@@ -224,73 +221,16 @@ 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_rqs_client = apify_client_async.request_queues()
-
-        # Normalize unnamed default storage in cases where not defined in `configuration.default_request_queue_id` to
-        # unnamed storage aliased as `__default__`
-        if not any([alias, name, id, configuration.default_request_queue_id]):
-            alias = '__default__'
-
-        if alias:
-            # Check if there is pre-existing alias mapping in the default KVS.
-            async with AliasResolver(storage_type=RequestQueue, 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 = RequestQueueMetadata.model_validate(
-                        await apify_rqs_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 = RequestQueueMetadata.model_validate(
-                await apify_rqs_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_request_queue_id
-            if not id:
-                raise ValueError(
-                    'RequestQueue "id", "name", or "alias" must be specified, '
-                    'or a default default_request_queue_id ID must be set in the configuration.'
-                )
-
-        # Use suitable client_key to make `hadMultipleClients` response of Apify API useful.
-        # It should persist across migrated or resurrected Actor runs on the Apify platform.
-        _api_max_client_key_length = 32
-        client_key = (configuration.actor_run_id or crypto_random_object_id(length=_api_max_client_key_length))[
-            :_api_max_client_key_length
-        ]
-
-        apify_rq_client = apify_client_async.request_queue(request_queue_id=id, client_key=client_key)
-
-        # Fetch its metadata.
-        metadata = await apify_rq_client.get()
-
-        # If metadata is None, it means the storage does not exist, so we create it.
-        if metadata is None:
-            id = RequestQueueMetadata.model_validate(
-                await apify_rqs_client.get_or_create(),
-            ).id
-            apify_rq_client = apify_client_async.request_queue(request_queue_id=id, client_key=client_key)
-
-        # Verify that the storage exists by fetching its metadata again.
-        metadata = await apify_rq_client.get()
-        if metadata is None:
-            raise ValueError(f'Opening request queue with id={id}, name={name}, and alias={alias} failed.')
-
-        metadata_model = RequestQueueMetadata.model_validate(metadata)
-
-        return cls(api_client=apify_rq_client, metadata=metadata_model, access=access)
+        _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(),
+            metadata=metadata,
+            access=access,
+        )
 
     @override
     async def purge(self) -> None: