feat: Support Actor schema storages with Alias mechanism (#797)

### Description

- Update `Configuration` to include `actor_storages` that is loaded from
`actor_storages_json` env variable.
- Update `AliasResolver` to be able to resolve alias mapping from
`Configuration`.

### Issues

- Closes: #762

### Testing

- Added E2E test
- Added unit tests
- Manual Actor test

### Checklist

- [x] CI passed

---------

Co-authored-by: Vlada Dusek <v.dusek96@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

src/apify/_configuration.py

@@ -8,7 +8,7 @@
 from typing import Annotated, Any
 
 from pydantic import AliasChoices, BeforeValidator, Field, model_validator
-from typing_extensions import Self, deprecated
+from typing_extensions import Self, TypedDict, deprecated
 
 from crawlee import service_locator
 from crawlee._utils.models import timedelta_ms
@@ -34,6 +34,43 @@ def _transform_to_list(value: Any) -> list[str] | None:
     return value if isinstance(value, list) else str(value).split(',')
 
 
+class ActorStorages(TypedDict):
+    """Mapping of storage aliases to their IDs, grouped by storage type.
+
+    Populated from the `ACTOR_STORAGES_JSON` env var that the Apify platform sets when an Actor declares
+    named storages in its `actor.json` schema. Each key maps a user-defined alias (e.g. `'custom'`)
+    to the platform-assigned storage ID.
+    """
+
+    key_value_stores: dict[str, str]
+    datasets: dict[str, str]
+    request_queues: dict[str, str]
+
+
+def _load_storage_keys(data: None | str | ActorStorages) -> ActorStorages | None:
+    """Parse the `ACTOR_STORAGES_JSON` value into a normalized `ActorStorages` dict.
+
+    The platform provides this as a JSON string with camelCase keys (`keyValueStores`, `requestQueues`, `datasets`).
+    This validator deserializes the JSON when needed and normalizes the keys to snake_case, falling back to empty
+    dicts for missing storage types.
+
+    Args:
+        data: Raw value - `None` when the env var is unset, a JSON string from the env var, or an already-parsed
+        `ActorStorages` dict when set programmatically.
+
+    Returns:
+        Normalized storage mapping, or `None` if the input is `None`.
+    """
+    if data is None:
+        return None
+    storage_mapping = json.loads(data) if isinstance(data, str) else data
+    return {
+        'key_value_stores': storage_mapping.get('keyValueStores', storage_mapping.get('key_value_stores', {})),
+        'datasets': storage_mapping.get('datasets', storage_mapping.get('datasets', {})),
+        'request_queues': storage_mapping.get('requestQueues', storage_mapping.get('request_queues', {})),
+    }
+
+
 @docs_group('Configuration')
 class Configuration(CrawleeConfiguration):
     """A class for specifying the configuration of an Actor.
@@ -446,6 +483,15 @@ class Configuration(CrawleeConfiguration):
         BeforeValidator(lambda data: json.loads(data) if isinstance(data, str) else data or None),
     ] = None
 
+    actor_storages: Annotated[
+        ActorStorages | None,
+        Field(
+            alias='actor_storages_json',
+            description='Mapping of storage aliases to their platform-assigned IDs.',
+        ),
+        BeforeValidator(_load_storage_keys),
+    ] = None
+
     @model_validator(mode='after')
     def disable_browser_sandbox_on_platform(self) -> Self:
         """Disable the browser sandbox mode when running on the Apify platform.

src/apify/storage_clients/_apify/_alias_resolving.py

@@ -2,6 +2,7 @@
 
 import logging
 from asyncio import Lock
+from functools import cached_property
 from logging import getLogger
 from typing import TYPE_CHECKING, ClassVar, Literal, overload
 
@@ -139,7 +140,6 @@ def __init__(
         self._storage_type = storage_type
         self._alias = alias
         self._configuration = configuration
-        self._additional_cache_key = hash_api_base_url_and_token(configuration)
 
     async def __aenter__(self) -> AliasResolver:
         """Context manager to prevent race condition in alias creation."""
@@ -183,15 +183,7 @@ async def _get_alias_map(cls, configuration: Configuration) -> dict[str, str]:
             default_kvs_client = await cls._get_default_kvs_client(configuration)
 
             record = await default_kvs_client.get_record(cls._ALIAS_MAPPING_KEY)
-
-            # get_record can return {key: ..., value: ..., content_type: ...}
-            if isinstance(record, dict):
-                if 'value' in record and isinstance(record['value'], dict):
-                    cls._alias_map = record['value']
-                else:
-                    cls._alias_map = record
-            else:
-                cls._alias_map = dict[str, str]()
+            cls._alias_map = record.get('value', {}) if record else {}
 
         return cls._alias_map
 
@@ -201,6 +193,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:
@@ -220,30 +224,22 @@ async def store_mapping(self, storage_id: str) -> None:
 
         try:
             record = await default_kvs_client.get_record(self._ALIAS_MAPPING_KEY)
-
-            # get_record can return {key: ..., value: ..., content_type: ...}
-            if isinstance(record, dict) and 'value' in record:
-                record = record['value']
-
-            # Update or create the record with the new alias mapping
-            if isinstance(record, dict):
-                record[self._storage_key] = storage_id
-            else:
-                record = {self._storage_key: storage_id}
+            value = record.get('value', {}) if record else {}
+            value[self._storage_key] = storage_id
 
             # Store the mapping back in the KVS.
-            await default_kvs_client.set_record(self._ALIAS_MAPPING_KEY, record)
+            await default_kvs_client.set_record(key=self._ALIAS_MAPPING_KEY, value=value)
         except Exception as exc:
             logger.warning(f'Error storing alias mapping for {self._alias}: {exc}')
 
-    @property
+    @cached_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,
                 self._alias,
-                self._additional_cache_key,
+                hash_api_base_url_and_token(self._configuration),
             ]
         )
 

