Review call changes 2

Pijukatel · Pijukatel · commit 2d61f1ee09c5 · 2025-09-17T18:05:21.000+02:00
diff --git a/src/apify/storage_clients/_apify/_dataset_client.py b/src/apify/storage_clients/_apify/_dataset_client.py
@@ -13,7 +13,7 @@
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 from crawlee.storages import Dataset
 
-from ._utils import _Alias
+from apify.storage_clients._apify._utils import Alias
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -128,14 +128,18 @@ async def open(
         alias = None if alias == 'default' else alias
 
         if alias:
-            # 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
-            )
+            # Check if there is pre-existing alias mapping in the default KVS.
+            async with Alias(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:
diff --git a/src/apify/storage_clients/_apify/_key_value_store_client.py b/src/apify/storage_clients/_apify/_key_value_store_client.py
@@ -13,7 +13,7 @@
 from crawlee.storages import KeyValueStore
 
 from ._models import ApifyKeyValueStoreMetadata, KeyValueStoreListKeysPage
-from ._utils import _Alias
+from ._utils import Alias
 from apify._crypto import create_hmac_signature
 
 if TYPE_CHECKING:
@@ -119,14 +119,19 @@ async def open(
         alias = None if alias == 'default' else alias
 
         if alias:
-            # 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)
+            # Check if there is pre-existing alias mapping in the default KVS.
+            async with Alias(storage_type=KeyValueStore, 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:
+                    # 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.store_mapping(storage_id=id)
 
         # If name is provided, get or create the storage by name.
         elif name:
diff --git a/src/apify/storage_clients/_apify/_request_queue_client.py b/src/apify/storage_clients/_apify/_request_queue_client.py
@@ -19,7 +19,7 @@
 from crawlee.storages import RequestQueue
 
 from ._models import CachedRequest, ProlongRequestLockResponse, RequestQueueHead
-from ._utils import _Alias
+from ._utils import Alias
 from apify import Request
 
 if TYPE_CHECKING:
@@ -194,14 +194,18 @@ async def open(
         apify_rqs_client = apify_client_async.request_queues()
 
         if alias:
-            # 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)
+            # Check if there is pre-existing alias mapping in the default KVS.
+            async with Alias(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:
diff --git a/src/apify/storage_clients/_apify/_storage_client.py b/src/apify/storage_clients/_apify/_storage_client.py
@@ -1,19 +1,15 @@
 from __future__ import annotations
 
-import logging
-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 ._utils import Alias
 from apify._configuration import Configuration as ApifyConfiguration
 from apify._utils import docs_group
 
@@ -27,11 +23,6 @@
 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.'
@@ -42,7 +33,7 @@ def get_additional_cache_key(self, configuration: CrawleeConfiguration) -> Hasha
         if isinstance(configuration, ApifyConfiguration):
             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)
+            return Alias.get_additional_cache_key(configuration)
         raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
 
     @override
@@ -56,10 +47,6 @@ 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__))
@@ -75,12 +62,6 @@ 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__))
@@ -96,66 +77,6 @@ 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.
-        """
-        if not Configuration.get_global_configuration().is_at_home:
-            logging.getLogger(__name__).warning(
-                'Alias storage limited retention is only supported on Apify platform. '
-                'No pre-existing storages are imported.'
-            )
-            cls._alias_storages_initialized = True
-            return
-
-        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
diff --git a/src/apify/storage_clients/_apify/_utils.py b/src/apify/storage_clients/_apify/_utils.py
