Simplify and avoid unnecessary API calls

Author: Pijukatel

src/apify/_actor.py

@@ -36,7 +36,6 @@
 from apify.events import ApifyEventManager, EventManager, LocalEventManager
 from apify.log import _configure_logging, logger
 from apify.storage_clients import ApifyStorageClient, SmartApifyStorageClient
-from apify.storage_clients._apify._alias_resolving import AliasResolver
 from apify.storage_clients._file_system import ApifyFileSystemStorageClient
 from apify.storages import Dataset, KeyValueStore, RequestQueue
 
@@ -204,10 +203,6 @@ async def __aenter__(self) -> Self:
         if not Actor.is_at_home():
             # Make sure that the input related KVS is initialized to ensure that the input aware client is used
             await self.open_key_value_store()
-        else:
-            # Load pre-existing non-default aliased storages from configuration
-            # Supported only on the Apify platform, where those storages are pre-created by the platform.
-            await AliasResolver.register_aliases(configuration=self.configuration)
         return self
 
     async def __aexit__(

src/apify/storage_clients/_apify/_alias_resolving.py

@@ -5,6 +5,8 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, ClassVar, Literal, overload
 
+from propcache import cached_property
+
 from apify_client import ApifyClientAsync
 
 from ._utils import hash_api_base_url_and_token
@@ -136,11 +138,9 @@ def __init__(
         alias: str,
         configuration: Configuration,
     ) -> None:
+        self._storage_type = storage_type
         self._alias = alias
         self._configuration = configuration
-        self._storage_key = self.get_storage_key(
-            storage_type=storage_type, alias=alias, additional_cache_key=hash_api_base_url_and_token(configuration)
-        )
 
     async def __aenter__(self) -> AliasResolver:
         """Context manager to prevent race condition in alias creation."""
@@ -194,6 +194,18 @@ async def resolve_id(self) -> str | None:
         Returns:
             Storage id if it exists, None otherwise.
         """
+        # First try to find the alias in the configuration mapping to avoid any API calls.
+        # This mapping is maintained by the Apify platform and does not have to be maintained in the default KVS.
+        if self._configuration.actor_storages and self._alias != 'default':
+            storage_maps = {
+                'Dataset': self._configuration.actor_storages.datasets,
+                'KeyValueStore': self._configuration.actor_storages.key_value_stores,
+                'RequestQueue': self._configuration.actor_storages.request_queues,
+            }
+            if storage_id := storage_maps.get(self._storage_type, {}).get(self._alias):
+                return storage_id
+
+        # Fallback to the mapping saved in the default KVS
         return (await self._get_alias_map(self._configuration)).get(self._storage_key, None)
 
     async def store_mapping(self, storage_id: str) -> None:
@@ -237,60 +249,12 @@ async def _get_default_kvs_client(configuration: Configuration) -> KeyValueStore
 
         return apify_client_async.key_value_store(key_value_store_id=configuration.default_key_value_store_id)
 
-    @classmethod
-    def get_storage_key(
-        cls, storage_type: Literal['Dataset', 'KeyValueStore', 'RequestQueue'], alias: str, additional_cache_key: str
-    ) -> str:
-        return cls._ALIAS_STORAGE_KEY_SEPARATOR.join(
+    @cached_property
+    def _storage_key(self) -> str:
+        return self._ALIAS_STORAGE_KEY_SEPARATOR.join(
             [
-                storage_type,
-                alias,
-                additional_cache_key,
+                self._storage_type,
+                self._alias,
+                hash_api_base_url_and_token(self._configuration),
             ]
         )
-
-    @classmethod
-    async def register_aliases(cls, configuration: Configuration) -> None:
-        """Load alias mapping from configuration to the default kvs."""
-        async with await cls._get_alias_init_lock():
-            # Skip if no mapping or just default storages
-            if configuration.actor_storages is None or set(
-                configuration.actor_storages.datasets.keys()
-                | configuration.actor_storages.key_value_stores.keys()
-                | configuration.actor_storages.request_queues.keys()
-            ) == {'default'}:
-                return
-
-            configuration_mapping = {}
-
-            if configuration.default_dataset_id != configuration.actor_storages.datasets.get('default'):
-                logger.warning(
-                    f'Conflicting default dataset ids: {configuration.default_dataset_id=},'
-                    f" {configuration.actor_storages.datasets.get('default')=}"
-                )
-            additional_cache_key = hash_api_base_url_and_token(configuration)
-
-            for mapping, storage_type in (
-                (configuration.actor_storages.key_value_stores, 'KeyValueStore'),
-                (configuration.actor_storages.datasets, 'Dataset'),
-                (configuration.actor_storages.request_queues, 'RequestQueue'),
-            ):
-                for storage_alias, storage_id in mapping.items():
-                    configuration_mapping[
-                        cls.get_storage_key(
-                            storage_type,
-                            '__default__' if storage_alias == 'default' else storage_alias,
-                            additional_cache_key,
-                        )
-                    ] = storage_id
-
-            # Bulk update the mapping in the default KVS with the configuration mapping.
-            client = await cls._get_default_kvs_client(configuration=configuration)
-            record = await client.get_record(cls._ALIAS_MAPPING_KEY)
-            existing_mapping = record.get('value', {}) if record else {}
-
-            # Update the existing mapping with the configuration mapping.
-            existing_mapping.update(configuration_mapping)
-            # Store the updated mapping back in the KVS and in memory.
-            await client.set_record(cls._ALIAS_MAPPING_KEY, existing_mapping)
-            cls._alias_map.update(existing_mapping)

tests/e2e/test_schema_storages/actor_source/actor.json

@@ -11,7 +11,43 @@
                     }
                 }
             },
-            "custom": {
+            "custom_d": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            }
+        },
+        "key_value_stores": {
+            "default": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            },
+            "custom_kvs": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            }
+        },
+        "request_queues": {
+            "default": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            },
+            "custom_rq": {
                 "actorSpecification": 1,
                 "fields": {
                     "properties": {

tests/e2e/test_schema_storages/actor_source/main.py

@@ -4,4 +4,12 @@
 async def main() -> None:
     async with Actor:
         assert Actor.configuration.actor_storages
-        assert (await Actor.open_dataset(alias='custom')).id == Actor.configuration.actor_storages.datasets['custom']
+        assert (await Actor.open_dataset(alias='custom_d')).id == Actor.configuration.actor_storages.datasets[
+            'custom_d'
+        ]
+        assert (await Actor.open_dataset(alias='custom_kvs')).id == Actor.configuration.actor_storages.datasets[
+            'custom_kvs'
+        ]
+        assert (await Actor.open_dataset(alias='custom_rq')).id == Actor.configuration.actor_storages.datasets[
+            'custom_rq'
+        ]

tests/integration/test_storages.py

@@ -1,17 +1,14 @@
 from __future__ import annotations
 
 import asyncio
-from typing import cast
 
 import pytest
 
 from crawlee import service_locator
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
 from apify import Actor, Configuration
-from apify._configuration import ActorStorages
 from apify.storage_clients import ApifyStorageClient, MemoryStorageClient, SmartApifyStorageClient
-from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
 @pytest.mark.parametrize(
@@ -128,53 +125,3 @@ async def test_actor_implicit_storage_init(apify_token: str) -> None:
         assert await Actor.open_dataset() is not await Actor.open_dataset(force_cloud=True)
         assert await Actor.open_key_value_store() is not await Actor.open_key_value_store(force_cloud=True)
         assert await Actor.open_request_queue() is not await Actor.open_request_queue(force_cloud=True)
-
-
-async def test_actor_storages_alias_resolving(apify_token: str) -> None:
-    """Test that `AliasResolver.register_aliases` correctly updates default KVS with Actor storages."""
-
-    # Actor storages
-    datasets = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
-    request_queues = {'default': 'default_request_queue_id', 'custom': 'custom_request_queue_id'}
-    key_value_stores = {'default': 'default_key_value_store_id', 'custom': 'custom_key_value_store_id'}
-
-    # Set up the configuration and storage client for the test
-    configuration = Configuration(
-        default_key_value_store_id='default_kvs_id',
-        token=apify_token,
-        actor_storages=ActorStorages(
-            datasets=datasets, request_queues=request_queues, key_value_stores=key_value_stores
-        ),
-    )
-    storage_client = ApifyStorageClient()
-    service_locator.set_configuration(configuration)
-    service_locator.set_storage_client(storage_client)
-
-    client_cache_key = cast('tuple', storage_client.get_storage_client_cache_key(configuration))[-1]
-    # Add some unrelated pre-existing alias mapping (it should be preserved after registering aliases)
-    pre_existing_mapping = {f'KeyValueStore,pre_existing_alias,{client_cache_key}': 'pre_existing_id'}
-
-    default_kvs = await KeyValueStore.open(configuration=configuration, storage_client=storage_client)
-    await default_kvs.set_value(AliasResolver._ALIAS_MAPPING_KEY, pre_existing_mapping)
-
-    # Construct the expected mapping
-    expected_mapping = {}
-    for storage_type, storage_map in (
-        ('Dataset', datasets),
-        ('KeyValueStore', key_value_stores),
-        ('RequestQueue', request_queues),
-    ):
-        for storage_alias, storage_id in storage_map.items():
-            expected_mapping[
-                ','.join(
-                    (storage_type, '__default__' if storage_alias == 'default' else storage_alias, client_cache_key)
-                )
-            ] = storage_id
-    expected_mapping.update(pre_existing_mapping)
-
-    try:
-        configuration.default_key_value_store_id = default_kvs.id
-        await AliasResolver.register_aliases(configuration=configuration)
-        assert await default_kvs.get_value(AliasResolver._ALIAS_MAPPING_KEY) == expected_mapping
-    finally:
-        await default_kvs.drop()

tests/unit/storage_clients/test_alias_resolver.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from apify._configuration import Configuration
+from apify._configuration import ActorStorages, Configuration
 from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
@@ -76,3 +76,27 @@ async def test_get_alias_map_returns_in_memory_map() -> None:
     AliasResolver._alias_map = {}
     result = await AliasResolver._get_alias_map(config)
     assert result == {}
+
+
+async def test_actor_storages_alias_resolving() -> None:
+    """Test that `AliasResolver.register_aliases` correctly updates default KVS with Actor storages."""
+
+    # Actor storages
+    datasets = {'default': 'default_dataset_id', 'custom': 'custom_Dataset_id'}
+    request_queues = {'default': 'default_request_queue_id', 'custom': 'custom_RequestQueue_id'}
+    key_value_stores = {'default': 'default_key_value_store_id', 'custom': 'custom_KeyValueStore_id'}
+
+    # Set up the configuration and storage client for the test
+    configuration = Configuration(
+        default_key_value_store_id='default_kvs_id',
+        actor_storages=ActorStorages(
+            datasets=datasets, request_queues=request_queues, key_value_stores=key_value_stores
+        ),
+    )
+
+    # Construct the expected mapping
+    for storage_type in ('Dataset', 'KeyValueStore', 'RequestQueue'):
+        assert (
+            await AliasResolver(storage_type=storage_type, alias='custom', configuration=configuration).resolve_id()
+            == f'custom_{storage_type}_id'
+        )