move retry logic to storage clients

Author: Mantisus

src/crawlee/_utils/retry.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from functools import wraps
+from typing import TYPE_CHECKING, ParamSpec, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
+def retry_on_error(
+    *exception_types: type[Exception],
+    max_attempts: int = 3,
+    base_delay: timedelta = timedelta(milliseconds=500),
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+    """Retry an async function with exponential backoff on specified exceptions.
+
+    Args:
+        *exception_types: Exception types to catch and retry on.
+        max_attempts: Maximum number of attempts including the first one.
+        base_delay: Base delay between retries; doubles on each subsequent attempt.
+    """
+
+    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+
+        if max_attempts < 1:
+            raise ValueError('max_attempts must be at least 1')
+
+        @wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            base_delay_seconds = base_delay.total_seconds()
+            for attempt in range(max_attempts):
+                try:
+                    return await func(*args, **kwargs)
+                except Exception as exc:  # noqa: PERF203
+                    if not isinstance(exc, exception_types) or attempt >= max_attempts - 1:
+                        raise
+                    await asyncio.sleep(base_delay_seconds * (2**attempt))
+            raise RuntimeError('Unreachable')
+
+        return wrapper
+
+    return decorator

src/crawlee/errors.py

@@ -18,7 +18,6 @@
     'RequestHandlerError',
     'ServiceConflictError',
     'SessionError',
-    'StorageWriteError',
     'UserDefinedErrorHandlerError',
 ]
 
@@ -117,12 +116,3 @@ 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,7 +15,6 @@
 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
 
@@ -223,24 +222,21 @@ async def purge(self) -> None:
     @override
     async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         async with self._lock:
-            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 = self._metadata.item_count
+            if isinstance(data, list):
+                for item in data:
                     new_item_count += 1
-                    await self._push_item(data, new_item_count)
+                    await self._push_item(item, new_item_count)
+            else:
+                new_item_count += 1
+                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,
-                )
-            except OSError as e:
-                raise StorageWriteError(e) from e
+            # 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,
+            )
 
     @override
     async def get_data(

src/crawlee/storage_clients/_file_system/_key_value_store_client.py

@@ -17,7 +17,6 @@
 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
 
@@ -329,20 +328,17 @@ 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:
-            try:
-                # Ensure the key-value store directory exists.
-                await asyncio.to_thread(self.path_to_kvs.mkdir, parents=True, exist_ok=True)
+            # 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)
-            except OSError as e:
-                raise StorageWriteError(e) from e
+            # Update the KVS metadata to record the access and modification.
+            await self._update_metadata(update_accessed_at=True, update_modified_at=True)
 
     @override
     async def delete_value(self, *, key: str) -> None:

src/crawlee/storage_clients/_redis/_dataset_client.py

@@ -6,7 +6,7 @@
 from redis.exceptions import RedisError
 from typing_extensions import NotRequired, override
 
-from crawlee.errors import StorageWriteError
+from crawlee._utils.retry import retry_on_error
 from crawlee.storage_clients._base import DatasetClient
 from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata
 
@@ -104,14 +104,17 @@ async def open(
             instance_kwargs={},
         )
 
+    @retry_on_error(RedisError)
     @override
     async def get_metadata(self) -> DatasetMetadata:
         return await self._get_metadata(DatasetMetadata)
 
+    @retry_on_error(RedisError)
     @override
     async def drop(self) -> None:
         await self._drop(extra_keys=[self._items_key])
 
+    @retry_on_error(RedisError)
     @override
     async def purge(self) -> None:
         await self._purge(
@@ -121,24 +124,22 @@ async def purge(self) -> None:
             ),
         )
 
+    @retry_on_error(RedisError)
     @override
     async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         if isinstance(data, dict):
             data = [data]
 
-        try:
-            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)
-                    ),
-                )
-
-        except RedisError as e:
-            raise StorageWriteError(e) from e
+        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)
+                ),
+            )
 
