diff --git a/src/apify/storage_clients/_apify/_dataset_client.py b/src/apify/storage_clients/_apify/_dataset_client.py index de9634fd..db6bec4d 100644 --- a/src/apify/storage_clients/_apify/_dataset_client.py +++ b/src/apify/storage_clients/_apify/_dataset_client.py @@ -1,13 +1,13 @@ from __future__ import annotations import asyncio +import json from logging import getLogger from typing import TYPE_CHECKING from typing_extensions import override from crawlee._utils.byte_size import ByteSize -from crawlee._utils.file import json_dumps from crawlee.storage_clients._base import DatasetClient from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata @@ -138,10 +138,6 @@ async def drop(self) -> None: @override async def push_data(self, data: Sequence[Mapping[str, JsonSerializable]] | Mapping[str, JsonSerializable]) -> None: - async def payloads_generator(items: Sequence[Mapping[str, JsonSerializable]]) -> AsyncIterator[str]: - for index, item in enumerate(items): - yield await self._check_and_serialize(item, index) - async with self._charge_lock(), self._lock: items = data if self._is_sequence_of_items(data) else [data] if not items: @@ -149,7 +145,9 @@ async def payloads_generator(items: Sequence[Mapping[str, JsonSerializable]]) -> limit = self._compute_limit_for_push(len(items)) items = items[:limit] - async for chunk in self._chunk_by_size(payloads_generator(items)): + offset = 0 + while offset < len(items): + chunk, offset = await asyncio.to_thread(self._serialize_chunk, items, offset) await self._api_client.push_items(items=chunk) await self._charge_for_items(count_items=limit) @@ -213,58 +211,43 @@ async def iterate_items( yield item @classmethod - async def _check_and_serialize(cls, item: Mapping[str, JsonSerializable], index: int | None = None) -> str: - """Serialize a given item to JSON, checks its serializability and size against a limit. + def _serialize_chunk(cls, items: Sequence[Mapping[str, JsonSerializable]], offset: int) -> tuple[str, int]: + """Serialize items starting at `offset` into one JSON array staying within the payload size limit. + + The array holds as many consecutive items as fit within `_EFFECTIVE_LIMIT_SIZE`, always at least one. Output + is compact JSON - it goes straight on the wire. This is CPU-bound and blocking; call it via `asyncio.to_thread`. Args: - item: The item to serialize. - index: Index of the item, used for error context. + items: The items to serialize. + offset: Index of the first item to serialize. Returns: - Serialized JSON string. + The JSON array string and the index of the first item that did not fit into it. Raises: - ValueError: If item is not JSON serializable or exceeds size limit. - """ - s = ' ' if index is None else f' at index {index} ' - - try: - payload = await json_dumps(item) - except Exception as exc: - raise ValueError(f'Data item{s}is not serializable to JSON.') from exc - - payload_size = ByteSize(len(payload.encode('utf-8'))) - if payload_size > cls._EFFECTIVE_LIMIT_SIZE: - raise ValueError(f'Data item{s}is too large (size: {payload_size}, limit: {cls._EFFECTIVE_LIMIT_SIZE})') - - return payload - - async def _chunk_by_size(self, items: AsyncIterator[str]) -> AsyncIterator[str]: - """Yield chunks of JSON arrays composed of input strings, respecting a size limit. - - Groups an iterable of JSON string payloads into larger JSON arrays, ensuring the total size - of each array does not exceed `EFFECTIVE_LIMIT_SIZE`. Each output is a JSON array string that - contains as many payloads as possible without breaching the size threshold, maintaining the - order of the original payloads. Assumes individual items are below the size limit. - - Args: - items: Iterable of JSON string payloads. - - Yields: - Strings representing JSON arrays of payloads, each staying within the size limit. + ValueError: If an item is not JSON serializable or on its own exceeds the size limit. """ - last_chunk_size = ByteSize(2) # Add 2 bytes for [] wrapper. - current_chunk = [] - - async for payload in items: - payload_size = ByteSize(len(payload.encode('utf-8'))) - - if last_chunk_size + payload_size <= self._EFFECTIVE_LIMIT_SIZE: - current_chunk.append(payload) - last_chunk_size += payload_size + ByteSize(1) # Add 1 byte for ',' separator. - else: - yield f'[{",".join(current_chunk)}]' - current_chunk = [payload] - last_chunk_size = payload_size + ByteSize(2) # Add 2 bytes for [] wrapper. - - yield f'[{",".join(current_chunk)}]' + limit = cls._EFFECTIVE_LIMIT_SIZE.bytes + payloads: list[str] = [] + chunk_size = 2 # Add 2 bytes for [] wrapper. + + for index in range(offset, len(items)): + try: + payload = json.dumps(items[index], ensure_ascii=False, separators=(',', ':'), default=str) + except Exception as exc: + raise ValueError(f'Data item at index {index} is not serializable to JSON.') from exc + + payload_size = len(payload.encode('utf-8')) + if payload_size > limit: + raise ValueError( + f'Data item at index {index} is too large ' + f'(size: {ByteSize(payload_size)}, limit: {cls._EFFECTIVE_LIMIT_SIZE})' + ) + + if payloads and chunk_size + payload_size > limit: + return f'[{",".join(payloads)}]', index + + payloads.append(payload) + chunk_size += payload_size + 1 # Add 1 byte for ',' separator. + + return f'[{",".join(payloads)}]', len(items) diff --git a/tests/unit/storage_clients/test_apify_dataset_client.py b/tests/unit/storage_clients/test_apify_dataset_client.py index 263d30e5..acdb529a 100644 --- a/tests/unit/storage_clients/test_apify_dataset_client.py +++ b/tests/unit/storage_clients/test_apify_dataset_client.py @@ -1,10 +1,13 @@ from __future__ import annotations import asyncio -from unittest.mock import AsyncMock +import json +from unittest.mock import AsyncMock, Mock import pytest +from crawlee._utils.byte_size import ByteSize + from apify.storage_clients._apify._dataset_client import ApifyDatasetClient @@ -31,3 +34,73 @@ async def test_drop_calls_api_delete() -> None: client, api_client = _make_dataset_client() await client.drop() api_client.delete.assert_awaited_once() + + +async def test_push_data_sends_compact_json() -> None: + """Pushed payloads carry no indentation or separator padding.""" + client, api_client = _make_dataset_client() + + await client.push_data([{'id': 1, 'name': 'first'}, {'id': 2, 'name': 'second'}]) + + chunk = api_client.push_items.await_args.kwargs['items'] + assert chunk == '[{"id":1,"name":"first"},{"id":2,"name":"second"}]' + + +async def test_push_data_serializes_in_a_single_thread_hop_per_chunk(monkeypatch: pytest.MonkeyPatch) -> None: + """Serialization is offloaded once per pushed chunk rather than once per item.""" + monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(200)) + to_thread = Mock(wraps=asyncio.to_thread) + monkeypatch.setattr(asyncio, 'to_thread', to_thread) + client, api_client = _make_dataset_client() + + await client.push_data([{'id': i} for i in range(500)]) + + assert api_client.push_items.await_count > 1 + assert to_thread.call_count == api_client.push_items.await_count + + +async def test_push_data_makes_progress_when_an_item_fills_a_whole_chunk(monkeypatch: pytest.MonkeyPatch) -> None: + """An item that fits the limit only without the array wrapper still yields one chunk per item.""" + items = [{'value': 'x' * 30} for _ in range(3)] + payloads = [json.dumps(item, ensure_ascii=False, separators=(',', ':')) for item in items] + monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(len(payloads[0].encode('utf-8')))) + client, api_client = _make_dataset_client() + + async with asyncio.timeout(5): + await client.push_data(items) + + chunks = [call.kwargs['items'] for call in api_client.push_items.await_args_list] + assert chunks == [f'[{payload}]' for payload in payloads] + + +async def test_push_data_splits_items_into_chunks_within_the_size_limit(monkeypatch: pytest.MonkeyPatch) -> None: + """Items are pushed in several chunks, each staying within the payload size limit.""" + monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(100)) + client, api_client = _make_dataset_client() + items = [{'value': 'x' * 30} for _ in range(5)] + + await client.push_data(items) + + chunks = [call.kwargs['items'] for call in api_client.push_items.await_args_list] + assert len(chunks) > 1 + assert all(len(chunk.encode('utf-8')) <= 100 for chunk in chunks) + assert [item for chunk in chunks for item in json.loads(chunk)] == items + + +async def test_push_data_rejects_an_oversized_item(monkeypatch: pytest.MonkeyPatch) -> None: + """An item exceeding the payload size limit raises with its index.""" + monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(100)) + client, _ = _make_dataset_client() + + with pytest.raises(ValueError, match='at index 1 is too large'): + await client.push_data([{'id': 1}, {'value': 'x' * 200}]) + + +async def test_push_data_rejects_a_non_serializable_item() -> None: + """An item that cannot be serialized to JSON raises with its index.""" + client, _ = _make_dataset_client() + circular: dict = {} + circular['self'] = circular + + with pytest.raises(ValueError, match='at index 0 is not serializable'): + await client.push_data(circular)