crawlee 0.6.13b31__py3-none-any.whl → 1.1.1b1__py3-none-any.whl

crawlee/storage_clients/_redis/_storage_client.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import warnings
+from typing import Literal
+
+from redis.asyncio import Redis
+from typing_extensions import override
+
+from crawlee._utils.docs import docs_group
+from crawlee.configuration import Configuration
+from crawlee.storage_clients._base import StorageClient
+
+from ._dataset_client import RedisDatasetClient
+from ._key_value_store_client import RedisKeyValueStoreClient
+from ._request_queue_client import RedisRequestQueueClient
+
+
+@docs_group('Storage clients')
+class RedisStorageClient(StorageClient):
+    """Redis implementation of the storage client.
+
+    This storage client provides access to datasets, key-value stores, and request queues that persist data
+    to a Redis database v8.0+. Each storage type uses Redis-specific data structures and key patterns for
+    efficient storage and retrieval.
+
+    The client accepts either a Redis connection string or a pre-configured Redis client instance.
+    Exactly one of these parameters must be provided during initialization.
+
+    Storage types use the following Redis data structures:
+    - **Datasets**: Redis JSON arrays for item storage with metadata in JSON objects
+    - **Key-value stores**: Redis hashes for key-value pairs with separate metadata storage
+    - **Request queues**: Redis lists for FIFO queuing, hashes for request data and in-progress tracking,
+      and Bloom filters for request deduplication
+
+    Warning:
+        This is an experimental feature. The behavior and interface may change in future versions.
+    """
+
+    def __init__(
+        self,
+        *,
+        connection_string: str | None = None,
+        redis: Redis | None = None,
+        queue_dedup_strategy: Literal['default', 'bloom'] = 'default',
+        queue_bloom_error_rate: float = 1e-7,
+    ) -> None:
+        """Initialize the Redis storage client.
+
+        Args:
+            connection_string: Redis connection string (e.g., "redis://localhost:6379").
+                Supports standard Redis URL format with optional database selection.
+            redis: Pre-configured Redis client instance.
+            queue_dedup_strategy: Strategy for request queue deduplication. Options are:
+                - 'default': Uses Redis sets for exact deduplication.
+                - 'bloom': Uses Redis Bloom filters for probabilistic deduplication with lower memory usage. When using
+                    this approach, approximately 1 in 1e-7 requests will be falsely considered duplicate.
+            queue_bloom_error_rate: Desired false positive rate for Bloom filter deduplication. Only relevant if
+                `queue_dedup_strategy` is set to 'bloom'.
+        """
+        match (redis, connection_string):
+            case (None, None):
+                raise ValueError('Either redis or connection_string must be provided.')
+            case (Redis(), None):
+                self._redis = redis
+            case (None, str()):
+                self._redis = Redis.from_url(connection_string)
+            case (Redis(), str()):
+                raise ValueError('Either redis or connection_string must be provided, not both.')
+
+        self._queue_dedup_strategy = queue_dedup_strategy
+        self._queue_bloom_error_rate = queue_bloom_error_rate
+
+        # Call the notification only once
+        warnings.warn(
+            (
+                'RedisStorageClient is experimental and its API, behavior, and key structure may change in future '
+                'releases.'
+            ),
+            category=UserWarning,
+            stacklevel=2,
+        )
+
+    @override
+    async def create_dataset_client(
+        self,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        alias: str | None = None,
+        configuration: Configuration | None = None,
+    ) -> RedisDatasetClient:
+        configuration = configuration or Configuration.get_global_configuration()
+
+        client = await RedisDatasetClient.open(
+            id=id,
+            name=name,
+            alias=alias,
+            redis=self._redis,
+        )
+
+        await self._purge_if_needed(client, configuration)
+        return client
+
+    @override
+    async def create_kvs_client(
+        self,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        alias: str | None = None,
+        configuration: Configuration | None = None,
+    ) -> RedisKeyValueStoreClient:
+        configuration = configuration or Configuration.get_global_configuration()
+
+        client = await RedisKeyValueStoreClient.open(
+            id=id,
+            name=name,
+            alias=alias,
+            redis=self._redis,
+        )
+
+        await self._purge_if_needed(client, configuration)
+        return client
+
+    @override
+    async def create_rq_client(
+        self,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        alias: str | None = None,
+        configuration: Configuration | None = None,
+    ) -> RedisRequestQueueClient:
+        configuration = configuration or Configuration.get_global_configuration()
+
+        client = await RedisRequestQueueClient.open(
+            id=id,
+            name=name,
+            alias=alias,
+            redis=self._redis,
+            dedup_strategy=self._queue_dedup_strategy,
+            bloom_error_rate=self._queue_bloom_error_rate,
+        )
+
+        await self._purge_if_needed(client, configuration)
+        return client

crawlee/storage_clients/_redis/_utils.py

@@ -0,0 +1,23 @@
+from collections.abc import Awaitable
+from pathlib import Path
+from typing import TypeVar, overload
+
+T = TypeVar('T')
+
+
+@overload
+async def await_redis_response(response: Awaitable[T]) -> T: ...
+@overload
+async def await_redis_response(response: T) -> T: ...
+
+
+async def await_redis_response(response: Awaitable[T] | T) -> T:
+    """Solve the problem of ambiguous typing for redis."""
+    return await response if isinstance(response, Awaitable) else response
+
+
+def read_lua_script(script_name: str) -> str:
+    """Read a Lua script from a file."""
+    file_path = Path(__file__).parent / 'lua_scripts' / script_name
+    with file_path.open('r', encoding='utf-8') as file:
+        return file.read()

crawlee/storage_clients/_redis/lua_scripts/atomic_bloom_add_requests.lua

