Update NDU creation logic based on updated Crawlee

Author: Pijukatel

src/apify/storage_clients/_apify/_dataset_client.py

@@ -11,8 +11,9 @@
 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 resolve_alias_to_id, store_alias_mapping
+from ._utils import _Alias
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -126,19 +127,15 @@ async def open(
         # Normalize 'default' alias to None
         alias = None if alias == 'default' else alias
 
-        # Handle alias resolution
         if alias:
-            # Try to resolve alias to existing storage ID
-            resolved_id = await resolve_alias_to_id(alias, 'dataset', configuration)
-            if resolved_id:
-                id = resolved_id
-            else:
-                # Create a new storage and store the alias mapping
-                new_storage_metadata = DatasetMetadata.model_validate(
-                    await apify_datasets_client.get_or_create(),
-                )
-                id = new_storage_metadata.id
-                await store_alias_mapping(alias, 'dataset', id, configuration)
+            # Create a new storage and store the alias mapping
+            new_storage_metadata = DatasetMetadata.model_validate(
+                await apify_datasets_client.get_or_create(),
+            )
+            id = new_storage_metadata.id
+            await _Alias(storage_type=Dataset, alias=alias, token=token, api_url=api_url).store_mapping_to_apify_kvs(
+                storage_id=id
+            )
 
         # If name is provided, get or create the storage by name.
         elif name:

src/apify/storage_clients/_apify/_key_value_store_client.py

@@ -10,9 +10,10 @@
 from apify_client import ApifyClientAsync
 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 resolve_alias_to_id, store_alias_mapping
+from ._utils import _Alias
 from apify._crypto import create_hmac_signature
 
 if TYPE_CHECKING:
@@ -117,19 +118,15 @@ async def open(
         # Normalize 'default' alias to None
         alias = None if alias == 'default' else alias
 
-        # Handle alias resolution
         if alias:
-            # Try to resolve alias to existing storage ID
-            resolved_id = await resolve_alias_to_id(alias, 'kvs', configuration)
-            if resolved_id:
-                id = resolved_id
-            else:
-                # Create a new storage and store the alias mapping
-                new_storage_metadata = ApifyKeyValueStoreMetadata.model_validate(
-                    await apify_kvss_client.get_or_create(),
-                )
-                id = new_storage_metadata.id
-                await store_alias_mapping(alias, 'kvs', id, configuration)
+            # Create a new storage and store the alias mapping
+            new_storage_metadata = ApifyKeyValueStoreMetadata.model_validate(
+                await apify_kvss_client.get_or_create(),
+            )
+            id = new_storage_metadata.id
+            await _Alias(
+                storage_type=KeyValueStore, alias=alias, token=token, api_url=api_url
+            ).store_mapping_to_apify_kvs(storage_id=id)
 
         # If name is provided, get or create the storage by name.
         elif name:

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -16,9 +16,10 @@
 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 CachedRequest, ProlongRequestLockResponse, RequestQueueHead
-from ._utils import resolve_alias_to_id, store_alias_mapping
+from ._utils import _Alias
 from apify import Request
 
 if TYPE_CHECKING:
@@ -192,22 +193,15 @@ async def open(
         )
         apify_rqs_client = apify_client_async.request_queues()
 
-        # Normalize 'default' alias to None
-        alias = None if alias == 'default' else alias
-
-        # Handle alias resolution
         if alias:
-            # Try to resolve alias to existing storage ID
-            resolved_id = await resolve_alias_to_id(alias, 'rq', configuration)
-            if resolved_id:
-                id = resolved_id
-            else:
-                # Create a new storage and store the alias mapping
-                new_storage_metadata = RequestQueueMetadata.model_validate(
-                    await apify_rqs_client.get_or_create(),
-                )
-                id = new_storage_metadata.id
-                await store_alias_mapping(alias, 'rq', id, configuration)
+            # Create a new storage and store the alias mapping
+            new_storage_metadata = RequestQueueMetadata.model_validate(
+                await apify_rqs_client.get_or_create(),
+            )
+            id = new_storage_metadata.id
+            await _Alias(
+                storage_type=RequestQueue, alias=alias, token=token, api_url=api_url
+            ).store_mapping_to_apify_kvs(storage_id=id)
 
         # If name is provided, get or create the storage by name.
         elif name:

src/apify/storage_clients/_apify/_storage_client.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
+from asyncio import Lock
 from typing import TYPE_CHECKING
 
 from typing_extensions import override
 
+from crawlee import service_locator
 from crawlee.storage_clients._base import StorageClient
 
 from ._dataset_client import ApifyDatasetClient
 from ._key_value_store_client import ApifyKeyValueStoreClient
 from ._request_queue_client import ApifyRequestQueueClient
+from ._utils import _ALIAS_MAPPING_KEY, _Alias
+from apify._configuration import Configuration
 from apify._configuration import Configuration as ApifyConfiguration
 from apify._utils import docs_group
 
@@ -22,6 +26,11 @@
 class ApifyStorageClient(StorageClient):
     """Apify storage client."""
 
+    _alias_storages_initialized = False
+    """Flag that indicates whether the pre-existing alias storages were already initialized."""
+    _alias_init_lock: Lock | None = None
+    """Lock for creating alias storages. Only one alias storage can be created at the time."""
+
     # 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.'
@@ -30,7 +39,9 @@ class ApifyStorageClient(StorageClient):
     @override
     def get_additional_cache_key(self, configuration: CrawleeConfiguration) -> Hashable:
         if isinstance(configuration, ApifyConfiguration):
-            return f'{configuration.api_base_url},{configuration.token}'
+            if configuration.api_base_url is None or configuration.token is None:
+                raise ValueError("'Configuration.api_base_url' and 'Configuration.token' must be set.")
+            return _Alias.get_additional_cache_key(configuration.api_base_url, configuration.token)
         raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
 
     @override
@@ -44,6 +55,10 @@ async def create_dataset_client(
     ) -> ApifyDatasetClient:
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
+            if alias:
+                await self._initialize_alias_storages()
+                async with self.get_alias_init_lock():
+                    return await ApifyDatasetClient.open(id=id, name=name, alias=alias, configuration=configuration)
             return await ApifyDatasetClient.open(id=id, name=name, alias=alias, configuration=configuration)
 
         raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
@@ -59,6 +74,12 @@ async def create_kvs_client(
     ) -> ApifyKeyValueStoreClient:
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
+            if alias:
+                await self._initialize_alias_storages()
+                async with self.get_alias_init_lock():
+                    return await ApifyKeyValueStoreClient.open(
+                        id=id, name=name, alias=alias, configuration=configuration
+                    )
             return await ApifyKeyValueStoreClient.open(id=id, name=name, alias=alias, configuration=configuration)
 
         raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
@@ -74,6 +95,58 @@ async def create_rq_client(
     ) -> ApifyRequestQueueClient:
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
+            if alias:
+                await self._initialize_alias_storages()
+                async with self.get_alias_init_lock():
+                    return await ApifyRequestQueueClient.open(
+                        id=id, name=name, alias=alias, configuration=configuration
+                    )
             return await ApifyRequestQueueClient.open(id=id, name=name, alias=alias, configuration=configuration)
 
         raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
+
+    @classmethod
+    def get_alias_init_lock(cls) -> Lock:
+        if not cls._alias_init_lock:
+            cls._alias_init_lock = Lock()
+        return cls._alias_init_lock
+
+    @classmethod
+    async def _initialize_alias_storages(cls) -> None:
+        """Initialize alias storages.
+
+        This method is called once to populate storage_instance_manager alias related cache. All existing alias
+        storages are saved in storage_instance_manager cache. If the alias storage is not there, it does not exist yet.
+        """
+        async with cls.get_alias_init_lock():
+            if cls._alias_storages_initialized:
+                return
+
+            cache = service_locator.storage_instance_manager._cache_by_storage_client[ApifyStorageClient]  # noqa: SLF001
+
+            default_kvs_client = await _Alias.get_default_kvs_client()
+
+            record = await default_kvs_client.get_record(key=_ALIAS_MAPPING_KEY)
+
+            if record is not None and 'value' in record:
+                # get_record can return {key: ..., value: ..., content_type: ...}
+                alias_export_map = record['value']
+
+                for export_key, storage_id in alias_export_map.value.items():
+                    exported_alias = _Alias.from_exported_string(export_key)
+
+                    # Re-create custom config used to open the storage
+                    custom_config = Configuration()
+                    custom_config.api_base_url = exported_alias.api_url
+                    custom_config.token = exported_alias.token
+
+                    # Populate the id cache by opening storage by id
+                    storage = await exported_alias.storage_type.open(
+                        id=storage_id, configuration=custom_config, storage_client=ApifyStorageClient()
+                    )
+                    # Populate the alias cache as well
+                    cache.by_alias[exported_alias.storage_type][exported_alias.alias][
+                        exported_alias.additional_cache_key
+                    ] = storage
+
+        cls._alias_storages_initialized = True