Add special caching for ApifyClient

TODO: Merge update from NDU storages

Author: Pijukatel

src/apify/storage_clients/_apify/_storage_client.py

@@ -9,72 +9,71 @@
 from ._dataset_client import ApifyDatasetClient
 from ._key_value_store_client import ApifyKeyValueStoreClient
 from ._request_queue_client import ApifyRequestQueueClient
+from apify._configuration import Configuration as ApifyConfiguration
 from apify._utils import docs_group
 
 if TYPE_CHECKING:
-    from crawlee.configuration import Configuration
+    from collections.abc import Hashable
+
+    from crawlee.configuration import Configuration as CrawleeConfiguration
 
 
 @docs_group('Storage clients')
 class ApifyStorageClient(StorageClient):
     """Apify storage client."""
 
+    # 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.'
+    )
+
+    @override
+    def get_additional_cache_key(self, configuration: CrawleeConfiguration) -> Hashable:
+        if isinstance(configuration, ApifyConfiguration):
+            return f'{configuration.api_base_url},{configuration.token}'
+        raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
+
     @override
     async def create_dataset_client(
         self,
         *,
         id: str | None = None,
         name: str | None = None,
-        configuration: Configuration | None = None,
+        alias: str | None = None,
+        configuration: CrawleeConfiguration | None = None,
     ) -> ApifyDatasetClient:
-        # Import here to avoid circular imports.
-        from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
-
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
             return await ApifyDatasetClient.open(id=id, name=name, configuration=configuration)
 
-        raise TypeError(
-            f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(configuration).__name__} instead.'
-        )
+        raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
 
     @override
     async def create_kvs_client(
         self,
         *,
         id: str | None = None,
         name: str | None = None,
-        configuration: Configuration | None = None,
+        alias: str | None = None,
+        configuration: CrawleeConfiguration | None = None,
     ) -> ApifyKeyValueStoreClient:
-        # Import here to avoid circular imports.
-        from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
-
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
             return await ApifyKeyValueStoreClient.open(id=id, name=name, configuration=configuration)
 
-        raise TypeError(
-            f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(configuration).__name__} instead.'
-        )
+        raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))
 
     @override
     async def create_rq_client(
         self,
         *,
         id: str | None = None,
         name: str | None = None,
-        configuration: Configuration | None = None,
+        alias: str | None = None,
+        configuration: CrawleeConfiguration | None = None,
     ) -> ApifyRequestQueueClient:
-        # Import here to avoid circular imports.
-        from apify import Configuration as ApifyConfiguration  # noqa: PLC0415
-
         configuration = configuration or ApifyConfiguration.get_global_configuration()
         if isinstance(configuration, ApifyConfiguration):
             return await ApifyRequestQueueClient.open(id=id, name=name, configuration=configuration)
 
-        raise TypeError(
-            f'Expected "configuration" to be an instance of "apify.Configuration", '
-            f'but got {type(configuration).__name__} instead.'
-        )
+        raise TypeError(self._lsp_violation_error_message_template.format(type(configuration).__name__))

src/apify/storage_clients/_file_system/_storage_client.py

@@ -27,9 +27,12 @@ async def create_kvs_client(
         *,
         id: str | None = None,
         name: str | None = None,
+        alias: str | None = None,
         configuration: Configuration | None = None,
     ) -> FileSystemKeyValueStoreClient:
         configuration = configuration or Configuration.get_global_configuration()
-        client = await ApifyFileSystemKeyValueStoreClient.open(id=id, name=name, configuration=configuration)
+        client = await ApifyFileSystemKeyValueStoreClient.open(
+            id=id, name=name, alias=alias, configuration=configuration
+        )
         await self._purge_if_needed(client, configuration)
         return client

tests/unit/actor/test_apify_storage.py

@@ -0,0 +1,54 @@
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+
+from crawlee.storages import Dataset, KeyValueStore, RequestQueue
+from crawlee.storages._base import Storage
+
+from apify import Configuration
+from apify.storage_clients import ApifyStorageClient
+from apify.storage_clients._apify import ApifyDatasetClient, ApifyKeyValueStoreClient, ApifyRequestQueueClient
+
+
+@pytest.mark.parametrize(
+    ('storage', '_storage_client'),
+    [
+        (Dataset, ApifyDatasetClient),
+        (KeyValueStore, ApifyKeyValueStoreClient),
+        (RequestQueue, ApifyRequestQueueClient),
+    ],
+)
+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."""
+    storage_names = iter(['1', '2', '3', '1', '3'])
+
+    apify_storage_client = ApifyStorageClient()
+
+    config_1 = Configuration(token='a')
+    config_2 = Configuration(token='b')
+    config_3 = Configuration(token='a', api_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')
+
+    mocked_open = AsyncMock(spec=_storage_client.open)
+    mocked_open.get_metadata = AsyncMock(storage_names)
+
+    with mock.patch.object(_storage_client, 'open', mocked_open):
+        storage_1 = await storage.open(storage_client=apify_storage_client, configuration=config_1)
+        storage_2 = await storage.open(storage_client=apify_storage_client, configuration=config_2)
+        storage_3 = await storage.open(storage_client=apify_storage_client, configuration=config_3)
+        storage_4 = await storage.open(storage_client=apify_storage_client, configuration=config_4)
+        storage_5 = await storage.open(storage_client=apify_storage_client, configuration=config_5)
+
+        # Different configuration results in different storage clients.
+        assert storage_1 is not storage_2
+        assert storage_1 is not storage_3
+        assert storage_2 is not storage_3
+
+        # Equivalent configuration results in same storage clients.
+        assert storage_1 is storage_4
+        assert storage_3 is storage_5

tests/unit/storage_clients/test_file_system.py

@@ -22,6 +22,7 @@ async def test_purge_preserves_input_file_and_metadata() -> None:
     kvs_storage_client = await ApifyFileSystemKeyValueStoreClient.open(
         id=None,
         name='test-kvs',
+        alias=None,
         configuration=configuration,
     )