diff --git a/src/crawlee/storage_clients/_base/_key_value_store_client.py b/src/crawlee/storage_clients/_base/_key_value_store_client.py index 2acb11c57c..53d49059eb 100644 --- a/src/crawlee/storage_clients/_base/_key_value_store_client.py +++ b/src/crawlee/storage_clients/_base/_key_value_store_client.py @@ -79,6 +79,27 @@ async def iterate_keys( if False: yield KeyValueStoreRecordMetadata() + async def iterate_entries( + self, + *, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[KeyValueStoreRecord]: + """Iterate over all the existing records in the key-value store, including their values. + + The backend method for the `KeyValueStore.iterate_entries` and `KeyValueStore.iterate_values` calls. + + The default implementation lists the keys with `iterate_keys` and reads each value separately with + `get_value` as the iteration advances, so only a single value is held in memory at a time. Whether a record + deleted while the iteration is in progress is yielded depends on the implementation. Backends that can read + the keys together with their values more efficiently should override this method; such an implementation + must still keep the memory usage bounded, e.g. by streaming or by batching on the known record sizes. + """ + async for metadata in self.iterate_keys(exclusive_start_key=exclusive_start_key, limit=limit): + record = await self.get_value(key=metadata.key) + if record is not None: + yield record + @abstractmethod async def get_public_url(self, *, key: str) -> str: """Get the public URL for the given key. diff --git a/src/crawlee/storage_clients/_memory/_key_value_store_client.py b/src/crawlee/storage_clients/_memory/_key_value_store_client.py index e984a9932a..13014099d7 100644 --- a/src/crawlee/storage_clients/_memory/_key_value_store_client.py +++ b/src/crawlee/storage_clients/_memory/_key_value_store_client.py @@ -151,9 +151,11 @@ async def iterate_keys( if limit is not None: keys = keys[:limit] - # Yield metadata for each key + # Yield metadata for each key, skipping records deleted while the iteration was suspended. for key in keys: - record = self._records[key] + record = self._records.get(key) + if record is None: + continue yield KeyValueStoreRecordMetadata( key=key, content_type=record.content_type, diff --git a/src/crawlee/storage_clients/_redis/_key_value_store_client.py b/src/crawlee/storage_clients/_redis/_key_value_store_client.py index e1190492fe..bad1b76234 100644 --- a/src/crawlee/storage_clients/_redis/_key_value_store_client.py +++ b/src/crawlee/storage_clients/_redis/_key_value_store_client.py @@ -10,10 +10,11 @@ from crawlee._utils.file import infer_mime_type from crawlee._utils.retry import retry_on_error from crawlee.storage_clients._base import KeyValueStoreClient +from crawlee.storage_clients._utils import batch_records_by_size from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata from ._client_mixin import MetadataUpdateParams, RedisClientMixin -from ._utils import await_redis_response +from ._utils import await_redis_response, expect_bytes if TYPE_CHECKING: from collections.abc import AsyncIterator @@ -40,6 +41,9 @@ class RedisKeyValueStoreClient(KeyValueStoreClient, RedisClientMixin): 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. + + Values are stored as raw bytes, so the Redis client must not decode responses. A Redis client created with + `decode_responses=True` is rejected when a value is read. """ _DEFAULT_NAME = 'default' @@ -51,6 +55,15 @@ class RedisKeyValueStoreClient(KeyValueStoreClient, RedisClientMixin): _CLIENT_TYPE = 'Key-value store' """Human-readable client type for error messages.""" + _ITERATE_ENTRIES_BATCH_MAX_KEYS = 100 + """Maximum number of records fetched with a single HMGET call in `iterate_entries`.""" + + _ITERATE_ENTRIES_BATCH_MAX_BYTES = 8 * 1024 * 1024 + """Maximum total size of the records fetched with a single HMGET call in `iterate_entries`. + + A single record larger than this is still fetched, but alone in its batch. + """ + def __init__(self, storage_name: str, storage_id: str, redis: Redis) -> None: """Initialize a new instance. @@ -178,15 +191,30 @@ async def get_value(self, *, key: str) -> KeyValueStoreRecord | 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] + value_bytes = expect_bytes(await await_redis_response(self._redis.hget(self._items_key, key))) + + return self._build_record(metadata_item, value_bytes) + + @staticmethod + def _build_record( + metadata_item: KeyValueStoreRecordMetadata, value_bytes: bytes | None + ) -> KeyValueStoreRecord | None: + """Deserialize a stored value based on its content type into a record. + + Returns None, after logging a warning, when the value is missing or cannot be decoded as the content type + claims. + """ + key = metadata_item.key if value_bytes is None: logger.warning(f'Value for key "{key}" is missing.') return None + # Handle None values + if metadata_item.content_type == 'application/x-none': + value = None # Handle JSON values - if 'application/json' in metadata_item.content_type: + elif 'application/json' in metadata_item.content_type: try: value = json.loads(value_bytes.decode('utf-8')) except (json.JSONDecodeError, UnicodeDecodeError): @@ -252,6 +280,45 @@ async def iterate_keys( **MetadataUpdateParams(update_accessed_at=True), ) + @override + async def iterate_entries( + self, + *, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[KeyValueStoreRecord]: + """Iterate over all the existing records in the key-value store, including their values. + + The values are fetched in batches with a single HMGET call per batch, instead of several round trips per record + as the default implementation does. The batches are bounded by the record sizes known from the metadata, so + a store with large values does not load too many of them at once. + """ + metadata_items = [ + item async for item in self.iterate_keys(exclusive_start_key=exclusive_start_key, limit=limit) + ] + + for batch in batch_records_by_size( + metadata_items, + max_records=self._ITERATE_ENTRIES_BATCH_MAX_KEYS, + max_bytes=self._ITERATE_ENTRIES_BATCH_MAX_BYTES, + ): + for record in await self._fetch_records(batch): + yield record + + @retry_on_error(RedisError) + async def _fetch_records(self, batch: list[KeyValueStoreRecordMetadata]) -> list[KeyValueStoreRecord]: + """Fetch the values of the given records with a single HMGET call and return the deserialized records.""" + keys = [item.key for item in batch] + values = [expect_bytes(v) for v in await await_redis_response(self._redis.hmget(self._items_key, keys))] + + # A missing value means the record was deleted after its metadata was listed, so it is skipped silently. + records = ( + self._build_record(metadata_item, value_bytes) + for metadata_item, value_bytes in zip(batch, values, strict=True) + if value_bytes is not None + ) + return [record for record in records if record is not None] + @override async def get_public_url(self, *, key: str) -> str: raise NotImplementedError('Public URLs are not supported for memory key-value stores.') diff --git a/src/crawlee/storage_clients/_redis/_storage_client.py b/src/crawlee/storage_clients/_redis/_storage_client.py index a6c39f5def..bc8c8759e0 100644 --- a/src/crawlee/storage_clients/_redis/_storage_client.py +++ b/src/crawlee/storage_clients/_redis/_storage_client.py @@ -23,8 +23,10 @@ class RedisStorageClient(StorageClient): 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. + 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. Key-value stores hold binary values, so their client needs + a Redis client that returns raw bytes, which is the default; a client created with `decode_responses=True` works + for datasets and request queues only. Storage types use the following Redis data structures: - **Datasets**: Redis JSON arrays for item storage with metadata in JSON objects diff --git a/src/crawlee/storage_clients/_redis/_utils.py b/src/crawlee/storage_clients/_redis/_utils.py index 92bd1afce1..5359191d74 100644 --- a/src/crawlee/storage_clients/_redis/_utils.py +++ b/src/crawlee/storage_clients/_redis/_utils.py @@ -18,6 +18,24 @@ async def await_redis_response(response: Awaitable[T] | T) -> T: return response +def expect_bytes(value: bytes | str | None) -> bytes | None: + """Narrow a Redis reply to raw bytes, rejecting a client that decodes responses. + + redis-py types every reply as `bytes | str | None`, because a client created with `decode_responses=True` returns + strings. The key-value store client stores binary values and needs the raw bytes back, so it rejects such + a client. + + Raises: + TypeError: If the reply is a string, i.e. the Redis client decodes responses. + """ + if isinstance(value, str): + raise TypeError( + 'The Redis client returned a decoded string instead of raw bytes. The Redis key-value store client ' + 'requires a Redis client created without `decode_responses=True`.' + ) + return value + + def read_lua_script(script_name: str) -> str: """Read a Lua script from a file.""" file_path = Path(__file__).parent / 'lua_scripts' / script_name diff --git a/src/crawlee/storage_clients/_sql/_key_value_store_client.py b/src/crawlee/storage_clients/_sql/_key_value_store_client.py index 91451fbc70..368379439a 100644 --- a/src/crawlee/storage_clients/_sql/_key_value_store_client.py +++ b/src/crawlee/storage_clients/_sql/_key_value_store_client.py @@ -13,6 +13,7 @@ from crawlee._utils.file import infer_mime_type from crawlee._utils.retry import retry_on_error from crawlee.storage_clients._base import KeyValueStoreClient +from crawlee.storage_clients._utils import batch_records_by_size from crawlee.storage_clients.models import ( KeyValueStoreMetadata, KeyValueStoreRecord, @@ -59,6 +60,15 @@ class SqlKeyValueStoreClient(KeyValueStoreClient, SqlClientMixin): _DEFAULT_NAME = 'default' """Default dataset name used when no name is provided.""" + _ITERATE_ENTRIES_BATCH_MAX_KEYS = 100 + """Maximum number of records listed or read with a single query in `iterate_entries`.""" + + _ITERATE_ENTRIES_BATCH_MAX_BYTES = 8 * 1024 * 1024 + """Maximum total size of the records read with a single query in `iterate_entries`. + + A single record larger than this is still read, but alone in its batch. + """ + _METADATA_TABLE = KeyValueStoreMetadataDb """SQLAlchemy model for key-value store metadata.""" @@ -202,21 +212,33 @@ async def get_value(self, *, key: str) -> KeyValueStoreRecord | None: if not record_db: return None - # Deserialize the value based on content type - value_bytes = record_db.value + return self._build_record( + key=record_db.key, + content_type=record_db.content_type, + size=record_db.size, + value_bytes=record_db.value, + ) + + @staticmethod + def _build_record( + *, key: str, content_type: str, size: int | None, value_bytes: bytes + ) -> KeyValueStoreRecord | None: + """Deserialize a stored value based on its content type into a record. + Returns None, after logging a warning, when the stored bytes cannot be decoded as the content type claims. + """ # Handle None values - if record_db.content_type == 'application/x-none': + if content_type == 'application/x-none': value = None # Handle JSON values - elif 'application/json' in record_db.content_type: + elif 'application/json' in 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 record_db.content_type.startswith('text/'): + elif content_type.startswith('text/'): try: value = value_bytes.decode('utf-8') except UnicodeDecodeError: @@ -226,12 +248,7 @@ async def get_value(self, *, key: str) -> KeyValueStoreRecord | None: else: value = value_bytes - return KeyValueStoreRecord( - key=record_db.key, - value=value, - content_type=record_db.content_type, - size=record_db.size, - ) + return KeyValueStoreRecord(key=key, value=value, content_type=content_type, size=size) @retry_on_error(SQLAlchemyError) @override @@ -282,6 +299,96 @@ async def iterate_keys( await self._add_buffer_record(session) + @override + async def iterate_entries( + self, + *, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[KeyValueStoreRecord]: + """Iterate over all the existing records in the key-value store, including their values. + + The records are read in keyset-paginated pages of metadata, and the values of each page are then fetched with + a single query per batch, instead of one query per record as the default implementation does. The batches are + bounded by the record sizes, so a store with large values does not load too many of them at once. Every query + runs in its own short session, so no transaction stays open while the consumer processes the records. + """ + last_key = exclusive_start_key + remaining = limit + + while remaining is None or remaining > 0: + page_size = self._ITERATE_ENTRIES_BATCH_MAX_KEYS + if remaining is not None: + page_size = min(page_size, remaining) + + page = await self._list_record_metadata(exclusive_start_key=last_key, limit=page_size) + if not page: + return + + for batch in batch_records_by_size( + page, + max_records=self._ITERATE_ENTRIES_BATCH_MAX_KEYS, + max_bytes=self._ITERATE_ENTRIES_BATCH_MAX_BYTES, + ): + for record in await self._fetch_records(batch): + yield record + + last_key = page[-1].key + if remaining is not None: + remaining -= len(page) + if len(page) < page_size: + return + + @retry_on_error(SQLAlchemyError) + async def _list_record_metadata( + self, *, exclusive_start_key: str | None, limit: int + ) -> list[KeyValueStoreRecordMetadata]: + """Read one page of record metadata, ordered by key and starting after `exclusive_start_key`.""" + stmt = ( + select(self._ITEM_TABLE.key, self._ITEM_TABLE.content_type, self._ITEM_TABLE.size) + .where(self._ITEM_TABLE.key_value_store_id == self._id) + .order_by(self._ITEM_TABLE.key) + .limit(limit) + ) + if exclusive_start_key is not None: + stmt = stmt.where(self._ITEM_TABLE.key > exclusive_start_key) + + async with self.get_session(with_simple_commit=True) as session: + result = await session.execute(stmt) + page = [ + KeyValueStoreRecordMetadata(key=row.key, content_type=row.content_type, size=row.size) for row in result + ] + await self._add_buffer_record(session) + + return page + + @retry_on_error(SQLAlchemyError) + async def _fetch_records(self, batch: list[KeyValueStoreRecordMetadata]) -> list[KeyValueStoreRecord]: + """Fetch the values of the given records with a single query and return the deserialized records.""" + stmt = ( + select( + self._ITEM_TABLE.key, + self._ITEM_TABLE.content_type, + self._ITEM_TABLE.size, + self._ITEM_TABLE.value, + ) + .where( + self._ITEM_TABLE.key_value_store_id == self._id, + self._ITEM_TABLE.key.in_([item.key for item in batch]), + ) + .order_by(self._ITEM_TABLE.key) + ) + + async with self.get_session(with_simple_commit=True) as session: + result = await session.execute(stmt) + rows = result.all() + + records = ( + self._build_record(key=row.key, content_type=row.content_type, size=row.size, value_bytes=row.value) + for row in rows + ) + return [record for record in records if record is not None] + @retry_on_error(SQLAlchemyError) @override async def record_exists(self, *, key: str) -> bool: diff --git a/src/crawlee/storage_clients/_utils.py b/src/crawlee/storage_clients/_utils.py new file mode 100644 index 0000000000..9fb65ba9ad --- /dev/null +++ b/src/crawlee/storage_clients/_utils.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from crawlee.storage_clients.models import KeyValueStoreRecordMetadata + + +def batch_records_by_size( + records: Iterable[KeyValueStoreRecordMetadata], + *, + max_records: int, + max_bytes: int, +) -> Iterator[list[KeyValueStoreRecordMetadata]]: + """Group record metadata into batches bounded by record count and by total record size. + + Each batch is meant to be read in a single call, so the size bound keeps a store with large values from loading + too many of them into memory at once. A record whose size alone exceeds `max_bytes` is still yielded, but alone in + its batch. A record with unknown size counts as empty. + """ + batch: list[KeyValueStoreRecordMetadata] = [] + batch_size = 0 + + for record in records: + record_size = record.size or 0 + if batch and (len(batch) >= max_records or batch_size + record_size > max_bytes): + yield batch + batch, batch_size = [], 0 + + batch.append(record) + batch_size += record_size + + if batch: + yield batch diff --git a/src/crawlee/storages/_dataset.py b/src/crawlee/storages/_dataset.py index 1f5a7297ba..0b4aece7ac 100644 --- a/src/crawlee/storages/_dataset.py +++ b/src/crawlee/storages/_dataset.py @@ -246,6 +246,21 @@ async def iterate_items( ): yield item + def __aiter__(self) -> AsyncIterator[Mapping[str, JsonSerializable]]: + """Iterate over all items in the dataset. + + Allows using the dataset directly in an `async for` loop. It is equivalent to calling `iterate_items` + with the default arguments. + + ### Usage + + ```python + async for item in dataset: + print(item) + ``` + """ + return self.iterate_items() + async def list_items( self, *, diff --git a/src/crawlee/storages/_key_value_store.py b/src/crawlee/storages/_key_value_store.py index 557f3432bb..1eb82087e1 100644 --- a/src/crawlee/storages/_key_value_store.py +++ b/src/crawlee/storages/_key_value_store.py @@ -213,6 +213,65 @@ async def iterate_keys( ): yield item + async def iterate_values( + self, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[Any]: + """Iterate over the values of the existing records in the KVS. + + The records are fetched lazily as the iteration advances, so only a bounded number of values is held in + memory at a time. By default this means one request per record on top of the key listing, unless the storage + client reads the values in bounded batches instead. + + Args: + exclusive_start_key: Key to start the iteration from. + limit: Maximum number of records to iterate over. None means no limit. + + Yields: + The value of each record. + """ + async for _, value in self.iterate_entries(exclusive_start_key=exclusive_start_key, limit=limit): + yield value + + async def iterate_entries( + self, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[tuple[str, Any]]: + """Iterate over the existing records in the KVS as `(key, value)` pairs. + + The records are fetched lazily as the iteration advances, so only a bounded number of values is held in + memory at a time. By default this means one request per record on top of the key listing, unless the storage + client reads the values in bounded batches instead. A record deleted while the iteration is in + progress may or may not be yielded, depending on whether its value was already read. + + Args: + exclusive_start_key: Key to start the iteration from. + limit: Maximum number of records to iterate over. None means no limit. + + Yields: + A `(key, value)` tuple for each record. + """ + async for record in self._client.iterate_entries(exclusive_start_key=exclusive_start_key, limit=limit): + yield record.key, record.value + + def __aiter__(self) -> AsyncIterator[str]: + """Iterate over all keys in the KVS. + + Allows using the key-value store directly in an `async for` loop, which yields the keys like iterating + over a `dict` does. Use `iterate_keys` for the key metadata, `iterate_values` for the values, or + `iterate_entries` for `(key, value)` pairs. + + ### Usage + + ```python + async for key in kvs: + print(key) + ``` + """ + return (metadata.key async for metadata in self.iterate_keys()) + async def list_keys( self, exclusive_start_key: str | None = None, diff --git a/tests/unit/storage_clients/_memory/test_memory_kvs_client.py b/tests/unit/storage_clients/_memory/test_memory_kvs_client.py index 4dfc44085e..4c7fce8058 100644 --- a/tests/unit/storage_clients/_memory/test_memory_kvs_client.py +++ b/tests/unit/storage_clients/_memory/test_memory_kvs_client.py @@ -78,3 +78,17 @@ async def test_memory_metadata_updates(kvs_client: MemoryKeyValueStoreClient) -> assert metadata.created_at == initial_created assert metadata.modified_at > initial_modified assert metadata.accessed_at > accessed_after_read + + +async def test_iterate_keys_skips_record_deleted_during_iteration(kvs_client: MemoryKeyValueStoreClient) -> None: + """Test that a record deleted while `iterate_keys` is suspended is skipped.""" + await kvs_client.set_value(key='key1', value='a') + await kvs_client.set_value(key='key2', value='b') + + collected_keys = [] + async for metadata in kvs_client.iterate_keys(): + if metadata.key == 'key1': + await kvs_client.delete_value(key='key2') + collected_keys.append(metadata.key) + + assert collected_keys == ['key1'] diff --git a/tests/unit/storage_clients/_redis/test_redis_kvs_client.py b/tests/unit/storage_clients/_redis/test_redis_kvs_client.py index b45dbbc973..69dbcc5d6d 100644 --- a/tests/unit/storage_clients/_redis/test_redis_kvs_client.py +++ b/tests/unit/storage_clients/_redis/test_redis_kvs_client.py @@ -6,15 +6,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeAsyncRedis from redis.exceptions import RedisError from crawlee.storage_clients import RedisStorageClient from crawlee.storage_clients._redis._utils import await_redis_response if TYPE_CHECKING: - from collections.abc import AsyncGenerator - - from fakeredis import FakeAsyncRedis + from collections.abc import AsyncGenerator, Iterator from crawlee.storage_clients._redis import RedisKeyValueStoreClient @@ -257,3 +256,117 @@ async def test_set_value_does_not_retry_on_unexpected_exception(kvs_client: Redi # Verify that retry logic was not attempted assert mock_sleep.call_count == 0 + + +@pytest.fixture +def hmget(kvs_client: RedisKeyValueStoreClient) -> Iterator[MagicMock]: + """Wrap the Redis `hmget` of the client in a mock that records the calls while still performing them.""" + with patch.object(kvs_client.redis, 'hmget', wraps=kvs_client.redis.hmget) as mock: + yield mock + + +def hmget_keys(hmget: MagicMock) -> list[list[str]]: + return [list(call.args[1]) for call in hmget.call_args_list] + + +async def test_iterate_entries_reads_values_in_batches(kvs_client: RedisKeyValueStoreClient, hmget: MagicMock) -> None: + """Test that `iterate_entries` fetches values with batched HMGET calls instead of `get_value` per key.""" + await kvs_client.set_value(key='a-json', value={'nested': [1, 2]}) + await kvs_client.set_value(key='b-text', value='plain text') + await kvs_client.set_value(key='c-bytes', value=b'\x00\x01binary', content_type='application/octet-stream') + await kvs_client.set_value(key='d-none', value=None) + + with patch.object(kvs_client, 'get_value', side_effect=AssertionError('get_value must not be called')): + records = [record async for record in kvs_client.iterate_entries()] + + assert [record.key for record in records] == ['a-json', 'b-text', 'c-bytes', 'd-none'] + assert [record.value for record in records] == [{'nested': [1, 2]}, 'plain text', b'\x00\x01binary', None] + assert records[0].content_type.startswith('application/json') + assert records[1].content_type.startswith('text/plain') + assert records[2].content_type == 'application/octet-stream' + + # All values fit in a single batch. + assert hmget_keys(hmget) == [['a-json', 'b-text', 'c-bytes', 'd-none']] + + +async def test_iterate_entries_batches_are_bounded_by_key_count( + kvs_client: RedisKeyValueStoreClient, hmget: MagicMock +) -> None: + """Test that `iterate_entries` splits the HMGET calls when a batch reaches the maximum number of keys.""" + for i in range(5): + await kvs_client.set_value(key=f'key{i}', value=f'value{i}') + + with patch.object(type(kvs_client), '_ITERATE_ENTRIES_BATCH_MAX_KEYS', 2): + records = [record async for record in kvs_client.iterate_entries()] + + assert [(record.key, record.value) for record in records] == [(f'key{i}', f'value{i}') for i in range(5)] + assert hmget_keys(hmget) == [['key0', 'key1'], ['key2', 'key3'], ['key4']] + + +async def test_iterate_entries_batches_are_bounded_by_size( + kvs_client: RedisKeyValueStoreClient, hmget: MagicMock +) -> None: + """Test that `iterate_entries` splits the HMGET calls by the record sizes known from the metadata. + + A record larger than the limit is still fetched, but alone in its batch. + """ + await kvs_client.set_value(key='small1', value='ab') + await kvs_client.set_value(key='small2', value='cd') + await kvs_client.set_value(key='large', value='x' * 100) + await kvs_client.set_value(key='small3', value='ef') + + with patch.object(type(kvs_client), '_ITERATE_ENTRIES_BATCH_MAX_BYTES', 10): + records = [record async for record in kvs_client.iterate_entries()] + + assert [record.key for record in records] == ['large', 'small1', 'small2', 'small3'] + assert hmget_keys(hmget) == [['large'], ['small1', 'small2', 'small3']] + + +async def test_iterate_entries_with_exclusive_start_key_and_limit(kvs_client: RedisKeyValueStoreClient) -> None: + """Test that `iterate_entries` applies `exclusive_start_key` and `limit`.""" + for i in range(6): + await kvs_client.set_value(key=f'key{i}', value=f'value{i}') + + with patch.object(kvs_client, 'get_value', side_effect=AssertionError('get_value must not be called')): + records = [record async for record in kvs_client.iterate_entries(exclusive_start_key='key1', limit=3)] + + assert [(record.key, record.value) for record in records] == [ + ('key2', 'value2'), + ('key3', 'value3'), + ('key4', 'value4'), + ] + + +async def test_iterate_entries_skips_value_deleted_after_listing( + kvs_client: RedisKeyValueStoreClient, caplog: pytest.LogCaptureFixture +) -> None: + """Test that `iterate_entries` silently skips a record whose value is gone by the time its batch is fetched.""" + await kvs_client.set_value(key='kept', value='a') + await kvs_client.set_value(key='removed', value='b') + await await_redis_response(kvs_client.redis.hdel(kvs_client._items_key, 'removed')) + + with caplog.at_level('WARNING'): + records = [record async for record in kvs_client.iterate_entries()] + + assert [(record.key, record.value) for record in records] == [('kept', 'a')] + assert 'missing' not in caplog.text + + +async def test_iterate_entries_empty_store(kvs_client: RedisKeyValueStoreClient) -> None: + """Test that `iterate_entries` on an empty store yields nothing.""" + records = [record async for record in kvs_client.iterate_entries()] + + assert records == [] + + +async def test_decoding_redis_client_is_rejected(suppress_user_warning: None) -> None: # noqa: ARG001 + """Test that reading values through a Redis client created with `decode_responses=True` raises a clear error.""" + storage_client = RedisStorageClient(redis=FakeAsyncRedis(decode_responses=True)) + kvs_client = await storage_client.create_kvs_client(name='decoding_kvs') + await kvs_client.set_value(key='key', value='value') + + with pytest.raises(TypeError, match='decode_responses'): + await kvs_client.get_value(key='key') + + with pytest.raises(TypeError, match='decode_responses'): + _ = [record async for record in kvs_client.iterate_entries()] diff --git a/tests/unit/storage_clients/_sql/test_sql_kvs_client.py b/tests/unit/storage_clients/_sql/test_sql_kvs_client.py index c09d8a13dc..1b31165964 100644 --- a/tests/unit/storage_clients/_sql/test_sql_kvs_client.py +++ b/tests/unit/storage_clients/_sql/test_sql_kvs_client.py @@ -16,7 +16,7 @@ from crawlee.storage_clients.models import KeyValueStoreMetadata if TYPE_CHECKING: - from collections.abc import AsyncGenerator + from collections.abc import AsyncGenerator, Iterator from pathlib import Path from sqlalchemy import Connection @@ -325,3 +325,113 @@ async def test_set_value_does_not_retry_on_unexpected_exception(kvs_client: SqlK # Verify that retry logic was not attempted assert mock_sleep.call_count == 0 + + +@pytest.fixture +def fetch_records(kvs_client: SqlKeyValueStoreClient) -> Iterator[AsyncMock]: + """Wrap `_fetch_records` of the client in a mock that records the value batches while still fetching them.""" + with patch.object(kvs_client, '_fetch_records', wraps=kvs_client._fetch_records) as mock: + yield mock + + +def fetched_keys(fetch_records: AsyncMock) -> list[list[str]]: + return [[item.key for item in call.args[0]] for call in fetch_records.await_args_list] + + +async def test_iterate_entries_reads_values_in_batches(kvs_client: SqlKeyValueStoreClient) -> None: + """Test that `iterate_entries` reads values in batches instead of calling `get_value` per key.""" + await kvs_client.set_value(key='a-json', value={'nested': [1, 2]}) + await kvs_client.set_value(key='b-text', value='plain text') + await kvs_client.set_value(key='c-bytes', value=b'\x00\x01binary', content_type='application/octet-stream') + await kvs_client.set_value(key='d-none', value=None) + + with patch.object(kvs_client, 'get_value', side_effect=AssertionError('get_value must not be called')): + records = [record async for record in kvs_client.iterate_entries()] + + assert [record.key for record in records] == ['a-json', 'b-text', 'c-bytes', 'd-none'] + assert [record.value for record in records] == [{'nested': [1, 2]}, 'plain text', b'\x00\x01binary', None] + assert records[0].content_type.startswith('application/json') + assert records[1].content_type.startswith('text/plain') + assert records[2].content_type == 'application/octet-stream' + assert all(record.size is not None for record in records) + + +async def test_iterate_entries_with_exclusive_start_key_and_limit(kvs_client: SqlKeyValueStoreClient) -> None: + """Test that `iterate_entries` applies `exclusive_start_key` and `limit` in the query.""" + for i in range(6): + await kvs_client.set_value(key=f'key{i}', value=f'value{i}') + + with patch.object(kvs_client, 'get_value', side_effect=AssertionError('get_value must not be called')): + records = [record async for record in kvs_client.iterate_entries(exclusive_start_key='key1', limit=3)] + + assert [(record.key, record.value) for record in records] == [ + ('key2', 'value2'), + ('key3', 'value3'), + ('key4', 'value4'), + ] + + +async def test_iterate_entries_limit_spans_multiple_pages( + kvs_client: SqlKeyValueStoreClient, fetch_records: AsyncMock +) -> None: + """Test that `iterate_entries` stops at `limit` when it falls in the middle of a later metadata page.""" + for i in range(6): + await kvs_client.set_value(key=f'key{i}', value=f'value{i}') + + with patch.object(type(kvs_client), '_ITERATE_ENTRIES_BATCH_MAX_KEYS', 2): + records = [record async for record in kvs_client.iterate_entries(limit=3)] + + assert [(record.key, record.value) for record in records] == [(f'key{i}', f'value{i}') for i in range(3)] + assert fetched_keys(fetch_records) == [['key0', 'key1'], ['key2']] + + +async def test_fetch_records_skips_record_deleted_after_listing(kvs_client: SqlKeyValueStoreClient) -> None: + """Test that a record deleted between the metadata listing and its value fetch is skipped.""" + await kvs_client.set_value(key='kept', value='a') + await kvs_client.set_value(key='removed', value='b') + batch = [item async for item in kvs_client.iterate_keys()] + await kvs_client.delete_value(key='removed') + + records = await kvs_client._fetch_records(batch) + + assert [(record.key, record.value) for record in records] == [('kept', 'a')] + + +async def test_iterate_entries_empty_store(kvs_client: SqlKeyValueStoreClient) -> None: + """Test that `iterate_entries` on an empty store yields nothing.""" + records = [record async for record in kvs_client.iterate_entries()] + + assert records == [] + + +async def test_iterate_entries_batches_are_bounded_by_key_count( + kvs_client: SqlKeyValueStoreClient, fetch_records: AsyncMock +) -> None: + """Test that `iterate_entries` splits the value queries when a batch reaches the maximum number of keys.""" + for i in range(5): + await kvs_client.set_value(key=f'key{i}', value=f'value{i}') + + with patch.object(type(kvs_client), '_ITERATE_ENTRIES_BATCH_MAX_KEYS', 2): + records = [record async for record in kvs_client.iterate_entries()] + + assert [(record.key, record.value) for record in records] == [(f'key{i}', f'value{i}') for i in range(5)] + assert fetched_keys(fetch_records) == [['key0', 'key1'], ['key2', 'key3'], ['key4']] + + +async def test_iterate_entries_batches_are_bounded_by_size( + kvs_client: SqlKeyValueStoreClient, fetch_records: AsyncMock +) -> None: + """Test that `iterate_entries` splits the value queries by the record sizes. + + A record larger than the limit is still read, but alone in its batch. + """ + await kvs_client.set_value(key='small1', value='ab') + await kvs_client.set_value(key='small2', value='cd') + await kvs_client.set_value(key='large', value='x' * 100) + await kvs_client.set_value(key='small3', value='ef') + + with patch.object(type(kvs_client), '_ITERATE_ENTRIES_BATCH_MAX_BYTES', 10): + records = [record async for record in kvs_client.iterate_entries()] + + assert [record.key for record in records] == ['large', 'small1', 'small2', 'small3'] + assert fetched_keys(fetch_records) == [['large'], ['small1', 'small2', 'small3']] diff --git a/tests/unit/storage_clients/test_batch_records_by_size.py b/tests/unit/storage_clients/test_batch_records_by_size.py new file mode 100644 index 0000000000..00f224ff13 --- /dev/null +++ b/tests/unit/storage_clients/test_batch_records_by_size.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from crawlee.storage_clients._utils import batch_records_by_size +from crawlee.storage_clients.models import KeyValueStoreRecordMetadata + + +def make_record(key: str, size: int | None) -> KeyValueStoreRecordMetadata: + return KeyValueStoreRecordMetadata(key=key, content_type='application/octet-stream', size=size) + + +def batch_keys(batches: list[list[KeyValueStoreRecordMetadata]]) -> list[list[str]]: + return [[record.key for record in batch] for batch in batches] + + +def test_batches_are_bounded_by_record_count() -> None: + """A new batch starts once the current one reaches `max_records`.""" + records = [make_record(f'k{i}', 1) for i in range(5)] + + batches = list(batch_records_by_size(records, max_records=2, max_bytes=1000)) + + assert batch_keys(batches) == [['k0', 'k1'], ['k2', 'k3'], ['k4']] + + +def test_batches_are_bounded_by_total_size() -> None: + """A new batch starts when the next record would push the total size over `max_bytes`.""" + records = [make_record('a', 4), make_record('b', 4), make_record('c', 4), make_record('d', 4)] + + batches = list(batch_records_by_size(records, max_records=100, max_bytes=10)) + + assert batch_keys(batches) == [['a', 'b'], ['c', 'd']] + + +def test_oversized_record_is_alone_in_its_batch() -> None: + """A record larger than `max_bytes` is yielded alone in its batch.""" + records = [make_record('small1', 2), make_record('large', 100), make_record('small2', 2), make_record('small3', 2)] + + batches = list(batch_records_by_size(records, max_records=100, max_bytes=10)) + + assert batch_keys(batches) == [['small1'], ['large'], ['small2', 'small3']] + + +def test_unknown_size_counts_as_empty() -> None: + """A record with unknown size does not count toward `max_bytes`.""" + records = [make_record('a', None), make_record('b', None), make_record('c', 10)] + + batches = list(batch_records_by_size(records, max_records=100, max_bytes=10)) + + assert batch_keys(batches) == [['a', 'b', 'c']] + + +def test_no_records_yield_no_batches() -> None: + """An empty input yields no batches.""" + assert list(batch_records_by_size([], max_records=100, max_bytes=10)) == [] diff --git a/tests/unit/storages/test_dataset.py b/tests/unit/storages/test_dataset.py index c18c3e54f3..7e0b64fbd5 100644 --- a/tests/unit/storages/test_dataset.py +++ b/tests/unit/storages/test_dataset.py @@ -280,6 +280,23 @@ async def test_iterate_items(dataset: Dataset) -> None: assert collected_items[-1]['id'] == 5 +async def test_async_iteration(dataset: Dataset) -> None: + """Test that the dataset can be used directly in an `async for` loop.""" + items = [{'id': i} for i in range(1, 6)] # 5 items + await dataset.push_data(items) + + collected_items = [item async for item in dataset] + + assert collected_items == items + + +async def test_async_iteration_empty_dataset(dataset: Dataset) -> None: + """Test that iterating over an empty dataset yields nothing.""" + collected_items = [item async for item in dataset] + + assert collected_items == [] + + async def test_iterate_items_with_options(dataset: Dataset) -> None: """Test iterating with offset, limit and desc parameters.""" # Add some items diff --git a/tests/unit/storages/test_key_value_store.py b/tests/unit/storages/test_key_value_store.py index 13b66b4f3a..1faadce748 100644 --- a/tests/unit/storages/test_key_value_store.py +++ b/tests/unit/storages/test_key_value_store.py @@ -9,11 +9,13 @@ from crawlee import service_locator from crawlee.configuration import Configuration from crawlee.storage_clients import FileSystemStorageClient, MemoryStorageClient, SqlStorageClient, StorageClient +from crawlee.storage_clients._memory import MemoryKeyValueStoreClient +from crawlee.storage_clients.models import KeyValueStoreRecord from crawlee.storages import KeyValueStore from crawlee.storages._storage_instance_manager import StorageInstanceManager if TYPE_CHECKING: - from collections.abc import AsyncGenerator + from collections.abc import AsyncGenerator, AsyncIterator from pathlib import Path @@ -267,6 +269,101 @@ async def test_iterate_keys_with_limit(kvs: KeyValueStore) -> None: assert len(collected_keys) == 5 +async def test_iterate_values(kvs: KeyValueStore) -> None: + """Test iterating over values in the key-value store.""" + await kvs.set_value('key1', 'value1') + await kvs.set_value('key2', {'nested': 2}) + await kvs.set_value('key3', [3]) + + collected_values = [value async for value in kvs.iterate_values()] + + assert len(collected_values) == 3 + assert 'value1' in collected_values + assert {'nested': 2} in collected_values + assert [3] in collected_values + + +async def test_iterate_entries(kvs: KeyValueStore) -> None: + """Test iterating over (key, value) pairs in the key-value store.""" + await kvs.set_value('key1', 'value1') + await kvs.set_value('key2', {'nested': 2}) + await kvs.set_value('key3', [3]) + + collected_entries = dict([entry async for entry in kvs.iterate_entries()]) + + assert collected_entries == {'key1': 'value1', 'key2': {'nested': 2}, 'key3': [3]} + + +async def test_iterate_entries_with_limit_and_exclusive_start_key(kvs: KeyValueStore) -> None: + """Test that `iterate_entries` passes `limit` and `exclusive_start_key` through to the key listing.""" + for i in range(10): + await kvs.set_value(f'key{i}', f'value{i}') + + all_keys = [metadata.key for metadata in await kvs.list_keys()] + start_key = all_keys[2] + + collected_entries = [entry async for entry in kvs.iterate_entries(exclusive_start_key=start_key, limit=3)] + + assert len(collected_entries) == 3 + expected_keys = all_keys[all_keys.index(start_key) + 1 :][:3] + assert [key for key, _ in collected_entries] == expected_keys + assert all(value == f'value{key.removeprefix("key")}' for key, value in collected_entries) + + +async def test_iterate_entries_empty_kvs(kvs: KeyValueStore) -> None: + """Test that iterating over an empty key-value store yields nothing.""" + collected_entries = [entry async for entry in kvs.iterate_entries()] + + assert collected_entries == [] + + +async def test_iterate_entries_uses_storage_client_implementation() -> None: + """Test that `iterate_entries` and `iterate_values` delegate to the storage client's `iterate_entries`.""" + + class OptimizedKeyValueStoreClient(MemoryKeyValueStoreClient): + async def iterate_entries( + self, + *, + exclusive_start_key: str | None = None, + limit: int | None = None, + ) -> AsyncIterator[KeyValueStoreRecord]: + # A single-pass implementation that never touches `get_value`. + keys = sorted(k for k in self._records if exclusive_start_key is None or k > exclusive_start_key) + for key in keys[:limit]: + record = self._records[key] + yield KeyValueStoreRecord( + key=key, + value=f'optimized-{record.value}', + content_type=record.content_type, + size=record.size, + ) + + async def get_value(self, *, key: str) -> KeyValueStoreRecord | None: + raise AssertionError(f'get_value must not be called for {key!r} when the client implements iterate_entries') + + client = await OptimizedKeyValueStoreClient.open(id=None, name=None, alias=None) + kvs = KeyValueStore(client, id=(await client.get_metadata()).id, name=None) + await kvs.set_value('key1', 'value1') + await kvs.set_value('key2', 'value2') + await kvs.set_value('key3', 'value3') + + entries = [entry async for entry in kvs.iterate_entries(exclusive_start_key='key1', limit=1)] + values = [value async for value in kvs.iterate_values()] + + assert entries == [('key2', 'optimized-value2')] + assert values == ['optimized-value1', 'optimized-value2', 'optimized-value3'] + + +async def test_async_iteration(kvs: KeyValueStore) -> None: + """Test that the key-value store can be used directly in an `async for` loop, yielding keys like a dict.""" + await kvs.set_value('key1', 'value1') + await kvs.set_value('key2', 'value2') + + collected_keys = [key async for key in kvs] + + assert sorted(collected_keys) == ['key1', 'key2'] + + async def test_drop( storage_client: StorageClient, ) -> None: