crawlee 1.0.3b6__py3-none-any.whl → 1.2.2b24__py3-none-any.whl

crawlee/storage_clients/_redis/_key_value_store_client.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import json
+from logging import getLogger
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import override
+
+from crawlee._utils.file import infer_mime_type
+from crawlee.storage_clients._base import KeyValueStoreClient
+from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
+
+from ._client_mixin import MetadataUpdateParams, RedisClientMixin
+from ._utils import await_redis_response
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from redis.asyncio import Redis
+
+logger = getLogger(__name__)
+
+
+class RedisKeyValueStoreClient(KeyValueStoreClient, RedisClientMixin):
+    """Redis implementation of the key-value store client.
+
+    This client persists key-value data to Redis using hash data structures for efficient storage and retrieval.
+    Keys are mapped to values with automatic content type detection and size tracking for metadata management.
+
+    The key-value store data is stored in Redis using the following key pattern:
+    - `key_value_stores:{name}:items` - Redis hash containing key-value pairs (values stored as binary data).
+    - `key_value_stores:{name}:metadata_items` - Redis hash containing metadata for each key.
+    - `key_value_stores:{name}:metadata` - Redis JSON object containing store metadata.
+
+    Values are serialized based on their type: JSON objects are stored as UTF-8 encoded JSON strings,
+    text values as UTF-8 encoded strings, and binary data as-is. The implementation automatically handles
+    content type detection and maintains metadata about each record including size and MIME type information.
+
+    All operations are atomic through Redis hash operations and pipeline transactions. The client supports
+    concurrent access through Redis's built-in atomic operations for hash fields.
+    """
+
+    _DEFAULT_NAME = 'default'
+    """Default Key-Value Store name key prefix when none provided."""
+
+    _MAIN_KEY = 'key_value_stores'
+    """Main Redis key prefix for Key-Value Store."""
+
+    _CLIENT_TYPE = 'Key-value store'
+    """Human-readable client type for error messages."""
+
+    def __init__(self, storage_name: str, storage_id: str, redis: Redis) -> None:
+        """Initialize a new instance.
+
+        Preferably use the `RedisKeyValueStoreClient.open` class method to create a new instance.
+        """
+        super().__init__(storage_name=storage_name, storage_id=storage_id, redis=redis)
+
+    @property
+    def _items_key(self) -> str:
+        """Return the Redis key for the items of KVS."""
+        return f'{self._MAIN_KEY}:{self._storage_name}:items'
+
+    @property
+    def _metadata_items_key(self) -> str:
+        """Return the Redis key for the items metadata of KVS."""
+        return f'{self._MAIN_KEY}:{self._storage_name}:metadata_items'
+
+    @classmethod
+    async def open(
+        cls,
+        *,
+        id: str | None,
+        name: str | None,
+        alias: str | None,
+        redis: Redis,
+    ) -> RedisKeyValueStoreClient:
+        """Open or create a new Redis key-value store client.
+
+        This method attempts to open an existing key-value store from the Redis database. If a store with the specified
+        ID or name exists, it loads the metadata from the database. If no existing store is found, a new one
+        is created.
+
+        Args:
+            id: The ID of the key-value store. If not provided, a random ID will be generated.
+            name: The name of the key-value store for named (global scope) storages.
+            alias: The alias of the key-value store for unnamed (run scope) storages.
+            redis: Redis client instance.
+
+        Returns:
+            An instance for the opened or created storage client.
+        """
+        return await cls._open(
+            id=id,
+            name=name,
+            alias=alias,
+            redis=redis,
+            metadata_model=KeyValueStoreMetadata,
+            extra_metadata_fields={},
+            instance_kwargs={},
+        )
+
+    @override
+    async def get_metadata(self) -> KeyValueStoreMetadata:
+        return await self._get_metadata(KeyValueStoreMetadata)
+
+    @override
+    async def drop(self) -> None:
+        await self._drop(extra_keys=[self._items_key, self._metadata_items_key])
+
+    @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),
+        )
+
+    @override
+    async def set_value(self, *, key: str, value: Any, content_type: str | None = None) -> None:
+        # Special handling for None values
+        if value is None:
+            content_type = 'application/x-none'  # Special content type to identify None values
+            value_bytes = b''
+        else:
+            content_type = content_type or infer_mime_type(value)
+
+            # Serialize the value to bytes.
+            if 'application/json' in content_type:
+                value_bytes = json.dumps(value, default=str, ensure_ascii=False).encode('utf-8')
+            elif isinstance(value, str):
+                value_bytes = value.encode('utf-8')
+            elif isinstance(value, (bytes, bytearray)):
+                value_bytes = value
+            else:
+                # Fallback: attempt to convert to string and encode.
+                value_bytes = str(value).encode('utf-8')
+
+        size = len(value_bytes)
+        item_metadata = KeyValueStoreRecordMetadata(
+            key=key,
+            content_type=content_type,
+            size=size,
+        )
+
+        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))
+
+    @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))
+
+        async with self._get_pipeline() as pipe:
+            await self._update_metadata(pipe, **MetadataUpdateParams(update_accessed_at=True))
+
+        if not isinstance(serialized_metadata_item, (str, bytes, bytearray)):
+            logger.warning(f'Metadata for key "{key}" is missing or invalid.')
+            return None
+
+        metadata_item = KeyValueStoreRecordMetadata.model_validate_json(serialized_metadata_item)
+
+        # Handle None values
+        if metadata_item.content_type == 'application/x-none':
+            return KeyValueStoreRecord(value=None, **metadata_item.model_dump())
+
+        # Query the record by key
+        # redis-py typing issue
+        value_bytes: bytes | None = await await_redis_response(self._redis.hget(self._items_key, key))  # ty: ignore[invalid-assignment]
+
+        if value_bytes is None:
+            logger.warning(f'Value for key "{key}" is missing.')
+            return None
+
+        # Handle JSON values
+        if 'application/json' in metadata_item.content_type:
+            try:
+                value = json.loads(value_bytes.decode('utf-8'))
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                logger.warning(f'Failed to decode JSON value for key "{key}"')
+                return None
+        # Handle text values
+        elif metadata_item.content_type.startswith('text/'):
+            try:
+                value = value_bytes.decode('utf-8')
+            except UnicodeDecodeError:
+                logger.warning(f'Failed to decode text value for key "{key}"')
+                return None
+        # Handle binary values
+        else:
+            value = value_bytes
+
+        return KeyValueStoreRecord(value=value, **metadata_item.model_dump())
+
+    @override
+    async def delete_value(self, *, key: str) -> None:
+        async with self._get_pipeline() as pipe:
+            await await_redis_response(pipe.hdel(self._items_key, key))
+            await await_redis_response(pipe.hdel(self._metadata_items_key, key))
+            await self._update_metadata(pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True))
+
+    @override
+    async def iterate_keys(
+        self,
+        *,
+        exclusive_start_key: str | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[KeyValueStoreRecordMetadata]:
+        items_data = await await_redis_response(self._redis.hgetall(self._metadata_items_key))
+
+        if not items_data:
+            return  # No items to iterate over
+
+        if not isinstance(items_data, dict):
+            raise TypeError('The items data was received in an incorrect format.')
+
+        # Get all keys, sorted alphabetically
+        keys = sorted(items_data.keys())  # ty: ignore[invalid-argument-type]
+
+        # Apply exclusive_start_key filter if provided
+        if exclusive_start_key is not None:
+            bytes_exclusive_start_key = exclusive_start_key.encode()
+            keys = [k for k in keys if k > bytes_exclusive_start_key]
+
+        # Apply limit if provided
+        if limit is not None:
+            keys = keys[:limit]
+
+        # Yield metadata for each key
+        for key in keys:
+            record = items_data[key]
+            yield KeyValueStoreRecordMetadata.model_validate_json(record)
+
+        async with self._get_pipeline() as pipe:
+            await self._update_metadata(
+                pipe,
+                **MetadataUpdateParams(update_accessed_at=True),
+            )
+
+    @override
+    async def get_public_url(self, *, key: str) -> str:
+        raise NotImplementedError('Public URLs are not supported for memory key-value stores.')
+
+    @override
+    async def record_exists(self, *, key: str) -> bool:
+        async with self._get_pipeline(with_execute=False) as pipe:
+            await await_redis_response(pipe.hexists(self._items_key, key))
+            await self._update_metadata(
+                pipe,
+                **MetadataUpdateParams(update_accessed_at=True),
+            )
+            results = await pipe.execute()
+
+        return bool(results[0])