fix: Read file-like KVS values before upload and reject unencodable ones

src/apify_client/_utils/encoding.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-import io
 import json
 from base64 import b64encode
 from functools import cache
+from inspect import isawaitable, iscoroutine
 from typing import TYPE_CHECKING, Any
 
 from apify_client._models import WebhookCreate, WebhookRepresentation
@@ -12,29 +12,51 @@
     from apify_client.types import WebhooksList
 
 
-def encode_key_value_store_record_value(value: Any, *, content_type: str | None = None) -> tuple[Any, str]:
+def encode_key_value_store_record_value(
+    value: Any, *, content_type: str | None = None
+) -> tuple[bytes | bytearray | str, str]:
     """Encode a value for storage in a key-value store record.
 
     Args:
-        value: The value to encode (can be dict, str, bytes, or file-like object).
+        value: The value to encode. Anything exposing a callable `read` is treated as a file-like object: `read`
+            is called with no arguments, so the value is consumed from its current position and buffered in
+            memory whole - the object is neither rewound nor closed, and async file-like objects are rejected.
+            Any other value is JSON-serialized unless it is already bytes or a string.
         content_type: The content type; if None, it's inferred from the value type.
 
     Returns:
         A tuple of (encoded_value, content_type).
+
+    Raises:
+        TypeError: If the value cannot be encoded into a body the transport accepts.
     """
+    # 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. Impit exposes
+    # no streaming `content=` API, so the value has to be buffered whole.
+    read = getattr(value, 'read', None)
+    if callable(read):
+        value = read()
+
+        if isawaitable(value):
+            if iscoroutine(value):
+                value.close()  # Prevent a "coroutine was never awaited" warning.
+            raise TypeError(
+                'Async file-like objects are not supported. Await the read yourself and pass the resulting '
+                'bytes or string.'
+            )
+
+        if not isinstance(value, (bytes, bytearray, str)):
+            raise TypeError(f'Reading the file-like value returned {type(value).__name__}, expected bytes or str.')
+
     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,
@@ -43,6 +65,14 @@ def encode_key_value_store_record_value(value: Any, *, content_type: str | None
             default=str,
         ).encode('utf-8')
 
+    # A non-JSON content type skips the serialization above, so anything that is not bytes-like would reach the
+    # transport unencoded and fail there with an opaque error.
+    if not isinstance(value, (bytes, bytearray, str)):
+        raise TypeError(
+            f'Cannot encode a {type(value).__name__} value as {content_type!r}. Pass bytes, a string, or a '
+            'file-like object, or use a JSON content type.'
+        )
+
     return (value, content_type)
 
 

