Draft of minimizing side effects

Explicit kvs to RecoverableState

Author: Pijukatel

src/crawlee/_utils/recoverable_state.py

@@ -4,6 +4,7 @@
 
 from pydantic import BaseModel
 
+from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
 from crawlee.events._types import Event, EventPersistStateData
 
 if TYPE_CHECKING:
@@ -37,6 +38,7 @@ def __init__(
         persistence_enabled: Literal[True, False, 'explicit_only'] = False,
         persist_state_kvs_name: str | None = None,
         persist_state_kvs_id: str | None = None,
+        persist_state_kvs: KeyValueStore | None = None,
         logger: logging.Logger,
     ) -> None:
         """Initialize a new recoverable state object.
@@ -51,16 +53,28 @@ def __init__(
                 If neither a name nor and id are supplied, the default store will be used.
             persist_state_kvs_id: The identifier of the KeyValueStore to use for persistence.
                 If neither a name nor and id are supplied, the default store will be used.
+            persist_state_kvs: KeyValueStore to use for persistence. If not provided, a system-wide KeyValueStore will
+                be used, based on service locator configuration.
             logger: A logger instance for logging operations related to state persistence
         """
+        raise_if_too_many_kwargs(persist_state_kvs_name=persist_state_kvs_name,
+                                 persist_state_kvs_id=persist_state_kvs_id,
+                                 key_value_store=persist_state_kvs)
+        if not persist_state_kvs:
+            logger.debug(
+                'No explicit key_value_store set for recoverable state. Recovery will use a system-wide KeyValueStore '
+                'based on service_locator configuration, potentially calling service_locator.set_storage_client in the '
+                'process. It is recommended to initialize RecoverableState with explicit key_value_store to avoid '
+                'global side effects.')
+
         self._default_state = default_state
         self._state_type: type[TStateModel] = self._default_state.__class__
         self._state: TStateModel | None = None
         self._persistence_enabled = persistence_enabled
         self._persist_state_key = persist_state_key
         self._persist_state_kvs_name = persist_state_kvs_name
         self._persist_state_kvs_id = persist_state_kvs_id
-        self._key_value_store: 'KeyValueStore | None' = None  # noqa: UP037
+        self._key_value_store: KeyValueStore | None = persist_state_kvs
         self._log = logger
 
     async def initialize(self) -> TStateModel:
@@ -79,9 +93,10 @@ async def initialize(self) -> TStateModel:
         # Import here to avoid circular imports.
         from crawlee.storages._key_value_store import KeyValueStore  # noqa: PLC0415
 
-        self._key_value_store = await KeyValueStore.open(
-            name=self._persist_state_kvs_name, id=self._persist_state_kvs_id
-        )
+        if not self._key_value_store:
+            self._key_value_store = await KeyValueStore.open(
+                name=self._persist_state_kvs_name, id=self._persist_state_kvs_id
+            )
 
         await self._load_saved_state()
 

src/crawlee/storage_clients/_file_system/_request_queue_client.py

