fix: Make batch_add_requests split batches by serialized payload size

Author: vdusek

src/apify_client/_resource_clients/request_queue.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import json
 import math
 from collections.abc import Iterable
 from queue import Queue
@@ -413,10 +414,14 @@ def batch_add_requests(
         payload_size_limit_bytes = _MAX_PAYLOAD_SIZE_BYTES - math.ceil(_MAX_PAYLOAD_SIZE_BYTES * _SAFETY_BUFFER_PERCENT)
 
         # Split the requests into batches, constrained by the max payload size and max requests per batch.
+        # Sizes are measured with the same JSON serialization the HTTP client applies to request bodies.
         batches = constrained_batches(
             requests_as_dicts,
             max_size=payload_size_limit_bytes,
             max_count=_RQ_MAX_REQUESTS_PER_BATCH,
+            get_len=lambda r: len(json.dumps(r, ensure_ascii=False, allow_nan=False, default=str).encode('utf-8')),
+            # An individually oversized request gets its own batch and is left for the API to reject.
+            strict=False,
         )
 
         # Put the batches into the queue for processing.
@@ -990,10 +995,14 @@ async def batch_add_requests(
         payload_size_limit_bytes = _MAX_PAYLOAD_SIZE_BYTES - math.ceil(_MAX_PAYLOAD_SIZE_BYTES * _SAFETY_BUFFER_PERCENT)
 
         # Split the requests into batches, constrained by the max payload size and max requests per batch.
+        # Sizes are measured with the same JSON serialization the HTTP client applies to request bodies.
         batches = constrained_batches(
             requests_as_dicts,
             max_size=payload_size_limit_bytes,
             max_count=_RQ_MAX_REQUESTS_PER_BATCH,
+            get_len=lambda r: len(json.dumps(r, ensure_ascii=False, allow_nan=False, default=str).encode('utf-8')),
+            # An individually oversized request gets its own batch and is left for the API to reject.
+            strict=False,
         )
 
         for batch in batches:

tests/unit/test_client_request_queue.py

@@ -1,18 +1,29 @@
 from __future__ import annotations
 
+import gzip
+import json
 import re
 from typing import TYPE_CHECKING
 
 import pytest
+from werkzeug.wrappers import Response
 
 from apify_client import ApifyClient, ApifyClientAsync
 from apify_client.errors import ApifyApiError
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from pytest_httpserver import HTTPServer
+    from werkzeug.wrappers import Request
 
     from apify_client._typeddicts import RequestDraftDict
 
+# The Apify API limit on the payload size of a batch-add request, which the client's batching must respect.
+_API_MAX_PAYLOAD_SIZE_BYTES = 9 * 1024 * 1024
+
+_EMPTY_BATCH_RESPONSE_CONTENT = '{"data": {"processedRequests": [], "unprocessedRequests": []}}'
+
 _PARTIALLY_ADDED_BATCH_RESPONSE_CONTENT = """{
   "data": {
     "processedRequests": [
@@ -96,6 +107,126 @@ def test_batch_not_processed_raises_exception_sync(httpserver: HTTPServer) -> No
         rq_client.batch_add_requests(requests=requests)
 
 
+def _make_large_requests() -> list[RequestDraftDict]:
+    """Return 3 requests of ~4 MB each, so that all of them together exceed the 9 MB payload limit."""
+    return [
+        {
+            'unique_key': f'http://example.com/{i}',
+            'url': f'http://example.com/{i}?filler={"x" * (4 * 1024 * 1024)}',
+            'method': 'GET',
+        }
+        for i in range(3)
+    ]
+
+
+def _payload_capturing_handler(payloads: list[bytes]) -> Callable[[Request], Response]:
+    """Return a handler that records each POST body (gzip-decompressed) and responds with an empty batch result."""
+
+    def handler(request: Request) -> Response:
+        payloads.append(gzip.decompress(request.get_data()))
+        return Response(_EMPTY_BATCH_RESPONSE_CONTENT, status=200, content_type='application/json')
+
+    return handler
+
+
+async def test_batch_add_requests_splits_batches_by_payload_size_async(httpserver: HTTPServer) -> None:
+    """Test that batches are split by serialized byte size so no POST payload exceeds the 9 MB API limit."""
+    server_url = httpserver.url_for('/').removesuffix('/')
+    client = ApifyClientAsync(
+        token='placeholder_token',
+        api_url=server_url,
+        api_public_url=server_url,
+    )
+
+    payloads = list[bytes]()
+    httpserver.expect_request(re.compile(r'.*'), method='POST').respond_with_handler(
+        _payload_capturing_handler(payloads)
+    )
+    rq_client = client.request_queue(request_queue_id='whatever')
+
+    await rq_client.batch_add_requests(requests=_make_large_requests())
+
+    assert len(payloads) > 1
+    assert all(len(payload) <= _API_MAX_PAYLOAD_SIZE_BYTES for payload in payloads)
+    assert sum(len(json.loads(payload)) for payload in payloads) == 3
+
+
+def test_batch_add_requests_splits_batches_by_payload_size_sync(httpserver: HTTPServer) -> None:
+    """Test that batches are split by serialized byte size so no POST payload exceeds the 9 MB API limit."""
+    server_url = httpserver.url_for('/').removesuffix('/')
+    client = ApifyClient(
+        token='placeholder_token',
+        api_url=server_url,
+        api_public_url=server_url,
+    )
+
+    payloads = list[bytes]()
+    httpserver.expect_request(re.compile(r'.*'), method='POST').respond_with_handler(
+        _payload_capturing_handler(payloads)
+    )
+    rq_client = client.request_queue(request_queue_id='whatever')
+
+    rq_client.batch_add_requests(requests=_make_large_requests())
+
+    assert len(payloads) > 1
+    assert all(len(payload) <= _API_MAX_PAYLOAD_SIZE_BYTES for payload in payloads)
+    assert sum(len(json.loads(payload)) for payload in payloads) == 3
+
+
+def _make_oversized_and_small_requests() -> list[RequestDraftDict]:
+    """Return a small request plus one whose serialized size alone exceeds the 9 MB payload limit."""
+    return [
+        {'unique_key': 'small', 'url': 'http://example.com/small', 'method': 'GET'},
+        {
+            'unique_key': 'oversized',
+            'url': f'http://example.com/oversized?filler={"x" * (10 * 1024 * 1024)}',
+            'method': 'GET',
+        },
+    ]
+
+
+async def test_batch_add_requests_sends_oversized_request_alone_async(httpserver: HTTPServer) -> None:
+    """Test that a request exceeding the payload limit is sent in its own batch for the API to judge, not rejected."""
+    server_url = httpserver.url_for('/').removesuffix('/')
+    client = ApifyClientAsync(
+        token='placeholder_token',
+        api_url=server_url,
+        api_public_url=server_url,
+    )
+
+    payloads = list[bytes]()
+    httpserver.expect_request(re.compile(r'.*'), method='POST').respond_with_handler(
+        _payload_capturing_handler(payloads)
+    )
+    rq_client = client.request_queue(request_queue_id='whatever')
+
+    await rq_client.batch_add_requests(requests=_make_oversized_and_small_requests())
+
+    assert sum(len(json.loads(payload)) for payload in payloads) == 2
+    assert any(len(payload) > _API_MAX_PAYLOAD_SIZE_BYTES for payload in payloads)
+
+
+def test_batch_add_requests_sends_oversized_request_alone_sync(httpserver: HTTPServer) -> None:
+    """Test that a request exceeding the payload limit is sent in its own batch for the API to judge, not rejected."""
+    server_url = httpserver.url_for('/').removesuffix('/')
+    client = ApifyClient(
+        token='placeholder_token',
+        api_url=server_url,
+        api_public_url=server_url,
+    )
+
+    payloads = list[bytes]()
+    httpserver.expect_request(re.compile(r'.*'), method='POST').respond_with_handler(
+        _payload_capturing_handler(payloads)
+    )
+    rq_client = client.request_queue(request_queue_id='whatever')
+
+    rq_client.batch_add_requests(requests=_make_oversized_and_small_requests())
+
+    assert sum(len(json.loads(payload)) for payload in payloads) == 2
+    assert any(len(payload) > _API_MAX_PAYLOAD_SIZE_BYTES for payload in payloads)
+
+
 def test_batch_processed_partially_sync(httpserver: HTTPServer) -> None:
     server_url = httpserver.url_for('/').removesuffix('/')
     client = ApifyClient(