From 9ac7967c6cc3759bf32e9c11ad38b8ef7c00f20b Mon Sep 17 00:00:00 2001 From: Matt Hammerly Date: Wed, 8 Jul 2026 16:46:48 -0700 Subject: [PATCH] feat(py-client): implement "many" api for batch requests --- clients/python/README.md | 48 ++ clients/python/docs/objectstore_client.rst | 8 + .../python/src/objectstore_client/client.py | 22 +- clients/python/src/objectstore_client/many.py | 697 ++++++++++++++++++ clients/python/tests/test_e2e.py | 128 +++- clients/python/tests/test_many.py | 450 +++++++++++ 6 files changed, 1351 insertions(+), 2 deletions(-) create mode 100644 clients/python/src/objectstore_client/many.py create mode 100644 clients/python/tests/test_many.py diff --git a/clients/python/README.md b/clients/python/README.md index ba0bed88..7ed102e2 100644 --- a/clients/python/README.md +++ b/clients/python/README.md @@ -151,6 +151,54 @@ existing_parts = resumed.list_parts() key = resumed.complete(new_parts + existing_parts) ``` +### Many API (batch operations) + +The Many API lets you enqueue multiple operations that the client executes using +Objectstore's batch endpoint, minimizing network overhead. Operations too large +to batch (oversized inserts) automatically fall back to individual requests, and +batch and individual requests run concurrently. + +`send()` returns an `OperationResults` list. Results are **not** guaranteed to be +in the order they were enqueued, since operations execute concurrently. Each +result is a `GetResult`, `PutResult`, `DeleteResult`, `HeadResult`, or (for errors +that can't be attributed to a specific operation) an `ErrorResult`. + +```python +from objectstore_client import Client, Usecase +from objectstore_client.many import GetResult, PutResult + +client = Client("http://localhost:8888") +session = client.session(Usecase("attachments"), org=42, project=1337) + +results = ( + session.many() + .put(b"file1 contents", key="file1") + .put(b"file2 contents", key="file2") + .get("file3") + .delete("file4") + .send() +) + +for result in results: + if result.error is not None: + ... # handle per-operation error + elif isinstance(result, GetResult): + payload = result.response.payload if result.response else None # None => 404 + elif isinstance(result, PutResult): + stored_key = result.key +``` + +If you don't need to inspect individual results and just want to raise if any +operation failed, use `raise_for_failures()`, which raises an `ExceptionGroup` of +all per-operation errors: + +```python +session.many().put(b"file1", key="file1").put(b"file2", key="file2").send().raise_for_failures() +``` + +Concurrency is configurable via `.max_individual_concurrency(n)` (default 5) and +`.max_batch_concurrency(n)` (default 3). + ### Authentication If your Objectstore instance enforces authorization, you must configure authentication diff --git a/clients/python/docs/objectstore_client.rst b/clients/python/docs/objectstore_client.rst index 28fdf9a9..97f5592d 100644 --- a/clients/python/docs/objectstore_client.rst +++ b/clients/python/docs/objectstore_client.rst @@ -33,6 +33,14 @@ objectstore\_client.errors module :show-inheritance: :undoc-members: +objectstore\_client.many module +------------------------------- + +.. automodule:: objectstore_client.many + :members: + :show-inheritance: + :undoc-members: + objectstore\_client.metadata module ----------------------------------- diff --git a/clients/python/src/objectstore_client/client.py b/clients/python/src/objectstore_client/client.py index 0779529f..2ea79e03 100644 --- a/clients/python/src/objectstore_client/client.py +++ b/clients/python/src/objectstore_client/client.py @@ -3,7 +3,7 @@ from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from io import BytesIO -from typing import IO, Any, Literal, NamedTuple, cast +from typing import IO, TYPE_CHECKING, Any, Literal, NamedTuple, cast from urllib.parse import urlparse import sentry_sdk @@ -32,6 +32,9 @@ from objectstore_client.multipart import MultipartUpload from objectstore_client.scope import Scope +if TYPE_CHECKING: + from objectstore_client.many import ManyBuilder + class GetResponse(NamedTuple): metadata: Metadata @@ -275,6 +278,23 @@ def _make_url(self, key: str | None, full: bool = False) -> str: return f"http://{self._pool.host}:{self._pool.port}{path}" return path + def _make_batch_url(self) -> str: + relative_path = f"/v1/objects:batch/{self._usecase.name}/{self._scope}/" + return self._base_path.rstrip("/") + relative_path + + def many(self) -> ManyBuilder: + """ + Creates a [ManyBuilder] associated with this session. + + A [ManyBuilder] enqueues multiple operations, which the client sends as + batch requests via a dedicated endpoint (minimizing network overhead), + falling back to individual requests for operations too large to batch. + """ + # Imported lazily to avoid a circular import at module load time. + from objectstore_client.many import ManyBuilder + + return ManyBuilder(self) + def _make_multipart_url( self, action: str | None, diff --git a/clients/python/src/objectstore_client/many.py b/clients/python/src/objectstore_client/many.py new file mode 100644 index 00000000..8e4959cd --- /dev/null +++ b/clients/python/src/objectstore_client/many.py @@ -0,0 +1,697 @@ +""" +Batch (bulk) operations for the Objectstore client. + +A :class:`ManyBuilder` lets you enqueue multiple get/put/delete/head operations +and dispatch them efficiently. Operations that fit within the per-part size +limit are grouped into multipart batch requests sent to the dedicated +``objects:batch`` endpoint; operations too large to batch (oversized inserts) +fall back to individual requests. Both kinds run concurrently using thread +pools whose sizes are configurable. + +Each request part carries ``x-sn-batch-operation-kind`` and (except for keyless +inserts) ``x-sn-batch-operation-key`` headers, and the server echoes an +``x-sn-batch-operation-index`` plus ``x-sn-batch-operation-status`` on each +response part so results can be correlated back to their operations. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from io import BytesIO +from typing import IO, TYPE_CHECKING, Literal +from urllib.parse import quote, unquote + +import zstandard +from urllib3.fields import RequestField +from urllib3.filepost import encode_multipart_formdata + +from objectstore_client.client import GetResponse +from objectstore_client.errors import RequestError, raise_for_status +from objectstore_client.metadata import ( + HEADER_EXPIRATION, + HEADER_FILENAME, + HEADER_META_PREFIX, + HEADER_ORIGIN, + Compression, + ExpirationPolicy, + Metadata, + format_expiration, +) + +if TYPE_CHECKING: + from objectstore_client.client import Session + +# Maximum number of operations to send in a single batch request. +MAX_BATCH_OPS = 1000 + +# Maximum (post-compression, estimated) size of a single part's body in a batch +# request. Inserts larger than this are sent as individual requests instead. +MAX_BATCH_PART_SIZE = 1024 * 1024 # 1 MB + +# Maximum total (estimated) body size to include in a single batch request. +MAX_BATCH_BODY_SIZE = 100 * 1024 * 1024 # 100 MB + +# Default maximum number of concurrent individual (non-batch) requests. +DEFAULT_INDIVIDUAL_CONCURRENCY = 5 + +# Default maximum number of concurrent batch requests. +DEFAULT_BATCH_CONCURRENCY = 3 + +HEADER_BATCH_OPERATION_KIND = "x-sn-batch-operation-kind" +HEADER_BATCH_OPERATION_KEY = "x-sn-batch-operation-key" +HEADER_BATCH_OPERATION_INDEX = "x-sn-batch-operation-index" +HEADER_BATCH_OPERATION_STATUS = "x-sn-batch-operation-status" + + +@dataclass +class _GetOp: + key: str + decompress: bool + accept_encoding: Sequence[str] | None + + +@dataclass +class _InsertOp: + body: bytes | IO[bytes] + key: str | None + compression: Compression | None + content_type: str | None + metadata: dict[str, str] | None + expiration_policy: ExpirationPolicy | None + origin: str | None + filename: str | None + + +@dataclass +class _DeleteOp: + key: str + + +@dataclass +class _HeadOp: + key: str + + +_Operation = _GetOp | _InsertOp | _DeleteOp | _HeadOp + + +@dataclass +class GetResult: + """Result of a ``get`` operation in a batch.""" + + key: str + response: GetResponse | None + """The fetched object, or ``None`` if it was not found (404).""" + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class PutResult: + """Result of a ``put`` operation in a batch.""" + + key: str + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class DeleteResult: + """Result of a ``delete`` operation in a batch.""" + + key: str + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class HeadResult: + """Result of a ``head`` operation in a batch.""" + + key: str + metadata: Metadata | None + """The object metadata, or ``None`` if it was not found (404).""" + error: Exception | None + """The error that occurred, or ``None`` if the operation succeeded.""" + + +@dataclass +class ErrorResult: + """ + An error that could not be attributed to a specific operation. + + This can happen when a response part has a missing or malformed index or + status header, references an unknown operation index, or when the batch + request itself fails before the server can respond per-operation. + """ + + error: Exception + + +OperationResult = GetResult | PutResult | DeleteResult | HeadResult | ErrorResult + + +class OperationResults(list["OperationResult"]): + """ + The results of a :meth:`ManyBuilder.send` call. + + This is a ``list`` of :data:`OperationResult`\\ s. The order is **not** + guaranteed to match the order in which operations were enqueued, since + operations execute concurrently. + """ + + def failures(self) -> list[OperationResult]: + """Returns the subset of results that carry an error.""" + return [r for r in self if r.error is not None] + + def raise_for_failures(self) -> None: + """ + Raises an :class:`ExceptionGroup` of all per-operation errors, if any. + + Does nothing when every operation succeeded. + """ + errors = [r.error for r in self if r.error is not None] + if errors: + raise ExceptionGroup("one or more batch operations failed", errors) + + +class ManyBuilder: + """ + Enqueues multiple operations to be executed together. + + Construct one via :meth:`~objectstore_client.client.Session.many`, enqueue + operations with :meth:`get`/:meth:`put`/:meth:`delete`/:meth:`head`, then + call :meth:`send` to execute them. + """ + + def __init__(self, session: Session): + self._session = session + self._ops: list[_Operation] = [] + self._max_individual_concurrency: int | None = None + self._max_batch_concurrency: int | None = None + + def get( + self, + key: str, + *, + decompress: bool = True, + accept_encoding: Sequence[str] | None = None, + ) -> ManyBuilder: + """Enqueues a ``get`` for ``key``. See :meth:`Session.get`.""" + self._ops.append(_GetOp(key, decompress, accept_encoding)) + return self + + def put( + self, + contents: bytes | IO[bytes], + *, + key: str | None = None, + compression: Compression | Literal["none"] | None = None, + content_type: str | None = None, + metadata: dict[str, str] | None = None, + expiration_policy: ExpirationPolicy | None = None, + origin: str | None = None, + filename: str | None = None, + ) -> ManyBuilder: + """Enqueues a ``put`` of ``contents``. See :meth:`Session.put`.""" + if compression and compression not in ("none", "zstd"): + raise ValueError(f"Invalid compression: {compression}") + self._ops.append( + _InsertOp( + body=contents, + key=key or None, + compression=compression, + content_type=content_type, + metadata=metadata, + expiration_policy=expiration_policy, + origin=origin, + filename=filename, + ) + ) + return self + + def delete(self, key: str) -> ManyBuilder: + """Enqueues a ``delete`` for ``key``. See :meth:`Session.delete`.""" + self._ops.append(_DeleteOp(key)) + return self + + def head(self, key: str) -> ManyBuilder: + """Enqueues a ``head`` for ``key``. See :meth:`Session.head`.""" + self._ops.append(_HeadOp(key)) + return self + + def max_individual_concurrency(self, concurrency: int) -> ManyBuilder: + """ + Sets the maximum number of concurrent individual (non-batch) requests. + + Operations that exceed the per-part size limit are sent individually. + Defaults to ``5``. + """ + self._max_individual_concurrency = concurrency + return self + + def max_batch_concurrency(self, concurrency: int) -> ManyBuilder: + """ + Sets the maximum number of concurrent batch requests. + + Batchable operations are grouped into chunks sent as multipart batch + requests. Defaults to ``3``. + """ + self._max_batch_concurrency = concurrency + return self + + def send(self) -> OperationResults: + """ + Executes all enqueued operations and returns their results. + + Batchable operations are chunked and sent as batch requests; oversized + inserts are sent individually. Both run concurrently. The results are + **not** guaranteed to be in the order the operations were enqueued. + """ + individual_concurrency = self._max_individual_concurrency or DEFAULT_INDIVIDUAL_CONCURRENCY + batch_concurrency = self._max_batch_concurrency or DEFAULT_BATCH_CONCURRENCY + + batchable, individual = self._partition() + chunks = list(_iter_batches(batchable)) + + results = OperationResults() + with ( + ThreadPoolExecutor(max_workers=individual_concurrency) as individual_pool, + ThreadPoolExecutor(max_workers=batch_concurrency) as batch_pool, + ): + individual_futures = [ + individual_pool.submit(self._execute_individual, op) + for op in individual + ] + batch_futures = [ + batch_pool.submit(self._run_batch, chunk) for chunk in chunks + ] + for individual_future in individual_futures: + results.append(individual_future.result()) + for batch_future in batch_futures: + results.extend(batch_future.result()) + + return results + + def _partition(self) -> tuple[list[tuple[_Operation, int]], list[_Operation]]: + """Splits operations into ``(batchable, individual)``. + + Each batchable entry is paired with its estimated post-compression size, + used for chunking. Only oversized (or unsized) inserts go individual. + """ + batchable: list[tuple[_Operation, int]] = [] + individual: list[_Operation] = [] + for op in self._ops: + if isinstance(op, _InsertOp): + size = self._insert_size(op) + if size is not None and size <= MAX_BATCH_PART_SIZE: + batchable.append((op, size)) + else: + individual.append(op) + else: + batchable.append((op, 0)) + return batchable, individual + + def _insert_size(self, op: _InsertOp) -> int | None: + """Estimates the post-compression body size of an insert. + + Returns ``None`` when the size cannot be determined (e.g. a + non-seekable stream), signalling that it must be sent individually. + """ + if isinstance(op.body, bytes): + raw = len(op.body) + else: + size = _stream_size(op.body) + if size is None: + return None + raw = size + + if self._effective_compression(op) == "zstd": + return _zstd_compress_bound(raw) + return raw + + def _effective_compression(self, op: _InsertOp) -> Compression: + compression = op.compression or self._session._usecase._compression + return compression + + # -- execution -------------------------------------------------------------- + + def _execute_individual(self, op: _Operation) -> OperationResult: + """Executes a single operation via the ordinary (non-batch) endpoints.""" + if isinstance(op, _InsertOp): + try: + key = self._session.put( + op.body, + key=op.key, + compression=op.compression, + content_type=op.content_type, + metadata=op.metadata, + expiration_policy=op.expiration_policy, + origin=op.origin, + filename=op.filename, + ) + return PutResult(key, None) + except Exception as exc: + return PutResult(op.key or "", exc) + elif isinstance(op, _GetOp): + try: + response = self._session.get( + op.key, + decompress=op.decompress, + accept_encoding=op.accept_encoding, + ) + return GetResult(op.key, response, None) + except RequestError as exc: + if exc.status == 404: + return GetResult(op.key, None, None) + return GetResult(op.key, None, exc) + except Exception as exc: + return GetResult(op.key, None, exc) + elif isinstance(op, _DeleteOp): + try: + self._session.delete(op.key) + return DeleteResult(op.key, None) + except Exception as exc: + return DeleteResult(op.key, exc) + else: + try: + metadata = self._session.head(op.key) + return HeadResult(op.key, metadata, None) + except Exception as exc: + return HeadResult(op.key, None, exc) + + def _run_batch(self, ops: list[_Operation]) -> list[OperationResult]: + """Runs one chunk as a batch request, converting failures to results.""" + try: + return self._send_batch(ops) + except Exception as exc: + return [_error_result_for(op, exc) for op in ops] + + def _send_batch(self, ops: list[_Operation]) -> list[OperationResult]: + fields = [self._build_part(op) for op in ops] + body, content_type = encode_multipart_formdata(fields) + + headers = self._session._make_headers() + headers["Content-Type"] = content_type + + response = self._session._pool.request( + "POST", + self._session._make_batch_url(), + body=body, + headers=headers, + preload_content=True, + decode_content=True, + ) + raise_for_status(response) + + boundary = _extract_boundary(response.headers.get("content-type", "")) + parts = _parse_multipart_response(response.data or b"", boundary) + return self._results_from_parts(ops, parts) + + def _build_part(self, op: _Operation) -> RequestField: + if isinstance(op, _GetOp): + return _simple_part("get", op.key) + if isinstance(op, _DeleteOp): + return _simple_part("delete", op.key) + if isinstance(op, _HeadOp): + return _simple_part("head", op.key) + return self._insert_part(op) + + def _insert_part(self, op: _InsertOp) -> RequestField: + headers: dict[str, str] = {HEADER_BATCH_OPERATION_KIND: "insert"} + if op.key is not None: + headers[HEADER_BATCH_OPERATION_KEY] = quote(op.key, safe="") + + payload = op.body if isinstance(op.body, bytes) else op.body.read() + + if self._effective_compression(op) == "zstd": + payload = zstandard.ZstdCompressor().compress(payload) + headers["Content-Encoding"] = "zstd" + + if op.content_type: + headers["Content-Type"] = op.content_type + expiration = op.expiration_policy or self._session._usecase._expiration_policy + if expiration: + headers[HEADER_EXPIRATION] = format_expiration(expiration) + if op.origin: + headers[HEADER_ORIGIN] = op.origin + if op.filename is not None: + headers[HEADER_FILENAME] = op.filename + if op.metadata: + for meta_key, meta_value in op.metadata.items(): + headers[f"{HEADER_META_PREFIX}{meta_key}"] = meta_value + + return RequestField("part", data=payload, headers=headers) + + # -- response correlation --------------------------------------------------- + + def _results_from_parts( + self, + ops: list[_Operation], + parts: list[tuple[dict[str, str], bytes]], + ) -> list[OperationResult]: + results: list[OperationResult] = [] + seen: set[int] = set() + for headers, part_body in parts: + index, result = self._result_from_part(ops, headers, part_body) + if index is not None: + seen.add(index) + results.append(result) + + for index in range(len(ops)): + if index not in seen: + results.append( + _error_result_for( + ops[index], + RequestError( + f"server did not return a response for operation at " + f"index {index}", + status=0, + response="", + ), + ) + ) + return results + + def _result_from_part( + self, + ops: list[_Operation], + headers: dict[str, str], + part_body: bytes, + ) -> tuple[int | None, OperationResult]: + index_raw = headers.get(HEADER_BATCH_OPERATION_INDEX) + if index_raw is None or not index_raw.isdigit(): + return None, ErrorResult( + RequestError( + f"missing or invalid {HEADER_BATCH_OPERATION_INDEX} header", + status=0, + response="", + ) + ) + index = int(index_raw) + + status = _parse_status(headers.get(HEADER_BATCH_OPERATION_STATUS)) + if status is None: + return None, ErrorResult( + RequestError( + f"missing or invalid {HEADER_BATCH_OPERATION_STATUS} header", + status=0, + response="", + ) + ) + + if index >= len(ops): + return None, ErrorResult( + RequestError( + f"response references unknown operation index {index}", + status=0, + response="", + ) + ) + op = ops[index] + + key_header = headers.get(HEADER_BATCH_OPERATION_KEY) + key = unquote(key_header) if key_header else op.key + + is_error = status >= 400 and not ( + isinstance(op, (_GetOp, _HeadOp)) and status == 404 + ) + + # For successful responses the key is always present; for errors it may + # be absent (e.g. a server-generated-key insert that failed before a key + # was assigned), in which case we fall back to a sentinel. + if key is None: + if is_error: + key = "" + else: + return index, ErrorResult( + RequestError( + f"missing {HEADER_BATCH_OPERATION_KEY} header", + status=0, + response="", + ) + ) + + if is_error: + message = part_body.decode("utf-8", "replace") + error = RequestError( + f"operation failed with HTTP status code {status}", + status=status, + response=message, + ) + return index, _typed_error(op, key, error) + + if isinstance(op, _GetOp): + if status == 404: + return index, GetResult(key, None, None) + metadata = Metadata.from_headers(headers) + payload = _maybe_decompress(part_body, metadata, op) + return index, GetResult(key, GetResponse(metadata, payload), None) + if isinstance(op, _InsertOp): + return index, PutResult(key, None) + if isinstance(op, _DeleteOp): + return index, DeleteResult(key, None) + # head + if status == 404: + return index, HeadResult(key, None, None) + return index, HeadResult(key, Metadata.from_headers(headers), None) + + +def _typed_error(op: _Operation, key: str, error: Exception) -> OperationResult: + if isinstance(op, _GetOp): + return GetResult(key, None, error) + if isinstance(op, _InsertOp): + return PutResult(key, error) + if isinstance(op, _DeleteOp): + return DeleteResult(key, error) + return HeadResult(key, None, error) + + +def _error_result_for(op: _Operation, error: Exception) -> OperationResult: + return _typed_error(op, op.key or "", error) + + +def _simple_part(kind: str, key: str) -> RequestField: + return RequestField( + "part", + data=b"", + headers={ + HEADER_BATCH_OPERATION_KIND: kind, + HEADER_BATCH_OPERATION_KEY: quote(key, safe=""), + }, + ) + + +def _maybe_decompress(payload: bytes, metadata: Metadata, op: _GetOp) -> IO[bytes]: + """Applies transparent zstd decompression, mirroring :meth:`Session.get`.""" + accept = op.accept_encoding + encoding_accepted = accept is not None and ( + "*" in accept or metadata.compression in accept + ) + if metadata.compression and op.decompress and not encoding_accepted: + if metadata.compression != "zstd": + raise NotImplementedError( + "Transparent decoding of anything but `zstd` is not implemented yet" + ) + metadata.compression = None + dctx = zstandard.ZstdDecompressor() + return dctx.stream_reader(BytesIO(payload), read_across_frames=True) + return BytesIO(payload) + + +def _stream_size(stream: IO[bytes]) -> int | None: + """Returns the number of readable bytes remaining, or ``None`` if unknown.""" + try: + if not stream.seekable(): + return None + position = stream.tell() + end = stream.seek(0, 2) # SEEK_END + stream.seek(position) + return end - position + except (OSError, ValueError): + return None + + +def _zstd_compress_bound(src_size: int) -> int: + """Python port of ``ZSTD_compressBound`` (worst-case compressed size).""" + margin = ((128 << 10) - src_size) >> 11 if src_size < (128 << 10) else 0 + return src_size + (src_size >> 8) + margin + + +def _iter_batches( + ops: list[tuple[_Operation, int]], +) -> Iterator[list[_Operation]]: + """Groups operations into batches respecting the count and size limits.""" + batch: list[_Operation] = [] + batch_size = 0 + for op, size in ops: + if batch and ( + len(batch) >= MAX_BATCH_OPS or batch_size + size > MAX_BATCH_BODY_SIZE + ): + yield batch + batch, batch_size = [], 0 + batch.append(op) + batch_size += size + if batch: + yield batch + + +def _parse_status(value: str | None) -> int | None: + """Parses a status header like ``"200 OK"`` into its numeric code.""" + if not value: + return None + token = value.split(" ", 1)[0] + try: + return int(token) + except ValueError: + return None + + +def _extract_boundary(content_type: str) -> str: + for segment in content_type.split(";"): + segment = segment.strip() + if segment.startswith("boundary="): + boundary = segment[len("boundary=") :].strip() + if len(boundary) >= 2 and boundary[0] == '"' and boundary[-1] == '"': + boundary = boundary[1:-1] + return boundary + raise RequestError( + "missing multipart boundary in response Content-Type", + status=0, + response=content_type, + ) + + +def _parse_multipart_response( + body: bytes, boundary: str +) -> list[tuple[dict[str, str], bytes]]: + """Parses a ``multipart/form-data`` body into ``(headers, body)`` parts.""" + delimiter = b"--" + boundary.encode("latin-1") + parts: list[tuple[dict[str, str], bytes]] = [] + + for segment in body.split(delimiter): + # Skip the preamble (empty) and the closing boundary (starts with "--"). + if not segment or segment.startswith(b"--"): + continue + # Each part segment is `\r\n\r\n\r\n\r\n`. + if segment.startswith(b"\r\n"): + segment = segment[2:] + header_blob, separator, part_body = segment.partition(b"\r\n\r\n") + if not separator: + continue + if part_body.endswith(b"\r\n"): + part_body = part_body[:-2] + + headers: dict[str, str] = {} + for line in header_blob.split(b"\r\n"): + if not line: + continue + name, _, value = line.partition(b":") + headers[name.decode("utf-8", "replace").strip().lower()] = value.decode( + "utf-8", "replace" + ).strip() + parts.append((headers, part_body)) + + return parts diff --git a/clients/python/tests/test_e2e.py b/clients/python/tests/test_e2e.py index 359be23d..624546a8 100644 --- a/clients/python/tests/test_e2e.py +++ b/clients/python/tests/test_e2e.py @@ -13,9 +13,18 @@ import pytest import urllib3 import zstandard -from objectstore_client import Client, Usecase +from objectstore_client import ( + Client, + Usecase, +) from objectstore_client.auth import Permission, TokenGenerator from objectstore_client.errors import RequestError +from objectstore_client.many import ( + DeleteResult, + GetResult, + HeadResult, + PutResult, +) from objectstore_client.metadata import TimeToLive from objectstore_client.multipart import CompletePart, MultipartCompleteError from objectstore_client.scope import Scope @@ -650,6 +659,123 @@ def test_multipart_resume(server_url: str) -> None: assert retrieved.payload.read() == b"firstsecond" +def test_batch_operations(server_url: str) -> None: + client = Client(server_url, token=TestTokenGenerator.get()) + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + session = client.session(usecase, org=12345, project=1337) + + # First batch: 4 PUTs. key-2 uses zstd, the rest are uncompressed. + results = ( + session.many() + .put(b"first object", key="key-1", compression="none", filename="report.pdf") + .put(b"second object", key="key-2", compression="zstd") + .put(b"third object", key="key-3", compression="none") + .put(b"fourth object", key="key-4", compression="none") + .send() + ) + assert len(results) == 4 + assert all(isinstance(r, PutResult) and r.error is None for r in results) + assert sorted(r.key for r in results if isinstance(r, PutResult)) == [ + "key-1", + "key-2", + "key-3", + "key-4", + ] + + # Second batch: GET key-1, GET key-2 (auto-decompress), HEAD key-3, DELETE key-4. + results = ( + session.many().get("key-1").get("key-2").head("key-3").delete("key-4").send() + ) + assert len(results) == 4 + + gets = {r.key: r for r in results if isinstance(r, GetResult)} + heads = {r.key: r for r in results if isinstance(r, HeadResult)} + deletes = {r.key: r for r in results if isinstance(r, DeleteResult)} + + get1 = gets["key-1"] + assert get1.response is not None + assert get1.response.metadata.compression is None + assert get1.response.metadata.filename == "report.pdf" + assert get1.response.payload.read() == b"first object" + + get2 = gets["key-2"] + assert get2.response is not None + assert get2.response.metadata.compression is None # transparently decompressed + assert get2.response.payload.read() == b"second object" + + assert heads["key-3"].metadata is not None + assert deletes["key-4"].error is None + + # key-4 is gone. + assert session.head("key-4") is None + + +def test_batch_insert_without_key(server_url: str) -> None: + client = Client(server_url, token=TestTokenGenerator.get()) + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + session = client.session(usecase, org=12345, project=1337) + + results = session.many().put(b"keyless object", compression="none").send() + assert len(results) == 1 + put = results[0] + assert isinstance(put, PutResult) + assert put.error is None + assert put.key + + retrieved = session.get(put.key) + assert retrieved.payload.read() == b"keyless object" + + +def test_batch_partial_failures(server_url: str) -> None: + # A read-only token: writes/deletes fail with 403, reads 404 for missing keys. + client = Client( + server_url, + token=TestTokenGenerator.create(permissions=[Permission.OBJECT_READ]), + ) + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + session = client.session(usecase, org=12345, project=1337) + + results = ( + session.many() + .get("nonexistent-1") + .put(b"should fail", key="write-key", compression="none") + .delete("delete-key") + .get("nonexistent-2") + .send() + ) + assert len(results) == 4 + + puts = [r for r in results if isinstance(r, PutResult)] + deletes = [r for r in results if isinstance(r, DeleteResult)] + gets = [r for r in results if isinstance(r, GetResult)] + + assert isinstance(puts[0].error, RequestError) + assert puts[0].error.status == 403 + assert isinstance(deletes[0].error, RequestError) + assert deletes[0].error.status == 403 + # Missing objects come back as successful "not found" results, not errors. + assert all(g.error is None and g.response is None for g in gets) + + +def test_batch_oversized_insert_falls_back_to_individual(server_url: str) -> None: + client = Client(server_url, token=TestTokenGenerator.get()) + usecase = Usecase("test-usecase", expiration_policy=TimeToLive(timedelta(days=1))) + session = client.session(usecase, org=12345, project=1337) + + big = b"x" * (2 * 1024 * 1024) # 2 MB > per-part batch limit + results = ( + session.many() + .put(big, key="big", compression="none") + .put(b"small", key="small", compression="none") + .send() + ) + assert len(results) == 2 + assert all(isinstance(r, PutResult) and r.error is None for r in results) + + assert session.get("big").payload.read() == big + assert session.get("small").payload.read() == b"small" + + def test_multipart_concurrent_part_uploads(server_url: str) -> None: from concurrent.futures import ThreadPoolExecutor diff --git a/clients/python/tests/test_many.py b/clients/python/tests/test_many.py new file mode 100644 index 00000000..89faf66b --- /dev/null +++ b/clients/python/tests/test_many.py @@ -0,0 +1,450 @@ +"""Unit tests for batch (``many``) operations. + +Since these tests cannot spin up the real Objectstore server, they exercise the +full encode → parse → correlate path against an in-memory fake connection pool +that speaks the same ``objects:batch`` multipart protocol as the server. This +validates request serialization, response parsing, and result correlation +without a live backend (end-to-end behavior is covered by the Rust e2e suite). +""" + +from __future__ import annotations + +import json +from typing import Any +from urllib.parse import unquote + +import pytest +import zstandard +from objectstore_client import Client, Usecase +from objectstore_client.errors import RequestError +from objectstore_client.many import ( + DeleteResult, + ErrorResult, + GetResult, + HeadResult, + PutResult, + _extract_boundary, + _InsertOp, + _iter_batches, + _parse_multipart_response, + _parse_status, + _zstd_compress_bound, +) + +BATCH_OP_HEADERS = { + "x-sn-batch-operation-kind", + "x-sn-batch-operation-key", + "x-sn-batch-operation-index", + "x-sn-batch-operation-status", + "content-disposition", +} + +_STATUS_REASONS = { + 200: "OK", + 201: "Created", + 204: "No Content", + 403: "Forbidden", + 404: "Not Found", +} + + +def _status_line(code: int) -> str: + return f"{code} {_STATUS_REASONS.get(code, '')}".strip() + + +class FakeResponse: + def __init__(self, status: int, headers: dict[str, str], data: bytes) -> None: + self.status = status + self.headers = headers + self.data = data + + def read(self) -> bytes: + return self.data + + +def _serialize_response( + parts: list[tuple[dict[str, str], bytes]], boundary: str +) -> bytes: + """Serializes response parts exactly as the server's multipart writer does.""" + delimiter = b"--" + boundary.encode() + out = bytearray() + for headers, body in parts: + out += delimiter + b"\r\n" + for name, value in headers.items(): + out += f"{name}: {value}".encode() + b"\r\n" + out += b"\r\n" + out += body + b"\r\n" + out += delimiter + b"--" + return bytes(out) + + +class FakeBatchPool: + """A minimal in-memory stand-in for the server's ``objects:batch`` endpoint.""" + + def __init__(self, *, deny_writes: bool = False) -> None: + self.headers: dict[str, str] = {} + self.store: dict[str, tuple[dict[str, str], bytes]] = {} + self.deny_writes = deny_writes + self._generated = 0 + + def request( + self, + method: str, + url: str, + *, + body: bytes, + headers: dict[str, str], + **_kwargs: Any, + ) -> FakeResponse: + assert method == "POST" + assert "objects:batch" in url + boundary = _extract_boundary(headers["Content-Type"]) + request_parts = _parse_multipart_response(body, boundary) + + response_parts: list[tuple[dict[str, str], bytes]] = [] + for index, (part_headers, part_body) in enumerate(request_parts): + response_parts.append(self._handle(index, part_headers, part_body)) + + response_boundary = "os-boundary-" + "0" * 32 + data = _serialize_response(response_parts, response_boundary) + return FakeResponse( + 200, + {"content-type": f'multipart/form-data; boundary="{response_boundary}"'}, + data, + ) + + def _handle( + self, index: int, headers: dict[str, str], body: bytes + ) -> tuple[dict[str, str], bytes]: + kind = headers["x-sn-batch-operation-kind"] + key_enc = headers.get("x-sn-batch-operation-key") + key = unquote(key_enc) if key_enc else None + + base = { + "x-sn-batch-operation-index": str(index), + "x-sn-batch-operation-kind": kind, + } + + def with_key(resp: dict[str, str], k: str) -> dict[str, str]: + from urllib.parse import quote + + resp["x-sn-batch-operation-key"] = quote(k, safe="") + return resp + + if kind == "insert": + if self.deny_writes: + resp = dict(base) + resp["x-sn-batch-operation-status"] = _status_line(403) + if key is not None: + with_key(resp, key) + return resp, json.dumps({"detail": "forbidden"}).encode() + if key is None: + self._generated += 1 + key = f"generated-{self._generated}" + meta = {k: v for k, v in headers.items() if k not in BATCH_OP_HEADERS} + self.store[key] = (meta, body) + resp = with_key(dict(base), key) + resp["x-sn-batch-operation-status"] = _status_line(201) + return resp, b"" + + assert key is not None + if kind == "delete": + if self.deny_writes: + resp = with_key(dict(base), key) + resp["x-sn-batch-operation-status"] = _status_line(403) + return resp, json.dumps({"detail": "forbidden"}).encode() + self.store.pop(key, None) + resp = with_key(dict(base), key) + resp["x-sn-batch-operation-status"] = _status_line(204) + return resp, b"" + + if kind == "get": + resp = with_key(dict(base), key) + if key not in self.store: + resp["x-sn-batch-operation-status"] = _status_line(404) + return resp, b"" + meta, stored = self.store[key] + resp["x-sn-batch-operation-status"] = _status_line(200) + resp["x-sn-time-created"] = "2024-01-01T00:00:00+00:00" + resp.update(meta) + resp.setdefault("content-type", "application/octet-stream") + return resp, stored + + # head + resp = with_key(dict(base), key) + if key not in self.store: + resp["x-sn-batch-operation-status"] = _status_line(404) + return resp, b"" + meta, _stored = self.store[key] + resp["x-sn-batch-operation-status"] = _status_line(204) + resp["x-sn-time-created"] = "2024-01-01T00:00:00+00:00" + resp.update(meta) + return resp, b"" + + +def _session(pool: FakeBatchPool) -> Any: + client = Client("http://localhost:8888") + client._pool = pool # type: ignore[assignment] + return client.session(Usecase("testing", compression="none"), org=42, project=1337) + + +# --- Full round-trip tests ----------------------------------------------------- + + +def test_batch_put_get_delete_head_roundtrip() -> None: + pool = FakeBatchPool() + session = _session(pool) + + results = ( + session.many() + .put(b"first", key="key-1") + .put(b"second", key="key-2", compression="zstd", content_type="text/plain") + .put(b"third", key="key-3", filename="report.pdf", metadata={"foo": "bar"}) + .send() + ) + assert len(results) == 3 + assert all(isinstance(r, PutResult) and r.error is None for r in results) + assert {r.key for r in results} == {"key-1", "key-2", "key-3"} + + results = ( + session.many() + .get("key-1") + .get("key-2") # stored zstd-compressed, transparently decompressed + .head("key-3") + .delete("key-1") + .get("missing") + .send() + ) + get1 = next(r for r in results if isinstance(r, GetResult) and r.key == "key-1") + assert get1.error is None and get1.response is not None + assert get1.response.payload.read() == b"first" + + get2 = next(r for r in results if isinstance(r, GetResult) and r.key == "key-2") + assert get2.response is not None + assert get2.response.metadata.compression is None # decompressed + assert get2.response.payload.read() == b"second" + + head3 = next(r for r in results if isinstance(r, HeadResult)) + assert head3.metadata is not None + assert head3.metadata.filename == "report.pdf" + assert head3.metadata.custom.get("foo") == "bar" + + delete1 = next(r for r in results if isinstance(r, DeleteResult)) + assert delete1.error is None + + missing = next( + r for r in results if isinstance(r, GetResult) and r.key == "missing" + ) + assert missing.response is None and missing.error is None + + +def test_batch_get_accept_encoding_passthrough() -> None: + pool = FakeBatchPool() + session = _session(pool) + + session.many().put(b"payload", key="c", compression="zstd").send() + + results = session.many().get("c", accept_encoding=["zstd"]).send() + get = results[0] + assert isinstance(get, GetResult) and get.response is not None + # Passthrough: still compressed, metadata preserved. + assert get.response.metadata.compression == "zstd" + raw = get.response.payload.read() + assert zstandard.ZstdDecompressor().decompress(raw) == b"payload" + + +def test_batch_keyless_insert_gets_server_key() -> None: + pool = FakeBatchPool() + session = _session(pool) + + results = session.many().put(b"no key here").send() + assert len(results) == 1 + put = results[0] + assert isinstance(put, PutResult) + assert put.error is None + assert put.key.startswith("generated-") + assert put.key in pool.store + + +def test_batch_partial_failures() -> None: + pool = FakeBatchPool(deny_writes=True) + session = _session(pool) + # Seed an object directly so the read succeeds. + pool.store["exists"] = ({"content-type": "application/octet-stream"}, b"hi") + + results = ( + session.many() + .put(b"denied", key="write-key") + .delete("delete-key") + .get("exists") + .send() + ) + + put = next(r for r in results if isinstance(r, PutResult)) + assert isinstance(put.error, RequestError) + assert put.error.status == 403 + + delete = next(r for r in results if isinstance(r, DeleteResult)) + assert isinstance(delete.error, RequestError) + assert delete.error.status == 403 + + get = next(r for r in results if isinstance(r, GetResult)) + assert get.error is None and get.response is not None + assert get.response.payload.read() == b"hi" + + +def test_raise_for_failures() -> None: + pool = FakeBatchPool(deny_writes=True) + session = _session(pool) + + results = session.many().put(b"x", key="k1").put(b"y", key="k2").send() + assert len(results.failures()) == 2 + + with pytest.raises(ExceptionGroup) as exc_info: + results.raise_for_failures() + assert len(exc_info.value.exceptions) == 2 + + +def test_batch_request_failure_produces_per_op_errors() -> None: + """A whole-batch HTTP failure maps to one error per enqueued operation.""" + + class FailingPool: + headers: dict[str, str] = {} + + def request(self, *_args: Any, **_kwargs: Any) -> FakeResponse: + return FakeResponse(500, {}, b"boom") + + session = _session(FakeBatchPool()) + session._pool = FailingPool() # type: ignore[assignment] + + results = session.many().get("a").delete("b").send() + assert len(results) == 2 + assert all(isinstance(r.error, RequestError) for r in results) # type: ignore[union-attr] + assert {r.error.status for r in results} == {500} # type: ignore[union-attr] + + +# --- Classification / chunking ------------------------------------------------- + + +def test_partition_oversized_insert_is_individual() -> None: + pool = FakeBatchPool() + session = _session(pool) + + builder = ( + session.many() + .put(b"small", key="small") + .put(b"x" * (2 * 1024 * 1024), key="big", compression="none") + .get("g") + ) + batchable, individual = builder._partition() + + assert len(individual) == 1 + assert isinstance(individual[0], _InsertOp) + assert individual[0].key == "big" + # small insert + get are batchable + assert {op.key for op, _ in batchable} == {"small", "g"} # type: ignore[union-attr] + + +def test_partition_unseekable_stream_is_individual() -> None: + import io + + class Unseekable(io.BytesIO): + def seekable(self) -> bool: + return False + + pool = FakeBatchPool() + session = _session(pool) + builder = session.many().put(Unseekable(b"data"), key="stream") + batchable, individual = builder._partition() + assert not batchable + assert len(individual) == 1 + + +def test_iter_batches_count_limit() -> None: + from objectstore_client.many import _DeleteOp + + ops: list[tuple[Any, int]] = [(_DeleteOp(f"k{i}"), 0) for i in range(1001)] + batches = list(_iter_batches(ops)) + assert [len(b) for b in batches] == [1000, 1] + + +def test_iter_batches_size_limit() -> None: + from objectstore_client.many import MAX_BATCH_BODY_SIZE, _DeleteOp + + one_mb = 1024 * 1024 + ops: list[tuple[Any, int]] = [(_DeleteOp(f"k{i}"), one_mb) for i in range(150)] + batches = list(_iter_batches(ops)) + per_batch = MAX_BATCH_BODY_SIZE // one_mb + assert len(batches) > 1 + for batch in batches[:-1]: + assert len(batch) == per_batch + + +# --- Helper functions ---------------------------------------------------------- + + +def test_zstd_compress_bound_is_upper_bound() -> None: + for size in (0, 1, 1000, 100_000, 5_000_000): + payload = b"a" * size + actual = len(zstandard.ZstdCompressor().compress(payload)) + assert actual <= _zstd_compress_bound(size) + + +def test_parse_status() -> None: + assert _parse_status("200 OK") == 200 + assert _parse_status("404 Not Found") == 404 + assert _parse_status("204") == 204 + assert _parse_status(None) is None + assert _parse_status("garbage") is None + + +def test_extract_boundary() -> None: + assert ( + _extract_boundary('multipart/form-data; boundary="os-boundary-abc"') + == "os-boundary-abc" + ) + assert _extract_boundary("multipart/form-data; boundary=plain") == "plain" + with pytest.raises(RequestError): + _extract_boundary("application/json") + + +def test_parse_multipart_response_roundtrip() -> None: + boundary = "os-boundary-xyz" + parts = [ + ({"x-sn-batch-operation-index": "0", "content-type": "text/plain"}, b"hello"), + ({"x-sn-batch-operation-index": "1"}, b""), + ({"x-sn-batch-operation-index": "2"}, b"\x00\x01\xff binary \r\n data"), + ] + body = _serialize_response(parts, boundary) + parsed = _parse_multipart_response(body, boundary) + + assert len(parsed) == 3 + assert parsed[0][0]["x-sn-batch-operation-index"] == "0" + assert parsed[0][0]["content-type"] == "text/plain" + assert parsed[0][1] == b"hello" + assert parsed[1][1] == b"" + assert parsed[2][1] == b"\x00\x01\xff binary \r\n data" + + +def test_unattributed_error_result_for_bad_index() -> None: + boundary = "os-boundary-xyz" + parts = [({"x-sn-batch-operation-status": "200 OK"}, b"")] # no index header + body = _serialize_response(parts, boundary) + + pool = FakeBatchPool() + + class BadIndexPool(FakeBatchPool): + def request(self, *_args: Any, **_kwargs: Any) -> FakeResponse: + return FakeResponse( + 200, + {"content-type": f'multipart/form-data; boundary="{boundary}"'}, + body, + ) + + session = _session(pool) + session._pool = BadIndexPool() # type: ignore[assignment] + + results = session.many().get("k").send() + # one unattributed error + one synthesized missing-response error + assert any(isinstance(r, ErrorResult) for r in results) + assert any(isinstance(r, GetResult) and r.error is not None for r in results)