Add retry logic for KeyValueStore.set_value and Dataset.push_data

Author: Mantisus

src/crawlee/errors.py

@@ -18,6 +18,7 @@
     'RequestHandlerError',
     'ServiceConflictError',
     'SessionError',
+    'StorageWriteError',
     'UserDefinedErrorHandlerError',
 ]
 
@@ -116,3 +117,12 @@ class ContextPipelineInterruptedError(Exception):
 @docs_group('Errors')
 class RequestCollisionError(Exception):
     """Raised when a request cannot be processed due to a conflict with required resources."""
+
+
+@docs_group('Errors')
+class StorageWriteError(Exception):
+    """Raised when a write operation to a storage fails."""
+
+    def __init__(self, cause: Exception) -> None:
+        super().__init__(str(cause))
+        self.cause = cause

src/crawlee/storage_clients/_file_system/_dataset_client.py

@@ -15,6 +15,7 @@
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee._utils.file import atomic_write, json_dumps
 from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -222,21 +223,24 @@ async def purge(self) -> None:
     @override
     async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         async with self._lock:
-            new_item_count = self._metadata.item_count
-            if isinstance(data, list):
-                for item in data:
+            try:
+                new_item_count = self._metadata.item_count
+                if isinstance(data, list):
+                    for item in data:
+                        new_item_count += 1
+                        await self._push_item(item, new_item_count)
+                else:
                     new_item_count += 1
-                    await self._push_item(item, new_item_count)
-            else:
-                new_item_count += 1
-                await self._push_item(data, new_item_count)
+                    await self._push_item(data, new_item_count)
 
-            # now update metadata under the same lock
-            await self._update_metadata(
-                update_accessed_at=True,
-                update_modified_at=True,
-                new_item_count=new_item_count,
-            )
+                # now update metadata under the same lock
+                await self._update_metadata(
+                    update_accessed_at=True,
+                    update_modified_at=True,
+                    new_item_count=new_item_count,
+                )
+            except OSError as e:
+                raise StorageWriteError(e) from e
 
     @override
     async def get_data(

src/crawlee/storage_clients/_file_system/_key_value_store_client.py

@@ -17,6 +17,7 @@
 from crawlee._utils.crypto import crypto_random_object_id
 from crawlee._utils.file import atomic_write, infer_mime_type, json_dumps
 from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
 
@@ -328,17 +329,20 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
         record_metadata_content = await json_dumps(record_metadata.model_dump())
 
         async with self._lock:
-            # Ensure the key-value store directory exists.
-            await asyncio.to_thread(self.path_to_kvs.mkdir, parents=True, exist_ok=True)
+            try:
+                # Ensure the key-value store directory exists.
+                await asyncio.to_thread(self.path_to_kvs.mkdir, parents=True, exist_ok=True)
 
-            # Write the value to the file.
-            await atomic_write(record_path, value_bytes)
+                # Write the value to the file.
+                await atomic_write(record_path, value_bytes)
 
-            # Write the record metadata to the file.
-            await atomic_write(record_metadata_filepath, record_metadata_content)
+                # Write the record metadata to the file.
+                await atomic_write(record_metadata_filepath, record_metadata_content)
 
-            # Update the KVS metadata to record the access and modification.
-            await self._update_metadata(update_accessed_at=True, update_modified_at=True)
+                # Update the KVS metadata to record the access and modification.
+                await self._update_metadata(update_accessed_at=True, update_modified_at=True)
+            except OSError as e:
+                raise StorageWriteError(e) from e
 
     @override
     async def delete_value(self, *, key: str) -> None:

src/crawlee/storage_clients/_redis/_dataset_client.py

@@ -3,8 +3,10 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, cast
 
+from redis.exceptions import RedisError
 from typing_extensions import NotRequired, override
 
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -125,13 +127,17 @@ async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
             data = [data]
 
         async with self._get_pipeline() as pipe:
-            pipe.json().arrappend(self._items_key, '$', *data)
-            await self._update_metadata(
-                pipe,
-                **_DatasetMetadataUpdateParams(
-                    update_accessed_at=True, update_modified_at=True, delta_item_count=len(data)
-                ),
-            )
+            try:
+                pipe.json().arrappend(self._items_key, '$', *data)
+                await self._update_metadata(
+                    pipe,
+                    **_DatasetMetadataUpdateParams(
+                        update_accessed_at=True, update_modified_at=True, delta_item_count=len(data)
+                    ),
+                )
+
+            except RedisError as e:
+                raise StorageWriteError(e) from e
 
     @override
     async def get_data(

src/crawlee/storage_clients/_redis/_key_value_store_client.py

@@ -4,9 +4,11 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 
+from redis.exceptions import RedisError
 from typing_extensions import override
 
 from crawlee._utils.file import infer_mime_type
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
 
@@ -143,17 +145,23 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
         )
 
         async with self._get_pipeline() as pipe:
-            # redis-py typing issue
-            await await_redis_response(pipe.hset(self._items_key, key, value_bytes))  # ty: ignore[invalid-argument-type]
-
-            await await_redis_response(
-                pipe.hset(
-                    self._metadata_items_key,
-                    key,
-                    item_metadata.model_dump_json(),
+            try:
+                # redis-py typing issue
+                await await_redis_response(pipe.hset(self._items_key, key, value_bytes))  # ty: ignore[invalid-argument-type]
+
+                await await_redis_response(
+                    pipe.hset(
+                        self._metadata_items_key,
+                        key,
+                        item_metadata.model_dump_json(),
+                    )
                 )
-            )
-            await self._update_metadata(pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True))
+                await self._update_metadata(
+                    pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True)
+                )
+
+            except RedisError as e:
+                raise StorageWriteError(e) from e
 
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:

src/crawlee/storage_clients/_sql/_dataset_client.py

@@ -6,8 +6,10 @@
 
 from sqlalchemy import Select, insert, select
 from sqlalchemy import func as sql_func
+from sqlalchemy.exc import SQLAlchemyError
 from typing_extensions import Self, override
 
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -145,10 +147,15 @@ async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         db_items = [{'dataset_id': self._id, 'data': item} for item in data]
         stmt = insert(self._ITEM_TABLE).values(db_items)
 
-        async with self.get_session(with_simple_commit=True) as session:
-            await session.execute(stmt)
+        async with self.get_session() as session:
+            try:
+                await session.execute(stmt)
 
-            await self._add_buffer_record(session, update_modified_at=True, delta_item_count=len(data))
+                await self._add_buffer_record(session, update_modified_at=True, delta_item_count=len(data))
+                await session.commit()
+            except SQLAlchemyError as e:
+                await session.rollback()
+                raise StorageWriteError(e) from e
 
     @override
     async def get_data(

src/crawlee/storage_clients/_sql/_key_value_store_client.py

@@ -7,9 +7,11 @@
 
 from sqlalchemy import CursorResult, delete, select
 from sqlalchemy import func as sql_func
+from sqlalchemy.exc import SQLAlchemyError
 from typing_extensions import Self, override
 
 from crawlee._utils.file import infer_mime_type
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import (
     KeyValueStoreMetadata,
@@ -175,10 +177,15 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
             conflict_cols=['key_value_store_id', 'key'],
         )
 
-        async with self.get_session(with_simple_commit=True) as session:
-            await session.execute(upsert_stmt)
+        async with self.get_session() as session:
+            try:
+                await session.execute(upsert_stmt)
 
-            await self._add_buffer_record(session, update_modified_at=True)
+                await self._add_buffer_record(session, update_modified_at=True)
+                await session.commit()
+            except SQLAlchemyError as e:
+                await session.rollback()
+                raise StorageWriteError(e) from e
 
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:

src/crawlee/storage_clients/_sql/_request_queue_client.py

@@ -483,6 +483,9 @@ async def fetch_next_request(self) -> Request | None:
 
             await self._add_buffer_record(session)
 
+        if not requests_db:
+            return None
+
         requests = [Request.model_validate_json(r.data) for r in requests_db if r.request_id in blocked_ids]
 
         if not requests:

src/crawlee/storages/_dataset.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import logging
+from datetime import timedelta
 from io import StringIO
 from typing import TYPE_CHECKING, overload
 
@@ -9,6 +11,7 @@
 from crawlee import service_locator
 from crawlee._utils.docs import docs_group
 from crawlee._utils.file import export_csv_to_stream, export_json_to_stream
+from crawlee.errors import StorageWriteError
 
 from ._base import Storage
 from ._key_value_store import KeyValueStore
@@ -134,7 +137,13 @@ async def drop(self) -> None:
     async def purge(self) -> None:
         await self._client.purge()
 
-    async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
+    async def push_data(
+        self,
+        data: list[dict[str, Any]] | dict[str, Any],
+        *,
+        max_attempts: int = 5,
+        wait_time_between_retries: timedelta = timedelta(seconds=1),
+    ) -> None:
         """Store an object or an array of objects to the dataset.
 
         The size of the data is limited by the receiving API and therefore `push_data()` will only
@@ -144,8 +153,26 @@ async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         Args:
             data: A JSON serializable data structure to be stored in the dataset. The JSON representation
                 of each item must be smaller than 9MB.
+            max_attempts: The maximum number of attempts to push data in case of failure.
+            wait_time_between_retries: The time to wait between retries in case of failure.
         """
-        await self._client.push_data(data=data)
+        if max_attempts < 1:
+            raise ValueError('max_attempts must be at least 1')
+
+        wait_time_between_retries_seconds = wait_time_between_retries.total_seconds()
+        last_exception: StorageWriteError | None = None
+
+        for attempt in range(max_attempts):
+            try:
+                await self._client.push_data(data=data)
+                break
+            except StorageWriteError as e:
+                last_exception = e
+                if attempt < max_attempts - 1:
+                    await asyncio.sleep(wait_time_between_retries_seconds)
+        else:
+            if last_exception:
+                logger.warning(f'Failed to push data after {max_attempts} attempts with error: {last_exception.cause}')
 
     async def get_data(
         self,

src/crawlee/storages/_key_value_store.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from collections.abc import AsyncIterator
+from datetime import timedelta
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, overload
 
@@ -12,6 +13,7 @@
 from crawlee._types import JsonSerializable  # noqa: TC001
 from crawlee._utils.docs import docs_group
 from crawlee._utils.recoverable_state import RecoverableState
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients.models import KeyValueStoreMetadata
 
 from ._base import Storage
@@ -175,15 +177,39 @@ async def set_value(
         key: str,
         value: Any,
         content_type: str | None = None,
+        *,
+        max_attempts: int = 5,
+        wait_time_between_retries: timedelta = timedelta(seconds=1),
     ) -> None:
         """Set a value in the KVS.
 
         Args:
             key: Key of the record to set.
             value: Value to set.
             content_type: The MIME content type string.
+            max_attempts: The maximum number of attempts to set the value in case of failure.
+            wait_time_between_retries: Time to wait between retries.
         """
-        await self._client.set_value(key=key, value=value, content_type=content_type)
+        if max_attempts < 1:
+            raise ValueError('max_attempts must be at least 1')
+
+        wait_time_between_retries_seconds = wait_time_between_retries.total_seconds()
+        last_exception: StorageWriteError | None = None
+
+        for attempt in range(max_attempts):
+            try:
+                await self._client.set_value(key=key, value=value, content_type=content_type)
+                break
+            except StorageWriteError as e:
+                last_exception = e
+                if attempt < max_attempts - 1:
+                    await asyncio.sleep(wait_time_between_retries_seconds)
+        else:
+            if last_exception:
+                logger.warning(
+                    f'Failed to set value for key "{key}" after {max_attempts} attempts '
+                    f'with error: {last_exception.cause}'
+                )
 
     async def delete_value(self, key: str) -> None:
         """Delete a value from the KVS.