@@ -92,6 +92,7 @@ def __init__(
         metadata: RequestQueueMetadata,
         path_to_rq: Path,
         lock: asyncio.Lock,
+        recoverable_state: RecoverableState[RequestQueueState],
     ) -> None:
         """Initialize a new instance.
 
@@ -114,12 +115,7 @@ def __init__(
         self._is_empty_cache: bool | None = None
         """Cache for is_empty result: None means unknown, True/False is cached state."""
 
-        self._state = RecoverableState[RequestQueueState](
-            default_state=RequestQueueState(),
-            persist_state_key=f'__RQ_STATE_{self._metadata.id}',
-            persistence_enabled=True,
-            logger=logger,
-        )
+        self._state = recoverable_state
         """Recoverable state to maintain request ordering, in-progress status, and handled status."""
 
     @override
@@ -136,6 +132,19 @@ def path_to_metadata(self) -> Path:
         """The full path to the request queue metadata file."""
         return self.path_to_rq / METADATA_FILENAME
 
+    @classmethod
+    async def _create_recoverable_state(cls, id: str, configuration: Configuration) -> RecoverableState:
+        from crawlee.storage_clients import FileSystemStorageClient
+        from crawlee.storages import KeyValueStore
+        kvs = await KeyValueStore.open(storage_client=FileSystemStorageClient(),configuration=configuration)
+        return RecoverableState[RequestQueueState](
+            default_state=RequestQueueState(),
+            persist_state_key=f'__RQ_STATE_{id}',
+            persist_state_kvs=kvs,
+            persistence_enabled=True,
+            logger=logger,
+        )
+
     @classmethod
     async def open(
         cls,
@@ -194,6 +203,9 @@ async def open(
                                 metadata=metadata,
                                 path_to_rq=rq_base_path / rq_dir,
                                 lock=asyncio.Lock(),
+                                recoverable_state=await cls._create_recoverable_state(id=id,
+                                                                                      configuration=configuration),
+
                             )
                             await client._state.initialize()
                             await client._discover_existing_requests()
@@ -230,6 +242,7 @@ async def open(
                     metadata=metadata,
                     path_to_rq=path_to_rq,
                     lock=asyncio.Lock(),
+                    recoverable_state=await cls._create_recoverable_state(id=metadata.id, configuration=configuration),
                 )
 
                 await client._state.initialize()
@@ -254,6 +267,7 @@ async def open(
                     metadata=metadata,
                     path_to_rq=path_to_rq,
                     lock=asyncio.Lock(),
+                    recoverable_state=await cls._create_recoverable_state(id=metadata.id, configuration=configuration),
                 )
                 await client._state.initialize()
                 await client._update_metadata()

src/crawlee/storages/_key_value_store.py

@@ -278,13 +278,9 @@ async def get_auto_saved_value(
             if key in cache:
                 return cache[key].current_value.root
 
-            cache[key] = recoverable_state = RecoverableState(
-                default_state=AutosavedValue(default_value),
-                persistence_enabled=True,
-                persist_state_kvs_id=self.id,
-                persist_state_key=key,
-                logger=logger,
-            )
+            cache[key] = recoverable_state = RecoverableState(default_state=AutosavedValue(default_value),
+                                                              persist_state_key=key, persistence_enabled=True,
+                                                              persist_state_kvs=self, logger=logger)
 
             await recoverable_state.initialize()
 

tests/unit/conftest.py

@@ -69,6 +69,9 @@ def _prepare_test_env() -> None:
         # Verify that the test environment was set up correctly.
         assert os.environ.get('CRAWLEE_STORAGE_DIR') == str(tmp_path)
 
+        # Clear global cache of autosaved values
+        KeyValueStore._autosaved_values = {}
+
     return _prepare_test_env
 
 

tests/unit/storage_clients/_file_system/test_fs_rq_client.py

@@ -6,9 +6,9 @@
 
 import pytest
 
-from crawlee import Request
+from crawlee import Request, service_locator
 from crawlee.configuration import Configuration
-from crawlee.storage_clients import FileSystemStorageClient
+from crawlee.storage_clients import FileSystemStorageClient, MemoryStorageClient
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
@@ -77,6 +77,11 @@ async def test_request_file_persistence(rq_client: FileSystemRequestQueueClient)
             assert 'url' in request_data
             assert request_data['url'].startswith('https://example.com/')
 
+async def test_opening_rq_does_not_have_side_effect_on_service_locator(
+    rq_client: FileSystemRequestQueueClient # noqa: ARG001
+) -> None:
+    service_locator.set_storage_client(MemoryStorageClient())
+
 
 async def test_drop_removes_directory(rq_client: FileSystemRequestQueueClient) -> None:
     """Test that drop removes the entire RQ directory from disk."""

tests/unit/storages/test_key_value_store.py

@@ -10,13 +10,15 @@
 
 from crawlee import service_locator
 from crawlee.configuration import Configuration
+from crawlee.storage_clients import FileSystemStorageClient, MemoryStorageClient, SqlStorageClient, StorageClient
 from crawlee.storages import KeyValueStore
 from crawlee.storages._storage_instance_manager import StorageInstanceManager
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
+    from pathlib import Path
+
 
-    from crawlee.storage_clients import StorageClient
 
 
 @pytest.fixture
@@ -1063,3 +1065,33 @@ async def test_name_default_not_allowed(storage_client: StorageClient) -> None:
         f'it is reserved for default alias.',
     ):
         await KeyValueStore.open(name=StorageInstanceManager._DEFAULT_STORAGE_ALIAS, storage_client=storage_client)
+
+@pytest.mark.parametrize('tested_storage_client', [
+    pytest.param(MemoryStorageClient(), id='tested=MemoryStorageClient'),
+    pytest.param(FileSystemStorageClient(), id='tested=FileSystemStorageClient'),
+    pytest.param(SqlStorageClient(), id='tested=SqlStorageClient'),
+                  ])
+@pytest.mark.parametrize('global_storage_client', [
+    pytest.param(MemoryStorageClient(), id='global=MemoryStorageClient'),
+    pytest.param(FileSystemStorageClient(), id='global=FileSystemStorageClient'),
+    pytest.param(SqlStorageClient(), id='global=SqlStorageClient'),
+                  ])
+async def test_get_auto_saved_value_various_global_clients(tmp_path: Path, tested_storage_client: StorageClient,
+                                    global_storage_client:StorageClient) -> None:
+    """Ensure that persistence is working for all clients regardless of what is set in service locator."""
+    service_locator.set_configuration(Configuration(
+        crawlee_storage_dir=str(tmp_path),  # type: ignore[call-arg]
+        purge_on_start=True,
+    ))
+    service_locator.set_storage_client(global_storage_client)
+
+    kvs = await KeyValueStore.open(storage_client=tested_storage_client)
+    values_kvs = {'key': 'some_value'}
+    test_key = 'test_key'
+
+    autosaved_value_kvs1 = await kvs.get_auto_saved_value(test_key)
+    assert autosaved_value_kvs1 == {}
+    autosaved_value_kvs1.update(values_kvs)
+    await kvs.persist_autosaved_values()
+
+    assert await kvs.get_value(test_key) == autosaved_value_kvs1