Add docs and compute_short_hash for additional cache key

Author: Pijukatel

src/apify/storage_clients/_apify/_storage_client.py

@@ -31,10 +31,12 @@ class ApifyStorageClient(StorageClient):
     @override
     def get_additional_cache_key(self, configuration: CrawleeConfiguration) -> Hashable:
         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)
-        raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
+
+        config_class = type(configuration)
+        raise TypeError(
+            self._lsp_violation_error_message_template.format(f'{config_class.__module__}.{config_class.__name__}')
+        )
 
     @override
     async def create_dataset_client(

src/apify/storage_clients/_apify/_utils.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING, ClassVar
 
 from apify_client import ApifyClientAsync
+from crawlee._utils.crypto import compute_short_hash
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
 from apify._configuration import Configuration
@@ -25,15 +26,18 @@
 class Alias:
     """Class for handling aliases.
 
-    It includes helper methods for serialization/deserialization and initialization from kvs.
+    The purpose of this is class is to ensure that alias storages are created with correct id. This is achieved by using
+    default kvs as a storage for global mapping of aliases to storage ids. Same mapping is also kept in memory to avoid
+    unnecessary calls to API and also have limited support of alias storages when not running on Apify platform. When on
+     Apify platform, the storages created with alias are accessible by the same alias even after migration or reboot.
     """
 
     _alias_map: ClassVar[dict[str, str]] = {}
     """Map containing pre-existing alias storages and their ids. Global for all instances."""
     _alias_init_lock: Lock | None = None
     """Lock for creating alias storages. Only one alias storage can be created at the time. Global for all instances."""
 
-    ALIAS_STORAGE_KEY_SEPARATOR = '|'
+    ALIAS_STORAGE_KEY_SEPARATOR = ','
     ALIAS_MAPPING_KEY = '__STORAGE_ALIASES_MAPPING'
 
     def __init__(self, storage_type: _StorageT, alias: str, configuration: Configuration) -> None:
@@ -55,12 +59,24 @@ async def __aexit__(
 
     @classmethod
     async def _get_alias_init_lock(cls) -> Lock:
+        """Get lock for controlling the creation of the alias storages.
+
+        The lock is shared for all instances of the Alias class.
+        It is created in async method to ensure that some event loop is already running.
+        """
         if cls._alias_init_lock is None:
             cls._alias_init_lock = Lock()
         return cls._alias_init_lock
 
     @classmethod
     async def get_alias_map(cls) -> dict[str, str]:
+        """Get the aliases and storage ids mapping from the default kvs.
+
+        Mapping is loaded from kvs only once and is shared for all instances of the Alias class.
+
+        Returns:
+            Map of aliases and storage ids.
+        """
         if not cls._alias_map:
             default_kvs_client = await get_default_kvs_client()
 
@@ -79,13 +95,17 @@ async def get_alias_map(cls) -> dict[str, str]:
 
     @classmethod
     def get_additional_cache_key(cls, configuration: Configuration) -> str:
-        """Get additional cache key based on api_url and 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 str((configuration.api_base_url, configuration.token))
+        """Get additional cache key based on configuration.
+
+        Use only api_public_base_url and token as the relevant for differentiating storages.
+        """
+        if configuration.api_public_base_url is None or configuration.token is None:
+            raise ValueError("'Configuration.api_public_base_url' and 'Configuration.token' must be set.")
+        return compute_short_hash(f'{configuration.api_public_base_url}{configuration.token}'.encode())
 
     @property
     def storage_key(self) -> str:
+        """Get a unique storage key used for storing the alias in the mapping."""
         return self.ALIAS_STORAGE_KEY_SEPARATOR.join(
             [
                 self.storage_type.__name__,
@@ -95,11 +115,19 @@ def storage_key(self) -> str:
         )
 
     async def resolve_id(self) -> str | None:
+        """Get id of the aliased storage.
+
+        Either locate the id in the in-memory mapping or create the new storage.
+
+        Returns:
+            Storage id if it exists, None otherwise.
+        """
         return (await self.get_alias_map()).get(self.storage_key, None)
 
     async def store_mapping(self, storage_id: str) -> None:
-        """Add alias and related storage id to the mapping in default kvs."""
-        self._alias_map[self.storage_key] = storage_id
+        """Add alias and related storage id to the mapping in default kvs and local in-memory mapping."""
+        # Update in-memory mapping
+        (await self.get_alias_map())[self.storage_key] = storage_id
         if not Configuration.get_global_configuration().is_at_home:
             logging.getLogger(__name__).warning(
                 'Alias storage limited retention is only supported on Apify platform. Storage is not exported.'
@@ -117,7 +145,6 @@ async def store_mapping(self, storage_id: str) -> None:
                 record = record['value']
 
             # Update or create the record with the new alias mapping
-
             if isinstance(record, dict):
                 record[self.storage_key] = storage_id
             else:
@@ -126,7 +153,7 @@ async def store_mapping(self, storage_id: str) -> None:
             # Store the mapping back in the KVS.
             await default_kvs_client.set_record(self.ALIAS_MAPPING_KEY, record)
         except Exception as exc:
-            logger.warning(f'Error accessing alias mapping for {self.alias}: {exc}')
+            logger.warning(f'Error storing alias mapping for {self.alias}: {exc}')
 
 
 async def get_default_kvs_client() -> KeyValueStoreClientAsync:

tests/integration/conftest.py

@@ -20,6 +20,7 @@
 from ._utils import generate_unique_resource_name
 from apify import Actor
 from apify._models import ActorRun
+from apify.storage_clients._apify._utils import Alias
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Iterator, Mapping
@@ -60,6 +61,10 @@ def _prepare_test_env() -> None:
         service_locator._storage_client = None
         service_locator.storage_instance_manager.clear_cache()
 
+        # Reset the Alias class state.
+        Alias._alias_map = {}
+        Alias._alias_init_lock = None
+
         # Verify that the test environment was set up correctly.
         assert os.environ.get(ApifyEnvVars.LOCAL_STORAGE_DIR) == str(tmp_path)
 

tests/integration/test_apify_storages.py

@@ -0,0 +1,34 @@
+import asyncio
+
+import pytest
+
+from crawlee import service_locator
+from crawlee.storages import Dataset, KeyValueStore, RequestQueue
+
+from apify import Configuration
+from apify.storage_clients import ApifyStorageClient
+
+
+@pytest.mark.parametrize(
+    'storage',
+    [Dataset, KeyValueStore, RequestQueue],
+)
+async def test_alias_concurrent_creation_local(
+    storage_type: Dataset | KeyValueStore | RequestQueue, apify_token: str
+) -> None:
+    """Test that storages created with same alias are created only once even when created concurrently."""
+    service_locator.set_configuration(Configuration(token=apify_token))
+    service_locator.set_storage_client(ApifyStorageClient())
+    tasks = [asyncio.create_task(storage_type.open(alias='test')) for _ in range(2)]
+
+    storages = await asyncio.gather(*tasks)
+    unique_storage_ids = {storage.id for storage in storages}
+    try:
+        # Only one aliased storage should be created.
+        assert len(unique_storage_ids) == 1
+
+        # Clean up
+        await storages[0].drop()
+    except AssertionError:
+        for storage in storages:
+            await storage.drop()

tests/unit/actor/test_apify_storage.py renamed to tests/unit/actor/test_apify_storages.py

@@ -24,7 +24,8 @@
 async def test_get_additional_cache_key(
     storage: Storage, _storage_client: ApifyDatasetClient | ApifyKeyValueStoreClient | ApifyRequestQueueClient
 ) -> None:
-    """Test that Storages based on `ApifyStorageClient` include `token` and `api_base_url` in additional cache key."""
+    """Test that Storages based on `ApifyStorageClient` include `token` and `api_public_base_url` in
+    additional cache key."""
 
     def create_metadata(id: str) -> StorageMetadata:
         now = datetime.now(tz=timezone.utc)
@@ -36,10 +37,10 @@ def create_metadata(id: str) -> StorageMetadata:
 
     config_1 = Configuration(token='a')
     config_2 = Configuration(token='b')
-    config_3 = Configuration(token='a', api_base_url='https://super_custom_api.com')
+    config_3 = Configuration(token='a', api_public_base_url='https://super_custom_api.com')
 
     config_4 = Configuration(token='a')
-    config_5 = Configuration(token='a', api_base_url='https://super_custom_api.com')
+    config_5 = Configuration(token='a', api_public_base_url='https://super_custom_api.com')
 
     mocked_client = AsyncMock(spec=type[_storage_client])
     mocked_client.get_metadata = AsyncMock(side_effect=lambda: create_metadata(next(storage_ids)))

tests/unit/conftest.py

@@ -17,6 +17,7 @@
 
 import apify._actor
 import apify.log
+from apify.storage_clients._apify._utils import Alias
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterator
@@ -72,6 +73,10 @@ def _prepare_test_env() -> None:
         service_locator._storage_client = None
         service_locator.storage_instance_manager.clear_cache()
 
+        # Reset the Alias class state.
+        Alias._alias_map = {}
+        Alias._alias_init_lock = None
+
         # Verify that the test environment was set up correctly.
         assert os.environ.get(ApifyEnvVars.LOCAL_STORAGE_DIR) == str(tmp_path)