Refactor storages to allow force_cloud feature

Author: vdusek

src/crawlee/storages/_base_storage.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from crawlee.base_storage_client import BaseStorageClient
+    from crawlee.configuration import Configuration
 
 
 class BaseStorage(ABC):
@@ -18,14 +23,24 @@ def name(self) -> str | None:
 
     @classmethod
     @abstractmethod
-    async def open(cls, *, id: str | None = None, name: str | None = None) -> BaseStorage:
+    async def open(
+        cls,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        configuration: Configuration | None = None,
+        storage_client: BaseStorageClient | None = None,
+    ) -> BaseStorage:
         """Open a storage, either restore existing or create a new one.
 
         Args:
             id: The storage ID.
             name: The storage name.
+            configuration: Configuration object used during the storage creation or restoration process.
+            storage_client: Underlying storage client to use. If not provided, the default global storage client
+                from the service locator will be used.
         """
 
     @abstractmethod
     async def drop(self) -> None:
-        """Drop the storage. Remove it from underlying storage and delete from cache."""
+        """Drops the storage, removing it from the underlying storage client and clearing the cache."""

src/crawlee/storages/_creation_management.py

@@ -3,7 +3,6 @@
 import asyncio
 from typing import TYPE_CHECKING, TypeVar
 
-from crawlee import service_container
 from crawlee.memory_storage_client import MemoryStorageClient
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
@@ -124,18 +123,17 @@ async def open_storage(
     storage_class: type[TResource],
     id: str | None = None,
     name: str | None = None,
+    configuration: Configuration,
+    storage_client: BaseStorageClient,
 ) -> TResource:
     """Open either a new storage or restore an existing one and return it."""
-    config = service_container.get_configuration()
-    storage_client = service_container.get_storage_client()
-
     # Try to restore the storage from cache by name
     if name:
         cached_storage = _get_from_cache_by_name(storage_class=storage_class, name=name)
         if cached_storage:
             return cached_storage
 
-    default_id = _get_default_storage_id(config, storage_class)
+    default_id = _get_default_storage_id(configuration, storage_class)
 
     if not id and not name:
         id = default_id