@@ -0,0 +1,36 @@
+local added_filter_key = KEYS[1]
+local queue_key = KEYS[2]
+local data_key = KEYS[3]
+
+local forefront = ARGV[1] == '1'
+local unique_keys = cjson.decode(ARGV[2])
+local requests_data = cjson.decode(ARGV[3])
+
+-- Add and check which unique keys are actually new using Bloom filter
+local bf_results = redis.call('bf.madd', added_filter_key, unpack(unique_keys))
+
+local actually_added = {}
+local hset_args = {}
+
+-- Process the results
+for i, unique_key in ipairs(unique_keys) do
+    if bf_results[i] == 1 then
+        -- This key was added by us (did not exist before)
+        table.insert(hset_args, unique_key)
+        table.insert(hset_args, requests_data[unique_key])
+        table.insert(actually_added, unique_key)
+    end
+end
+
+-- Add only those that are actually new
+if #actually_added > 0 then
+    redis.call('hset', data_key, unpack(hset_args))
+
+    if forefront then
+        redis.call('lpush', queue_key, unpack(actually_added))
+    else
+        redis.call('rpush', queue_key, unpack(actually_added))
+    end
+end
+
+return cjson.encode(actually_added)

crawlee/storage_clients/_redis/lua_scripts/atomic_fetch_request.lua

@@ -0,0 +1,49 @@
+local queue_key = KEYS[1]
+local in_progress_key = KEYS[2]
+local data_key = KEYS[3]
+local client_id = ARGV[1]
+local blocked_until_timestamp = ARGV[2]
+local batch_size = tonumber(ARGV[3])
+
+-- Pop batch unique_key from queue
+local batch_result = redis.call('LMPOP', 1, queue_key, 'LEFT', 'COUNT', batch_size)
+if not batch_result then
+    return nil
+end
+local unique_keys = batch_result[2]
+
+-- Get requests data
+local requests_data = redis.call('HMGET', data_key, unpack(unique_keys))
+if not requests_data then
+    -- Data missing, skip this request
+    return nil
+end
+
+-- Prepare results and update in_progress
+local final_result = {}
+local in_progress_hmset = {}
+local pending_decrement = 0
+local in_progress_data = cjson.encode({
+    client_id = client_id,
+    blocked_until_timestamp = tonumber(blocked_until_timestamp)
+})
+for i = 1, #unique_keys do
+    local unique_key = unique_keys[i]
+    local request_data = requests_data[i]
+
+    if request_data then
+        -- Add to in_progress hash
+        table.insert(in_progress_hmset, unique_key)
+        table.insert(in_progress_hmset, in_progress_data)
+
+        table.insert(final_result, request_data)
+    end
+end
+
+-- Update in_progress hash
+if #in_progress_hmset > 0 then
+    redis.call('HMSET', in_progress_key, unpack(in_progress_hmset))
+end
+
+-- Return result with requests data
+return final_result

crawlee/storage_clients/_redis/lua_scripts/atomic_set_add_requests.lua

@@ -0,0 +1,37 @@
+local added_filter_key = KEYS[1]
+local queue_key = KEYS[2]
+local data_key = KEYS[3]
+
+local forefront = ARGV[1] == '1'
+local unique_keys = cjson.decode(ARGV[2])
+local requests_data = cjson.decode(ARGV[3])
+
+-- Add and check which unique keys are actually new using Redis set
+local actually_added = {}
+local hset_args = {}
+
+-- Process each unique key
+for _, unique_key in ipairs(unique_keys) do
+    -- Try to add the key to the set, returns 1 if added, 0 if already existed
+    local set_result = redis.call('sadd', added_filter_key, unique_key)
+
+    if set_result == 1 then
+        -- This key was added by us (did not exist before)
+        table.insert(hset_args, unique_key)
+        table.insert(hset_args, requests_data[unique_key])
+        table.insert(actually_added, unique_key)
+    end
+end
+
+-- Add only those that are actually new
+if #actually_added > 0 then
+    redis.call('hset', data_key, unpack(hset_args))
+
+    if forefront then
+        redis.call('lpush', queue_key, unpack(actually_added))
+    else
+        redis.call('rpush', queue_key, unpack(actually_added))
+    end
+end
+
+return cjson.encode(actually_added)

crawlee/storage_clients/_redis/lua_scripts/reclaim_stale_requests.lua

@@ -0,0 +1,34 @@
+local in_progress_key = KEYS[1]
+local queue_key = KEYS[2]
+local data_key = KEYS[3]
+local current_time = tonumber(ARGV[1])
+
+local max_reclaim = 1000
+
+local cursor = "0"
+local count = 0
+
+repeat
+    local result = redis.call('hscan', in_progress_key, cursor, 'COUNT', 100)
+    cursor = result[1]
+    local entries = result[2]
+
+    for i = 1, #entries, 2 do
+        if count >= max_reclaim then
+            break
+        end
+
+        local unique_key = entries[i]
+        local data = cjson.decode(entries[i + 1])
+
+        -- Check if timed out
+        if current_time > data.blocked_until_timestamp then
+            -- Atomically remove from in_progress and add back to queue
+            redis.call('hdel', in_progress_key, unique_key)
+            redis.call('rpush', queue_key, unique_key)
+            count = count + 1
+        end
+    end
+until cursor == "0" or count >= max_reclaim
+
+return count

crawlee/storage_clients/_sql/__init__.py

@@ -0,0 +1,6 @@
+from ._dataset_client import SqlDatasetClient
+from ._key_value_store_client import SqlKeyValueStoreClient
+from ._request_queue_client import SqlRequestQueueClient
+from ._storage_client import SqlStorageClient
+
+__all__ = ['SqlDatasetClient', 'SqlKeyValueStoreClient', 'SqlRequestQueueClient', 'SqlStorageClient']