+    @retry_on_error(RedisError)
     @override
     async def get_data(
         self,

src/crawlee/storage_clients/_redis/_key_value_store_client.py

@@ -8,7 +8,7 @@
 from typing_extensions import override
 
 from crawlee._utils.file import infer_mime_type
-from crawlee.errors import StorageWriteError
+from crawlee._utils.retry import retry_on_error
 from crawlee.storage_clients._base import KeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
 
@@ -102,21 +102,25 @@ async def open(
             instance_kwargs={},
         )
 
+    @retry_on_error(RedisError)
     @override
     async def get_metadata(self) -> KeyValueStoreMetadata:
         return await self._get_metadata(KeyValueStoreMetadata)
 
+    @retry_on_error(RedisError)
     @override
     async def drop(self) -> None:
         await self._drop(extra_keys=[self._items_key, self._metadata_items_key])
 
+    @retry_on_error(RedisError)
     @override
     async def purge(self) -> None:
         await self._purge(
             extra_keys=[self._items_key, self._metadata_items_key],
             metadata_kwargs=MetadataUpdateParams(update_accessed_at=True, update_modified_at=True),
         )
 
+    @retry_on_error(RedisError)
     @override
     async def set_value(self, *, key: str, value: Any, content_type: str | None = None) -> None:
         # Special handling for None values
@@ -143,25 +147,20 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
             content_type=content_type,
             size=size,
         )
-        try:
-            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(),
-                    )
-                )
-                await self._update_metadata(
-                    pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True)
+        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(),
                 )
+            )
+            await self._update_metadata(pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True))
 
-        except RedisError as e:
-            raise StorageWriteError(e) from e
-
+    @retry_on_error(RedisError)
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:
         serialized_metadata_item = await await_redis_response(self._redis.hget(self._metadata_items_key, key))
@@ -207,6 +206,7 @@ async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:
 
         return KeyValueStoreRecord(value=value, **metadata_item.model_dump())
 
+    @retry_on_error(RedisError)
     @override
     async def delete_value(self, *, key: str) -> None:
         async with self._get_pipeline() as pipe:
@@ -258,6 +258,7 @@ async def iterate_keys(
     async def get_public_url(self, *, key: str) -> str:
         raise NotImplementedError('Public URLs are not supported for memory key-value stores.')
 
+    @retry_on_error(RedisError)
     @override
     async def record_exists(self, *, key: str) -> bool:
         async with self._get_pipeline(with_execute=False) as pipe:

src/crawlee/storage_clients/_redis/_request_queue_client.py

@@ -6,10 +6,12 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, Literal
 
+from redis.exceptions import RedisError
 from typing_extensions import NotRequired, override
 
 from crawlee import Request
 from crawlee._utils.crypto import crypto_random_object_id
+from crawlee._utils.retry import retry_on_error
 from crawlee.storage_clients._base import RequestQueueClient
 from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
 
@@ -207,10 +209,12 @@ async def open(
             instance_kwargs={'dedup_strategy': dedup_strategy, 'bloom_error_rate': bloom_error_rate},
         )
 
+    @retry_on_error(RedisError)
     @override
     async def get_metadata(self) -> RequestQueueMetadata:
         return await self._get_metadata(RequestQueueMetadata)
 
+    @retry_on_error(RedisError)
     @override
     async def drop(self) -> None:
         if self._dedup_strategy == 'bloom':
@@ -222,6 +226,7 @@ async def drop(self) -> None:
         extra_keys.extend([self._queue_key, self._data_key, self._in_progress_key])
         await self._drop(extra_keys=extra_keys)
 
+    @retry_on_error(RedisError)
     @override
     async def purge(self) -> None:
         if self._dedup_strategy == 'bloom':
@@ -349,6 +354,7 @@ async def add_batch_of_requests(
             unprocessed_requests=[],
         )
 
+    @retry_on_error(RedisError)
     @override
     async def fetch_next_request(self) -> Request | None:
         if self._pending_fetch_cache:
@@ -377,6 +383,7 @@ async def fetch_next_request(self) -> Request | None:
 
         return requests[0]
 
+    @retry_on_error(RedisError)
     @override
     async def get_request(self, unique_key: str) -> Request | None:
         request_data = await await_redis_response(self._redis.hget(self._data_key, unique_key))
@@ -386,6 +393,7 @@ async def get_request(self, unique_key: str) -> Request | None:
 
         return None
 
+    @retry_on_error(RedisError)
     @override
     async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
         # Check if the request is in progress.
@@ -424,6 +432,7 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
             was_already_handled=True,
         )
 
+    @retry_on_error(RedisError)
     @override
     async def reclaim_request(
         self,
@@ -469,6 +478,7 @@ async def reclaim_request(
             was_already_handled=False,
         )
 
+    @retry_on_error(RedisError)
     @override
     async def is_empty(self) -> bool:
         """Check if the queue is empty.

src/crawlee/storage_clients/_sql/_client_mixin.py

@@ -205,6 +205,7 @@ async def get_session(self, *, with_simple_commit: bool = False) -> AsyncIterato
                 except SQLAlchemyError as e:
                     logger.warning(f'Error occurred during session transaction: {e}')
                     await session.rollback()
+                    raise
             else:
                 yield session