@@ -150,7 +148,7 @@ async def open_storage(
             return cached_storage
 
     # Purge on start if configured
-    if config.purge_on_start:
+    if configuration.purge_on_start:
         await storage_client.purge_on_start()
 
     # Lock and create new storage
@@ -169,7 +167,7 @@ async def open_storage(
             resource_collection_client = _get_resource_collection_client(storage_class, storage_client)
             storage_info = await resource_collection_client.get_or_create(name=name)
 
-        storage = storage_class(id=storage_info.id, name=storage_info.name)
+        storage = storage_class(id=storage_info.id, name=storage_info.name, storage_client=storage_client)
 
         # Cache the storage by ID and name
         _add_to_cache_by_id(storage.id, storage)

src/crawlee/storages/_dataset.py

@@ -20,8 +20,9 @@
     from collections.abc import AsyncIterator, Callable
 
     from crawlee._types import JsonSerializable, PushDataKwargs
+    from crawlee.base_storage_client import BaseStorageClient
     from crawlee.base_storage_client._models import DatasetItemsListPage
-
+    from crawlee.configuration import Configuration
 
 logger = logging.getLogger(__name__)
 
@@ -192,13 +193,11 @@ class Dataset(BaseStorage):
     _EFFECTIVE_LIMIT_SIZE = _MAX_PAYLOAD_SIZE - (_MAX_PAYLOAD_SIZE * _SAFETY_BUFFER_PERCENT)
     """Calculated payload limit considering safety buffer."""
 
-    def __init__(self, id: str, name: str | None) -> None:
-        storage_client = service_container.get_storage_client()
-
+    def __init__(self, id: str, name: str | None, storage_client: BaseStorageClient) -> None:
         self._id = id
         self._name = name
 
-        # Get resource clients from storage client
+        # Get resource clients from the storage client.
         self._resource_client = storage_client.dataset(self._id)
         self._resource_collection_client = storage_client.datasets()
 
@@ -214,10 +213,26 @@ def name(self) -> str | None:
 
     @override
     @classmethod
-    async def open(cls, *, id: str | None = None, name: str | None = None) -> Dataset:
+    async def open(
+        cls,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        configuration: Configuration | None = None,
+        storage_client: BaseStorageClient | None = None,
+    ) -> Dataset:
         from crawlee.storages._creation_management import open_storage
 
-        return await open_storage(storage_class=cls, id=id, name=name)
+        configuration = configuration or service_container.get_configuration()
+        storage_client = storage_client or service_container.get_storage_client()
+
+        return await open_storage(
+            storage_class=cls,
+            id=id,
+            name=name,
+            configuration=configuration,
+            storage_client=storage_client,
+        )
 
     @override
     async def drop(self) -> None:
@@ -241,7 +256,7 @@ async def push_data(self, data: JsonSerializable, **kwargs: Unpack[PushDataKwarg
         # Handle singular items
         if not isinstance(data, list):
             items = await self.check_and_serialize(data)
-            return await self._resource_client.push_items(items, **kwargs)  # type: ignore[no-any-return] # Mypy is broken
+            return await self._resource_client.push_items(items, **kwargs)
 
         # Handle lists
         payloads_generator = (await self.check_and_serialize(item, index) for index, item in enumerate(data))
@@ -264,7 +279,7 @@ async def get_data(self, **kwargs: Unpack[GetDataKwargs]) -> DatasetItemsListPag
         Returns:
             List page containing filtered and paginated dataset items.
         """
-        return await self._resource_client.list_items(**kwargs)  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.list_items(**kwargs)
 
     async def write_to_csv(self, destination: TextIO, **kwargs: Unpack[ExportDataCsvKwargs]) -> None:
         """Exports the entire dataset into an arbitrary stream.

src/crawlee/storages/_key_value_store.py

@@ -7,12 +7,14 @@
 
 from crawlee import service_container
 from crawlee._utils.docs import docs_group
-from crawlee.base_storage_client._models import KeyValueStoreKeyInfo, KeyValueStoreMetadata
+from crawlee.base_storage_client import BaseStorageClient, KeyValueStoreKeyInfo, KeyValueStoreMetadata
 from crawlee.storages._base_storage import BaseStorage
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
+    from crawlee.configuration import Configuration
+
 T = TypeVar('T')
 
 
@@ -51,9 +53,7 @@ class KeyValueStore(BaseStorage):
     ```
     """
 
-    def __init__(self, id: str, name: str | None) -> None:
-        storage_client = service_container.get_storage_client()
-
+    def __init__(self, id: str, name: str | None, storage_client: BaseStorageClient) -> None:
         self._id = id
         self._name = name
 
@@ -72,14 +72,30 @@ def name(self) -> str | None:
 
     async def get_info(self) -> KeyValueStoreMetadata | None:
         """Get an object containing general information about the key value store."""
-        return await self._resource_client.get()  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.get()
 
     @override
     @classmethod
-    async def open(cls, *, id: str | None = None, name: str | None = None) -> KeyValueStore:
+    async def open(
+        cls,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        configuration: Configuration | None = None,
+        storage_client: BaseStorageClient | None = None,
+    ) -> KeyValueStore:
         from crawlee.storages._creation_management import open_storage
 
-        return await open_storage(storage_class=cls, id=id, name=name)
+        configuration = configuration or service_container.get_configuration()
+        storage_client = storage_client or service_container.get_storage_client()
+
+        return await open_storage(
+            storage_class=cls,
+            id=id,
+            name=name,
+            configuration=configuration,
+            storage_client=storage_client,
+        )
 
     @override
     async def drop(self) -> None:
@@ -142,9 +158,9 @@ async def set_value(
             content_type: Content type of the record.
         """
         if value is None:
-            return await self._resource_client.delete_record(key)  # type: ignore[no-any-return] # Mypy is broken
+            return await self._resource_client.delete_record(key)
 
-        return await self._resource_client.set_record(key, value, content_type)  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.set_record(key, value, content_type)
 
     async def get_public_url(self, key: str) -> str:
         """Get the public URL for the given key.
@@ -155,4 +171,4 @@ async def get_public_url(self, key: str) -> str:
         Returns:
             The public URL for the given key.
         """
-        return await self._resource_client.get_public_url(key)  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.get_public_url(key)

src/crawlee/storages/_request_queue.py

@@ -15,7 +15,7 @@
 from crawlee._utils.lru_cache import LRUCache
 from crawlee._utils.requests import unique_key_to_request_id
 from crawlee._utils.wait import wait_for_all_tasks_for_finish
-from crawlee.base_storage_client._models import ProcessedRequest, RequestQueueMetadata
+from crawlee.base_storage_client import BaseStorageClient, ProcessedRequest, RequestQueueMetadata
 from crawlee.events._types import Event
 from crawlee.storages._base_storage import BaseStorage
 from crawlee.storages._request_provider import RequestProvider
@@ -24,7 +24,7 @@
     from collections.abc import Sequence
 
     from crawlee._request import Request
-
+    from crawlee.configuration import Configuration
 
 logger = getLogger(__name__)
 
@@ -104,10 +104,9 @@ class RequestQueue(BaseStorage, RequestProvider):
     _STORAGE_CONSISTENCY_DELAY = timedelta(seconds=3)
     """Expected delay for storage to achieve consistency, guiding the timing of subsequent read operations."""
 
-    def __init__(self, id: str, name: str | None) -> None:
+    def __init__(self, id: str, name: str | None, storage_client: BaseStorageClient) -> None:
         config = service_container.get_configuration()
         event_manager = service_container.get_event_manager()
-        storage_client = service_container.get_storage_client()
 
         self._id = id
         self._name = name
@@ -148,10 +147,26 @@ def name(self) -> str | None:
 
     @override
     @classmethod
-    async def open(cls, *, id: str | None = None, name: str | None = None) -> RequestQueue:
+    async def open(
+        cls,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        configuration: Configuration | None = None,
+        storage_client: BaseStorageClient | None = None,
+    ) -> RequestQueue:
         from crawlee.storages._creation_management import open_storage
 
-        return await open_storage(storage_class=cls, id=id, name=name)
+        configuration = configuration or service_container.get_configuration()
+        storage_client = storage_client or service_container.get_storage_client()
+
+        return await open_storage(
+            storage_class=cls,
+            id=id,
+            name=name,
+            configuration=configuration,
+            storage_client=storage_client,
+        )
 
     @override
     async def drop(self, *, timeout: timedelta | None = None) -> None:
@@ -204,7 +219,7 @@ async def add_request(
         ):
             self._assumed_total_count += 1
 
-        return processed_request  # type: ignore[no-any-return] # Mypy is broken
+        return processed_request
 
     @override
     async def add_requests_batched(
@@ -260,7 +275,7 @@ async def get_request(self, request_id: str) -> Request | None:
         Returns:
             The retrieved request, or `None`, if it does not exist.
         """
-        return await self._resource_client.get_request(request_id)  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.get_request(request_id)
 
     async def fetch_next_request(self) -> Request | None:
         """Return the next request in the queue to be processed.
@@ -373,7 +388,7 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
             self._assumed_handled_count += 1
 
         self._cache_request(unique_key_to_request_id(request.unique_key), processed_request)
-        return processed_request  # type: ignore[no-any-return] # Mypy is broken
+        return processed_request
 
     async def reclaim_request(
         self,
@@ -417,7 +432,7 @@ async def reclaim_request(
             except Exception as err:
                 logger.debug(f'Failed to delete request lock for request {request.id}', exc_info=err)
 
-        return processed_request  # type: ignore[no-any-return] # Mypy is broken
+        return processed_request
 
     async def is_empty(self) -> bool:
         """Check whether the queue is empty.
@@ -483,7 +498,7 @@ async def is_finished(self) -> bool:
 
     async def get_info(self) -> RequestQueueMetadata | None:
         """Get an object containing general information about the request queue."""
-        return await self._resource_client.get()  # type: ignore[no-any-return] # Mypy is broken
+        return await self._resource_client.get()
 
     @override
     async def get_handled_count(self) -> int:
@@ -658,7 +673,7 @@ async def _prolong_request_lock(self, request_id: str) -> datetime | None:
             )
             return None
         else:
-            return res.lock_expires_at  # type: ignore[no-any-return] # Mypy is broken
+            return res.lock_expires_at
 
     async def _clear_possible_locks(self) -> None:
         self._queue_paused_for_migration = True