diff --git a/src/apify_client/_resource_clients/key_value_store.py b/src/apify_client/_resource_clients/key_value_store.py index 391c0969..13af637f 100644 --- a/src/apify_client/_resource_clients/key_value_store.py +++ b/src/apify_client/_resource_clients/key_value_store.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import re from contextlib import asynccontextmanager, contextmanager from http import HTTPStatus @@ -801,7 +802,10 @@ async def set_record( content_type: The content type of the saved value. timeout: Timeout for the API HTTP request. """ - value, content_type = encode_key_value_store_record_value(value, content_type=content_type) + # Encoding may read a file or serialize a large payload (blocking), so run it off the event loop. + value, content_type = await asyncio.to_thread( + encode_key_value_store_record_value, value, content_type=content_type + ) headers = {'content-type': content_type} diff --git a/src/apify_client/_utils/encoding.py b/src/apify_client/_utils/encoding.py index b4adf385..3f14fef8 100644 --- a/src/apify_client/_utils/encoding.py +++ b/src/apify_client/_utils/encoding.py @@ -1,6 +1,5 @@ from __future__ import annotations -import io import json from base64 import b64encode from functools import cache @@ -22,19 +21,21 @@ def encode_key_value_store_record_value(value: Any, *, content_type: str | None Returns: A tuple of (encoded_value, content_type). """ + # Read file-like values into memory; the transport only accepts bytes-like bodies. Detect them by a + # callable `read` (not `io.IOBase`) so duck-typed file-likes are read, not JSON-serialized. + read = getattr(value, 'read', None) + if callable(read): + value = read() + if not content_type: - if isinstance(value, (bytes, bytearray, io.IOBase)): + if isinstance(value, (bytes, bytearray)): content_type = 'application/octet-stream' elif isinstance(value, str): content_type = 'text/plain; charset=utf-8' else: content_type = 'application/json; charset=utf-8' - if ( - 'application/json' in content_type - and not isinstance(value, (bytes, bytearray, io.IOBase)) - and not isinstance(value, str) - ): + if 'application/json' in content_type and not isinstance(value, (bytes, bytearray, str)): # Don't use indentation to reduce size. value = json.dumps( value, diff --git a/tests/unit/test_key_value_store.py b/tests/unit/test_key_value_store.py new file mode 100644 index 00000000..cbf8d6c3 --- /dev/null +++ b/tests/unit/test_key_value_store.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import gzip +import io +from typing import TYPE_CHECKING + +import brotli +import pytest +from werkzeug import Request, Response + +from apify_client import ApifyClient, ApifyClientAsync + +if TYPE_CHECKING: + from pytest_httpserver import HTTPServer + + from apify_client.types import HttpCompressionAlgorithm + +_MOCKED_KVS_ID = 'test_kvs_id' +_RECORD_PATH = f'/v2/key-value-stores/{_MOCKED_KVS_ID}/records/f' + + +@pytest.fixture( + params=[ + pytest.param(('gzip', 'gzip'), id='gzip'), + pytest.param(('brotli', 'br'), id='brotli'), + ] +) +def compression_case(request: pytest.FixtureRequest) -> tuple[HttpCompressionAlgorithm, str]: + """Run each test over both supported request-body compression algorithms, as (algorithm, content-encoding).""" + return request.param + + +def decode_body(request: Request) -> bytes: + """Decompress a captured request body according to its `Content-Encoding`.""" + raw = request.get_data() + encoding = request.headers.get('Content-Encoding') + if encoding == 'gzip': + return gzip.decompress(raw) + if encoding == 'br': + return brotli.decompress(raw) + return raw + + +def test_set_record_reads_file_like_value_sync( + httpserver: HTTPServer, compression_case: tuple[HttpCompressionAlgorithm, str] +) -> None: + """A file-like value is read and its bytes are uploaded, not passed through unread.""" + algorithm, content_encoding = compression_case + captured_requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + captured_requests.append(request) + return Response(status=201) + + httpserver.expect_request(_RECORD_PATH, method='PUT').respond_with_handler(capture_request) + + api_url = httpserver.url_for('/').removesuffix('/') + client = ApifyClient(token='test_token', api_url=api_url, compression=algorithm) + + client.key_value_store(_MOCKED_KVS_ID).set_record('f', io.BytesIO(b'buffer data')) + + assert len(captured_requests) == 1 + assert captured_requests[0].headers['content-encoding'] == content_encoding + assert decode_body(captured_requests[0]) == b'buffer data' + assert captured_requests[0].headers['content-type'] == 'application/octet-stream' + + +async def test_set_record_reads_file_like_value_async( + httpserver: HTTPServer, compression_case: tuple[HttpCompressionAlgorithm, str] +) -> None: + """A file-like value is read and its bytes are uploaded, not passed through unread.""" + algorithm, content_encoding = compression_case + captured_requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + captured_requests.append(request) + return Response(status=201) + + httpserver.expect_request(_RECORD_PATH, method='PUT').respond_with_handler(capture_request) + + api_url = httpserver.url_for('/').removesuffix('/') + client = ApifyClientAsync(token='test_token', api_url=api_url, compression=algorithm) + + await client.key_value_store(_MOCKED_KVS_ID).set_record('f', io.BytesIO(b'buffer data')) + + assert len(captured_requests) == 1 + assert captured_requests[0].headers['content-encoding'] == content_encoding + assert decode_body(captured_requests[0]) == b'buffer data' + assert captured_requests[0].headers['content-type'] == 'application/octet-stream' + + +def test_set_record_reads_stringio_value_sync( + httpserver: HTTPServer, compression_case: tuple[HttpCompressionAlgorithm, str] +) -> None: + """A text file-like value is read and uploaded as text/plain through the HTTP stack.""" + algorithm, content_encoding = compression_case + captured_requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + captured_requests.append(request) + return Response(status=201) + + httpserver.expect_request(_RECORD_PATH, method='PUT').respond_with_handler(capture_request) + + api_url = httpserver.url_for('/').removesuffix('/') + client = ApifyClient(token='test_token', api_url=api_url, compression=algorithm) + + client.key_value_store(_MOCKED_KVS_ID).set_record('f', io.StringIO('buffer data')) + + assert len(captured_requests) == 1 + assert captured_requests[0].headers['content-encoding'] == content_encoding + assert decode_body(captured_requests[0]) == b'buffer data' + assert captured_requests[0].headers['content-type'] == 'text/plain; charset=utf-8' diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a52d3fb4..803e11e2 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -249,10 +249,30 @@ def test_encode_key_value_store_record_value( def test_encode_key_value_store_record_value_bytesio() -> None: - """Test that BytesIO is encoded as octet-stream.""" + """Test that BytesIO is read into bytes and encoded as octet-stream.""" buffer = io.BytesIO(b'buffer data') value, content_type = encode_key_value_store_record_value(buffer) - assert value == buffer + assert value == b'buffer data' + assert content_type == 'application/octet-stream' + + +def test_encode_key_value_store_record_value_stringio() -> None: + """Test that StringIO is read into text and encoded as text/plain.""" + buffer = io.StringIO('buffer data') + value, content_type = encode_key_value_store_record_value(buffer) + assert value == 'buffer data' + assert content_type == 'text/plain; charset=utf-8' + + +def test_encode_key_value_store_record_value_duck_typed_file_like() -> None: + """Test that a duck-typed file-like value (a callable `read`, not an `io.IOBase`) is read into bytes.""" + + class Reader: + def read(self) -> bytes: + return b'buffer data' + + value, content_type = encode_key_value_store_record_value(Reader()) + assert value == b'buffer data' assert content_type == 'application/octet-stream'