fix: storage manager & purging the defaults (#150)

### Description

- Auto-purging of default storages was already implemented. I just
checked it how/if it works, how it works with more instances of
storages, and updated the StorageManager to be singleton-like.

### Related issues

- #87

### Testing

- Covered by unit-tests

Author: vdusek

src/crawlee/memory_storage_client/memory_storage_client.py

@@ -3,6 +3,7 @@
 import asyncio
 import contextlib
 import os
+from logging import getLogger
 from pathlib import Path
 
 import aioshutil
@@ -19,6 +20,8 @@
 from crawlee.memory_storage_client.request_queue_client import RequestQueueClient
 from crawlee.memory_storage_client.request_queue_collection_client import RequestQueueCollectionClient
 
+logger = getLogger(__name__)
+
 
 class MemoryStorageClient(BaseStorageClient):
     """Represents an in-memory storage client for managing datasets, key-value stores, and request queues.
@@ -47,6 +50,7 @@ def __init__(self, configuration: Configuration | None = None) -> None:
         self.datasets_handled: list[DatasetClient] = []
         self.key_value_stores_handled: list[KeyValueStoreClient] = []
         self.request_queues_handled: list[RequestQueueClient] = []
+
         self._purged_on_start = False  # Indicates whether a purge was already performed on this instance.
         self._purge_lock = asyncio.Lock()
 
@@ -134,6 +138,7 @@ def request_queues(self) -> RequestQueueCollectionClient:
     async def purge_on_start(self) -> None:
         # Optimistic, non-blocking check
         if self._purged_on_start is True:
+            logger.debug('Storage was already purged on start.')
             return
 
         async with self._purge_lock:
@@ -142,10 +147,10 @@ async def purge_on_start(self) -> None:
                 # Mypy doesn't understand that the _purged_on_start can change while we're getting the async lock
                 return  # type: ignore[unreachable]
 
-            await self._purge_inner()
+            await self._purge_default_storages()
             self._purged_on_start = True
 
-    async def _purge_inner(self) -> None:
+    async def _purge_default_storages(self) -> None:
         """Cleans up the storage directories, preparing the environment for a new run.
 
         It aims to remove residues from previous executions to avoid data contamination between runs.
@@ -163,21 +168,23 @@ async def _purge_inner(self) -> None:
                     self._TEMPORARY_DIR_NAME
                 ) or key_value_store_folder.name.startswith('__OLD'):
                     await self._batch_remove_files(key_value_store_folder.path)
-                elif key_value_store_folder.name == 'default':
+                elif key_value_store_folder.name == self.default_storage_id:
                     await self._handle_default_key_value_store(key_value_store_folder.path)
 
         # Datasets
         if await ospath.exists(self.datasets_directory):
             dataset_folders = await scandir(self.datasets_directory)
             for dataset_folder in dataset_folders:
-                if dataset_folder.name == 'default' or dataset_folder.name.startswith(self._TEMPORARY_DIR_NAME):
+                if dataset_folder.name == self.default_storage_id or dataset_folder.name.startswith(
+                    self._TEMPORARY_DIR_NAME
+                ):
                     await self._batch_remove_files(dataset_folder.path)
 
         # Request queues
         if await ospath.exists(self.request_queues_directory):
             request_queue_folders = await scandir(self.request_queues_directory)
             for request_queue_folder in request_queue_folders:
-                if request_queue_folder.name == 'default' or request_queue_folder.name.startswith(
+                if request_queue_folder.name == self.default_storage_id or request_queue_folder.name.startswith(
                     self._TEMPORARY_DIR_NAME
                 ):
                     await self._batch_remove_files(request_queue_folder.path)

src/crawlee/storage_client_manager.py

@@ -11,22 +11,11 @@
 class StorageClientManager:
     """A class for managing storage clients."""
 
-    def __init__(
-        self,
-        *,
-        local_client: BaseStorageClient | None = None,
-        cloud_client: BaseStorageClient | None = None,
-    ) -> None:
-        """Create a new instance.
+    _local_client: BaseStorageClient = MemoryStorageClient()
+    _cloud_client: BaseStorageClient | None = None
 
-        Args:
-            local_client: The storage client to be used in the local environment.
-            cloud_client: The storage client to be used in the cloud environment.
-        """
-        self._local_client = local_client or MemoryStorageClient()
-        self._cloud_client = cloud_client
-
-    def get_storage_client(self, *, in_cloud: bool = False) -> BaseStorageClient:
+    @classmethod
+    def get_storage_client(cls, *, in_cloud: bool = False) -> BaseStorageClient:
         """Get the storage client instance for the current environment.
 
         Args:
@@ -36,8 +25,26 @@ def get_storage_client(self, *, in_cloud: bool = False) -> BaseStorageClient:
             The current storage client instance.
         """
         if in_cloud:
-            if self._cloud_client is None:
+            if cls._cloud_client is None:
                 raise RuntimeError('Running in cloud environment, but cloud client was not provided.')
-            return self._cloud_client
+            return cls._cloud_client
 
-        return self._local_client
+        return cls._local_client
+
+    @classmethod
+    def set_cloud_client(cls, cloud_client: BaseStorageClient) -> None:
+        """Set the cloud storage client instance.
+
+        Args:
+            cloud_client: The cloud storage client instance.
+        """
+        cls._cloud_client = cloud_client
+
+    @classmethod
+    def set_local_client(cls, local_client: BaseStorageClient) -> None:
+        """Set the local storage client instance.
+
+        Args:
+            local_client: The local storage client instance.
+        """
+        cls._local_client = local_client

src/crawlee/storages/_creation_management.py

@@ -149,7 +149,7 @@ async def open_storage(
 ) -> Dataset | KeyValueStore | RequestQueue:
     """Open a either a new key-value store or restore existing one and return it."""
     configuration = configuration or Configuration()
-    storage_client = StorageClientManager().get_storage_client(in_cloud=configuration.in_cloud)
+    storage_client = StorageClientManager.get_storage_client(in_cloud=configuration.in_cloud)
 
     # Try to restore the storage from cache by ID
     if name:

src/crawlee/storages/base_storage.py

@@ -27,6 +27,7 @@ def name(self) -> str | None:
     @abstractmethod
     async def open(
         cls,
+        *,
         id: str | None = None,
         name: str | None = None,
         configuration: Configuration | None = None,

src/crawlee/storages/dataset.py

@@ -80,6 +80,7 @@ def name(self) -> str | None:
     @classmethod
     async def open(
         cls,
+        *,
         id: str | None = None,
         name: str | None = None,
         configuration: Configuration | None = None,

src/crawlee/storages/key_value_store.py

@@ -68,6 +68,7 @@ def name(self) -> str | None:
     @classmethod
     async def open(
         cls,
+        *,
         id: str | None = None,
         name: str | None = None,
         configuration: Configuration | None = None,

src/crawlee/storages/request_queue.py

@@ -113,6 +113,7 @@ def name(self) -> str | None:
     @classmethod
     async def open(
         cls,
+        *,
         id: str | None = None,
         name: str | None = None,
         configuration: Configuration | None = None,

tests/unit/conftest.py

@@ -6,6 +6,7 @@
 
 from crawlee.configuration import Configuration
 from crawlee.memory_storage_client import MemoryStorageClient
+from crawlee.storage_client_manager import StorageClientManager
 from crawlee.storages import _creation_management
 
 if TYPE_CHECKING:
@@ -15,6 +16,9 @@
 @pytest.fixture()
 def reset_default_instances(monkeypatch: pytest.MonkeyPatch) -> Callable[[], None]:
     def reset() -> None:
+        StorageClientManager._local_client = MemoryStorageClient()
+        StorageClientManager._cloud_client = None
+
         monkeypatch.setattr(_creation_management, '_cache_dataset_by_id', {})
         monkeypatch.setattr(_creation_management, '_cache_dataset_by_name', {})
         monkeypatch.setattr(_creation_management, '_cache_kvs_by_id', {})

tests/unit/memory_storage_client/test_memory_storage_client.py

@@ -132,7 +132,7 @@ async def test_purge_datasets(tmp_path: Path) -> None:
     assert default_dataset_info.name in folders_before_purge
     assert non_default_dataset_info.name in folders_before_purge
 
-    await ms._purge_inner()
+    await ms._purge_default_storages()
     folders_after_purge = os.listdir(ms.datasets_directory)
     assert default_dataset_info.name not in folders_after_purge
     assert non_default_dataset_info.name in folders_after_purge
@@ -164,7 +164,7 @@ async def test_purge_key_value_stores(tmp_path: Path) -> None:
     assert 'INPUT.json' in default_folder_files_before_purge
     assert 'test.json' in default_folder_files_before_purge
 
-    await ms._purge_inner()
+    await ms._purge_default_storages()
     folders_after_purge = os.listdir(ms.key_value_stores_directory)
     assert default_kvs_info.name in folders_after_purge
     assert non_default_kvs_info.name in folders_after_purge
@@ -189,7 +189,7 @@ async def test_purge_request_queues(tmp_path: Path) -> None:
     folders_before_purge = os.listdir(ms.request_queues_directory)
     assert default_rq_info.name in folders_before_purge
     assert non_default_rq_info.name in folders_before_purge
-    await ms._purge_inner()
+    await ms._purge_default_storages()
     folders_after_purge = os.listdir(ms.request_queues_directory)
     assert default_rq_info.name not in folders_after_purge
     assert non_default_rq_info.name in folders_after_purge

tests/unit/test_storage_client_manager.py

@@ -10,26 +10,24 @@
 
 
 def test_returns_memory_storage_client_as_default() -> None:
-    manager = StorageClientManager()
-    storage_client = manager.get_storage_client()
+    storage_client = StorageClientManager.get_storage_client()
     assert isinstance(storage_client, MemoryStorageClient), 'Should return the memory storage client by default'
 
 
 def test_returns_provided_local_client_for_non_cloud_environment() -> None:
     local_client = Mock(spec=BaseStorageClient)
-    manager = StorageClientManager(local_client=local_client)
-    storage_client = manager.get_storage_client()
+    StorageClientManager.set_local_client(local_client)
+    storage_client = StorageClientManager.get_storage_client()
     assert storage_client == local_client, 'Should return the local client when not in cloud'
 
 
 def test_returns_provided_cloud_client_for_cloud_environment() -> None:
     cloud_client = Mock(spec=BaseStorageClient)
-    manager = StorageClientManager(cloud_client=cloud_client)
-    storage_client = manager.get_storage_client(in_cloud=True)
+    StorageClientManager.set_cloud_client(cloud_client)
+    storage_client = StorageClientManager.get_storage_client(in_cloud=True)
     assert storage_client == cloud_client, 'Should return the cloud client when in cloud'
 
 
 def test_raises_error_when_no_cloud_client_provided() -> None:
-    manager = StorageClientManager()
     with pytest.raises(RuntimeError, match='cloud client was not provided'):
-        manager.get_storage_client(in_cloud=True)
+        StorageClientManager.get_storage_client(in_cloud=True)