tests/e2e/test_schema_storages/__init__.py

@@ -0,0 +1,24 @@
+{
+    "actorSpecification": 1,
+    "version": "0.0",
+    "storages": {
+        "datasets": {
+            "default": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            },
+            "custom": {
+                "actorSpecification": 1,
+                "fields": {
+                    "properties": {
+                        "id": { "type": "string" }
+                    }
+                }
+            }
+        }
+    }
+}

tests/e2e/test_schema_storages/actor_source/actor.json

@@ -0,0 +1,9 @@
+from apify import Actor
+
+
+async def main() -> None:
+    async with Actor:
+        assert Actor.configuration.actor_storages
+        dataset = await Actor.open_dataset(alias='custom')
+        expected_id = Actor.configuration.actor_storages['datasets']['custom']
+        assert dataset.id == expected_id

tests/e2e/test_schema_storages/actor_source/main.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..conftest import MakeActorFunction, RunActorFunction
+
+_ACTOR_SOURCE_DIR = Path(__file__).parent / 'actor_source'
+
+
+def read_actor_source(filename: str) -> str:
+    return (_ACTOR_SOURCE_DIR / filename).read_text()
+
+
+async def test_configuration_storages(make_actor: MakeActorFunction, run_actor: RunActorFunction) -> None:
+    actor = await make_actor(
+        label='schema_storages',
+        source_files={
+            'src/main.py': read_actor_source('main.py'),
+            '.actor/actor.json': read_actor_source('actor.json'),
+        },
+    )
+    run_result = await run_actor(actor)
+
+    assert run_result.status == 'SUCCEEDED'

tests/e2e/test_schema_storages/test_schema_storages.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 import pytest
@@ -286,7 +287,6 @@ def test_max_total_charge_usd_decimal_parsing(monkeypatch: pytest.MonkeyPatch) -
 
 def test_actor_pricing_info_from_json_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test that actor_pricing_info is parsed from JSON env var."""
-    import json
 
     pricing_json = json.dumps(
         {
@@ -300,3 +300,24 @@ def test_actor_pricing_info_from_json_env_var(monkeypatch: pytest.MonkeyPatch) -
     config = ApifyConfiguration()
     assert config.actor_pricing_info is not None
     assert config.actor_pricing_info.pricing_model == 'PAY_PER_EVENT'
+
+
+def test_actor_storage_json_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that actor_storages_json is parsed from JSON env var."""
+    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'}
+
+    actor_storages_json = json.dumps(
+        {
+            'datasets': datasets,
+            'requestQueues': request_queues,
+            'keyValueStores': key_value_stores,
+        }
+    )
+    monkeypatch.setenv('ACTOR_STORAGES_JSON', actor_storages_json)
+    config = ApifyConfiguration()
+    assert config.actor_storages
+    assert config.actor_storages['datasets'] == datasets
+    assert config.actor_storages['request_queues'] == request_queues
+    assert config.actor_storages['key_value_stores'] == key_value_stores

tests/unit/actor/test_configuration.py

@@ -76,3 +76,24 @@ 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_configuration_storages_alias_resolving() -> None:
+    """Test that `AliasResolver` correctly resolves ids of storages registered in Configuration."""
+
+    # 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 with the storage mapping
+    configuration = Configuration(
+        actor_storages={'datasets': datasets, 'request_queues': request_queues, 'key_value_stores': key_value_stores}
+    )
+
+    # Check that id of each non-default storage saved in the mapping is resolved
+    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'
+        )