refactor: Give more instance caching flexibility to the storage clients (#1425)

### Description

- `StorageClient` type is not an individual storage key and is instead
used in an additional cache key (can be overridden).
- This allows `StorageClients` to have better control over caching.

### Issues


- Enables: apify/apify-sdk-python#513

### Testing

- Unit tests

Author: Pijukatel

src/crawlee/storage_clients/_base/_storage_client.py

@@ -30,12 +30,13 @@ class StorageClient(ABC):
     (where applicable), and consistent access patterns across all storage types it supports.
     """
 
-    def get_additional_cache_key(self, configuration: Configuration) -> Hashable:  # noqa: ARG002
-        """Return a cache key that can differentiate between different storages of this client.
+    def get_storage_client_cache_key(self, configuration: Configuration) -> Hashable:  # noqa: ARG002
+        """Return a cache key that can differentiate between different storages of this and other clients.
 
-        Can be based on configuration or on the client itself. By default, returns an empty string.
+        Can be based on configuration or on the client itself. By default, returns a module and name of the client
+        class.
         """
-        return ''
+        return f'{self.__class__.__module__}.{self.__class__.__name__}'
 
     @abstractmethod
     async def create_dataset_client(

src/crawlee/storage_clients/_file_system/_storage_client.py

@@ -35,9 +35,9 @@ class FileSystemStorageClient(StorageClient):
     """
 
     @override
-    def get_additional_cache_key(self, configuration: Configuration) -> Hashable:
+    def get_storage_client_cache_key(self, configuration: Configuration) -> Hashable:
         # Even different client instances should return same storage if the storage_dir is the same.
-        return configuration.storage_dir
+        return super().get_storage_client_cache_key(configuration), configuration.storage_dir
 
     @override
     async def create_dataset_client(

src/crawlee/storages/_dataset.py

@@ -110,16 +110,15 @@ async def open(
         client_opener_coro = storage_client.create_dataset_client(
             id=id, name=name, alias=alias, configuration=configuration
         )
-        additional_cache_key = storage_client.get_additional_cache_key(configuration=configuration)
+        storage_client_cache_key = storage_client.get_storage_client_cache_key(configuration=configuration)
 
         return await service_locator.storage_instance_manager.open_storage_instance(
             cls,
             id=id,
             name=name,
             alias=alias,
             client_opener_coro=client_opener_coro,
-            storage_client_type=storage_client.__class__,
-            additional_cache_key=additional_cache_key,
+            storage_client_cache_key=storage_client_cache_key,
         )
 
     @override

src/crawlee/storages/_key_value_store.py

@@ -122,16 +122,15 @@ async def open(
         client_opener_coro = storage_client.create_kvs_client(
             id=id, name=name, alias=alias, configuration=configuration
         )
-        additional_cache_key = storage_client.get_additional_cache_key(configuration=configuration)
+        additional_cache_key = storage_client.get_storage_client_cache_key(configuration=configuration)
 
         return await service_locator.storage_instance_manager.open_storage_instance(
             cls,
             id=id,
             name=name,
-            client_opener_coro=client_opener_coro,
             alias=alias,
-            storage_client_type=storage_client.__class__,
-            additional_cache_key=additional_cache_key,
+            client_opener_coro=client_opener_coro,
+            storage_client_cache_key=additional_cache_key,
         )
 
     @override

src/crawlee/storages/_request_queue.py

@@ -126,16 +126,15 @@ async def open(
         storage_client = service_locator.get_storage_client() if storage_client is None else storage_client
 
         client_opener_coro = storage_client.create_rq_client(id=id, name=name, alias=alias, configuration=configuration)
-        additional_cache_key = storage_client.get_additional_cache_key(configuration=configuration)
+        additional_cache_key = storage_client.get_storage_client_cache_key(configuration=configuration)
 
         return await service_locator.storage_instance_manager.open_storage_instance(
             cls,
             id=id,
             name=name,
             alias=alias,
             client_opener_coro=client_opener_coro,
-            storage_client_type=storage_client.__class__,
-            additional_cache_key=additional_cache_key,
+            storage_client_cache_key=additional_cache_key,
         )
 
     @override

src/crawlee/storages/_storage_instance_manager.py

@@ -8,42 +8,56 @@
 from crawlee.storage_clients._base import DatasetClient, KeyValueStoreClient, RequestQueueClient
 
 if TYPE_CHECKING:
-    from crawlee.storage_clients import StorageClient
-
     from ._base import Storage
 
 T = TypeVar('T', bound='Storage')
 
 
 @dataclass
-class _StorageClientCache:
-    """Cache for specific storage client.
-
-    Example:
-        Storage=Dataset, id='123', additional_cache_key="some_path" will be located in
-        storage =  by_id[Dataset]['123'][some_path]
-    """
+class _StorageCache:
+    """Cache for storage instances."""
 
     by_id: defaultdict[type[Storage], defaultdict[str, defaultdict[Hashable, Storage]]] = field(
         default_factory=lambda: defaultdict(lambda: defaultdict(lambda: defaultdict()))
     )
-    """Cache for storage instances by ID, separated by storage type and additional hash key."""
+    """Cache for storage instances by ID. Example: by_id[Dataset]['some_id']['some_additional_cache_key']."""
 
     by_name: defaultdict[type[Storage], defaultdict[str, defaultdict[Hashable, Storage]]] = field(
         default_factory=lambda: defaultdict(lambda: defaultdict(lambda: defaultdict()))
     )
-    """Cache for storage instances by name, separated by storage type and additional hash key."""
+    """Cache for storage instances by name. Example: by_name[Dataset]['some_name']['some_additional_cache_key']"""
 
     by_alias: defaultdict[type[Storage], defaultdict[str, defaultdict[Hashable, Storage]]] = field(
         default_factory=lambda: defaultdict(lambda: defaultdict(lambda: defaultdict()))
     )
-    """Cache for storage instances by alias, separated by storage type and additional hash key."""
+    """Cache for storage instances by alias. Example: by_alias[Dataset]['some_alias']['some_additional_cache_key']"""
+
+    def remove_from_cache(self, storage_instance: Storage) -> None:
+        """Remove a storage instance from the cache.
+
+        Args:
+            storage_instance: The storage instance to remove.
+        """
+        storage_type = type(storage_instance)
+
+        # Remove from ID cache
+        for additional_key in self.by_id[storage_type][storage_instance.id]:
+            del self.by_id[storage_type][storage_instance.id][additional_key]
+            break
 
+        # Remove from name cache or alias cache. It can never be in both.
+        if storage_instance.name is not None:
+            for additional_key in self.by_name[storage_type][storage_instance.name]:
+                del self.by_name[storage_type][storage_instance.name][additional_key]
+                break
+        else:
+            for alias_key in self.by_alias[storage_type]:
+                for additional_key in self.by_alias[storage_type][alias_key]:
+                    del self.by_alias[storage_type][alias_key][additional_key]
+                    break
 
-StorageClientType = DatasetClient | KeyValueStoreClient | RequestQueueClient
-"""Type alias for the storage client types."""
 
-ClientOpenerCoro = Coroutine[None, None, StorageClientType]
+ClientOpenerCoro = Coroutine[None, None, DatasetClient | KeyValueStoreClient | RequestQueueClient]
 """Type alias for the client opener function."""
 
 
@@ -58,7 +72,7 @@ class StorageInstanceManager:
     """Reserved alias for default unnamed storage."""
 
     def __init__(self) -> None:
-        self._cache_by_storage_client: dict[type[StorageClient], _StorageClientCache] = defaultdict(_StorageClientCache)
+        self._cache: _StorageCache = _StorageCache()
 
     async def open_storage_instance(
         self,
@@ -67,9 +81,8 @@ async def open_storage_instance(
         id: str | None,
         name: str | None,
         alias: str | None,
-        storage_client_type: type[StorageClient],
         client_opener_coro: ClientOpenerCoro,
-        additional_cache_key: Hashable = '',
+        storage_client_cache_key: Hashable = '',
     ) -> T:
         """Open a storage instance with caching support.
 
@@ -78,9 +91,8 @@ async def open_storage_instance(
             id: Storage ID.
             name: Storage name. (global scope, persists across runs).
             alias: Storage alias (run scope, creates unnamed storage).
-            storage_client_type: Type of storage client to use.
             client_opener_coro: Coroutine to open the storage client when storage instance not found in cache.
-            additional_cache_key: Additional optional key to differentiate cache entries.
+            storage_client_cache_key: Additional optional key from storage client to differentiate cache entries.
 
         Returns:
             The storage instance.
@@ -105,45 +117,31 @@ async def open_storage_instance(
                 alias = self._DEFAULT_STORAGE_ALIAS
 
             # Check cache
-            if id is not None and (
-                cached_instance := self._cache_by_storage_client[storage_client_type]
-                .by_id[cls][id]
-                .get(additional_cache_key)
-            ):
+            if id is not None and (cached_instance := self._cache.by_id[cls][id].get(storage_client_cache_key)):
                 if isinstance(cached_instance, cls):
                     return cached_instance
                 raise RuntimeError('Cached instance type mismatch.')
 
-            if name is not None and (
-                cached_instance := self._cache_by_storage_client[storage_client_type]
-                .by_name[cls][name]
-                .get(additional_cache_key)
-            ):
+            if name is not None and (cached_instance := self._cache.by_name[cls][name].get(storage_client_cache_key)):
                 if isinstance(cached_instance, cls):
                     return cached_instance
                 raise RuntimeError('Cached instance type mismatch.')
 
             if alias is not None and (
-                cached_instance := self._cache_by_storage_client[storage_client_type]
-                .by_alias[cls][alias]
-                .get(additional_cache_key)
+                cached_instance := self._cache.by_alias[cls][alias].get(storage_client_cache_key)
             ):
                 if isinstance(cached_instance, cls):
                     return cached_instance
                 raise RuntimeError('Cached instance type mismatch.')
 
             # Check for conflicts between named and alias storages
-            if alias and (
-                self._cache_by_storage_client[storage_client_type].by_name[cls][alias].get(additional_cache_key)
-            ):
+            if alias and (self._cache.by_name[cls][alias].get(storage_client_cache_key)):
                 raise ValueError(
                     f'Cannot create alias storage "{alias}" because a named storage with the same name already exists. '
                     f'Use a different alias or drop the existing named storage first.'
                 )
 
-            if name and (
-                self._cache_by_storage_client[storage_client_type].by_alias[cls][name].get(additional_cache_key)
-            ):
+            if name and (self._cache.by_alias[cls][name].get(storage_client_cache_key)):
                 raise ValueError(
                     f'Cannot create named storage "{name}" because an alias storage with the same name already exists. '
                     f'Use a different name or drop the existing alias storage first.'
@@ -160,17 +158,15 @@ async def open_storage_instance(
 
             # Cache the instance.
             # Always cache by id.
-            self._cache_by_storage_client[storage_client_type].by_id[cls][instance.id][additional_cache_key] = instance
+            self._cache.by_id[cls][instance.id][storage_client_cache_key] = instance
 
             # Cache named storage.
             if instance_name is not None:
-                self._cache_by_storage_client[storage_client_type].by_name[cls][instance_name][additional_cache_key] = (
-                    instance
-                )
+                self._cache.by_name[cls][instance_name][storage_client_cache_key] = instance
 
             # Cache unnamed storage.
             if alias is not None:
-                self._cache_by_storage_client[storage_client_type].by_alias[cls][alias][additional_cache_key] = instance
+                self._cache.by_alias[cls][alias][storage_client_cache_key] = instance
 
             return instance
 
@@ -185,25 +181,8 @@ def remove_from_cache(self, storage_instance: Storage) -> None:
         Args:
             storage_instance: The storage instance to remove.
         """
-        storage_type = type(storage_instance)
-
-        for storage_client_cache in self._cache_by_storage_client.values():
-            # Remove from ID cache
-            for additional_key in storage_client_cache.by_id[storage_type][storage_instance.id]:
-                del storage_client_cache.by_id[storage_type][storage_instance.id][additional_key]
-                break
-
-            # Remove from name cache or alias cache. It can never be in both.
-            if storage_instance.name is not None:
-                for additional_key in storage_client_cache.by_name[storage_type][storage_instance.name]:
-                    del storage_client_cache.by_name[storage_type][storage_instance.name][additional_key]
-                    break
-            else:
-                for alias_key in storage_client_cache.by_alias[storage_type]:
-                    for additional_key in storage_client_cache.by_alias[storage_type][alias_key]:
-                        del storage_client_cache.by_alias[storage_type][alias_key][additional_key]
-                        break
+        self._cache.remove_from_cache(storage_instance)
 
     def clear_cache(self) -> None:
         """Clear all cached storage instances."""
-        self._cache_by_storage_client = defaultdict(_StorageClientCache)
+        self._cache = _StorageCache()