tests/unit/test_key_value_store.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import gzip
+import io
+from typing import TYPE_CHECKING, Any
+
+import brotli
+import pytest
+from werkzeug import Request, Response
+
+from apify_client import ApifyClient, ApifyClientAsync
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    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'
+
+
+class DuckTypedReader:
+    """A file-like object that is not an `io.IOBase`, so only duck-typed detection picks it up."""
+
+    def read(self) -> bytes:
+        return b'buffer data'
+
+
+# The values are built by a factory because reading consumes them, and each case runs once per compression
+# algorithm. Each case is (value factory, expected uploaded body, expected content type).
+_FILE_LIKE_VALUE_CASES = [
+    pytest.param(lambda: io.BytesIO(b'buffer data'), b'buffer data', 'application/octet-stream', id='bytes io'),
+    pytest.param(lambda: io.StringIO('buffer data'), b'buffer data', 'text/plain; charset=utf-8', id='string io'),
+    pytest.param(DuckTypedReader, b'buffer data', 'application/octet-stream', id='duck-typed reader'),
+]
+
+
+@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
+
+
+@pytest.fixture
+def api_url(httpserver: HTTPServer) -> str:
+    """The base URL of the mock server, in the form the clients expect."""
+    return httpserver.url_for('/').removesuffix('/')
+
+
+@pytest.fixture
+def captured_records(httpserver: HTTPServer) -> list[Request]:
+    """Answer record uploads with a 201 and collect the requests the client sent."""
+    requests: list[Request] = []
+
+    def capture_request(request: Request) -> Response:
+        requests.append(request)
+        return Response(status=201)
+
+    httpserver.expect_request(_RECORD_PATH, method='PUT').respond_with_handler(capture_request)
+    return requests
+
+
+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
+
+
+@pytest.mark.parametrize(('make_value', 'expected_body', 'expected_content_type'), _FILE_LIKE_VALUE_CASES)
+def test_set_record_reads_file_like_value_sync(
+    *,
+    api_url: str,
+    captured_records: list[Request],
+    compression_case: tuple[HttpCompressionAlgorithm, str],
+    make_value: Callable[[], Any],
+    expected_body: bytes,
+    expected_content_type: str,
+) -> None:
+    """A file-like value is read and its bytes are uploaded, not passed through unread."""
+    algorithm, content_encoding = compression_case
+    client = ApifyClient(token='test_token', api_url=api_url, compression=algorithm)
+
+    client.key_value_store(_MOCKED_KVS_ID).set_record('f', make_value())
+
+    assert len(captured_records) == 1
+    assert captured_records[0].headers['content-encoding'] == content_encoding
+    assert decode_body(captured_records[0]) == expected_body
+    assert captured_records[0].headers['content-type'] == expected_content_type
+
+
+@pytest.mark.parametrize(('make_value', 'expected_body', 'expected_content_type'), _FILE_LIKE_VALUE_CASES)
+async def test_set_record_reads_file_like_value_async(
+    *,
+    api_url: str,
+    captured_records: list[Request],
+    compression_case: tuple[HttpCompressionAlgorithm, str],
+    make_value: Callable[[], Any],
+    expected_body: bytes,
+    expected_content_type: str,
+) -> None:
+    """A file-like value is read and its bytes are uploaded, not passed through unread."""
+    algorithm, content_encoding = compression_case
+    client = ApifyClientAsync(token='test_token', api_url=api_url, compression=algorithm)
+
+    await client.key_value_store(_MOCKED_KVS_ID).set_record('f', make_value())
+
+    assert len(captured_records) == 1
+    assert captured_records[0].headers['content-encoding'] == content_encoding
+    assert decode_body(captured_records[0]) == expected_body
+    assert captured_records[0].headers['content-type'] == expected_content_type

tests/unit/test_utils.py

@@ -223,8 +223,7 @@ def test__clean_json_payload(input_dict: dict, expected: dict) -> None:
 def test_encode_key_value_store_record_value_dict() -> None:
     """Test that dictionaries are encoded as JSON."""
     value, content_type = encode_key_value_store_record_value({'key': 'value'})
-    assert b'"key"' in value
-    assert b'"value"' in value
+    assert value == b'{"key": "value"}'
     assert content_type == 'application/json; charset=utf-8'
 
 
@@ -249,13 +248,61 @@ 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'
+
+
+def test_encode_key_value_store_record_value_async_file_like_raises() -> None:
+    """Test that an async file-like value is rejected instead of storing the repr of an un-awaited coroutine."""
+
+    class AsyncReader:
+        async def read(self) -> bytes:
+            return b'buffer data'
+
+    with pytest.raises(TypeError, match='Async file-like objects are not supported'):
+        encode_key_value_store_record_value(AsyncReader())
+
+
+def test_encode_key_value_store_record_value_non_bytes_read_raises() -> None:
+    """Test that a `read` returning neither bytes nor str is rejected instead of being JSON-serialized."""
+
+    class EmptyNonBlockingReader:
+        def read(self) -> None:
+            """Mimic a non-blocking raw stream with no data available."""
+
+    with pytest.raises(TypeError, match='returned NoneType, expected bytes or str'):
+        encode_key_value_store_record_value(EmptyNonBlockingReader())
+
+
+def test_encode_key_value_store_record_value_non_encodable_with_explicit_content_type_raises() -> None:
+    """Test that a non-bytes-like value with a non-JSON content type is rejected before it reaches the transport."""
+    with pytest.raises(TypeError, match="Cannot encode a dict value as 'image/png'"):
+        encode_key_value_store_record_value({'a': 1}, content_type='image/png')
+
+
 def test_response_to_dict() -> None:
     """Test parsing response as dictionary."""
     mock_response = Mock()