diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py new file mode 100644 index 000000000000..3386dccc3803 --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py @@ -0,0 +1,356 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Zero-copy deserializer for ``BidiReadObjectResponse`` messages. + +The generated gRPC stub parses every response with proto-plus over upb, which +copies each data chunk twice (into the upb arena, then into a Python ``bytes`` +for ``checksummed_data.content``) on the event-loop thread. Data responses are +structurally tiny apart from the chunk itself, so this module walks the wire +format directly and exposes ``content`` as a ``memoryview`` of the raw message +when :func:`is_supported` holds (or a single ``bytes`` slice when the installed +``google-crc32c`` only accepts ``bytes``). Responses carrying object metadata +(the first message of a stream opened without a read handle) or anything +unexpected fall back to the generated parser; the rest are returned as the +duck-typed :class:`_ZeroCopyBidiReadObjectResponse`. + +:func:`wrapped_rpc` builds the matching ``BidiReadObject`` callable for a +transport; :class:`~google.cloud.storage.asyncio.async_read_object_stream._AsyncReadObjectStream` +prefers it over the generated one whenever :func:`is_enabled` holds. +""" + +import os +import weakref + +import google_crc32c +from google.protobuf.message import DecodeError + +from google.cloud import _storage_v2 +from google.cloud._storage_v2.services.storage.transports.base import ( + DEFAULT_CLIENT_INFO, +) +from google.cloud._storage_v2.services.storage.transports.grpc_asyncio import ( + StorageGrpcAsyncIOTransport, +) + +# Off by default; set ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1`` (or ``true``, +# ``yes``, ``on``) to use the wire parser instead of the generated one. +_ENV_VAR = "GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ" +_ENABLED_VALUES = ("1", "true", "yes", "on") + +# Protobuf wire types (low 3 bits of a field tag): varint, 8-byte fixed, +# length-delimited, 4-byte fixed. Groups (3, 4) are left to the generated parser. +_WT_VARINT, _WT_I64, _WT_LEN, _WT_I32 = 0, 1, 2, 5 + +_ReadRangePb = _storage_v2.ReadRange.pb() + + +class _ChecksummedData: + __slots__ = ("content", "crc32c", "_has_crc32c") + + def __init__(self): + # An absent content field is b"" like the generated parser; the last + # message of a range that ends on a chunk boundary looks like that. + self.content = b"" + self.crc32c = 0 + self._has_crc32c = False + + @property + def _pb(self): + return self + + def HasField(self, name): + # "_crc32c" is the synthetic oneof of the proto3 ``optional`` field. + if name in ("crc32c", "_crc32c"): + return self._has_crc32c + raise ValueError(f"Unknown field {name!r}") + + +class _ObjectRangeData: + __slots__ = ( + "checksummed_data", + "read_range", + "range_end", + "_has_checksummed_data", + "_has_read_range", + ) + + def __init__(self): + self.checksummed_data = _ChecksummedData() + self.read_range = _ReadRangePb() + self.range_end = False + self._has_checksummed_data = False + self._has_read_range = False + + @property + def _pb(self): + return self + + def HasField(self, name): + if name == "checksummed_data": + return self._has_checksummed_data + if name == "read_range": + return self._has_read_range + raise ValueError(f"Unknown field {name!r}") + + +class _ZeroCopyBidiReadObjectResponse: + """Duck-typed stand-in for a ``BidiReadObjectResponse`` without metadata. + + Only the attributes the async read path uses are provided: + ``object_data_ranges``, ``read_handle``, ``metadata`` (always ``None``), + ``HasField`` and ``_pb`` (itself, so code that unwraps proto-plus messages + keeps working). + """ + + __slots__ = ("object_data_ranges", "read_handle", "metadata") + + def __init__(self, ranges, read_handle): + self.object_data_ranges = ranges + self.read_handle = read_handle + self.metadata = None + + @property + def _pb(self): + return self + + def HasField(self, name): + if name == "read_handle": + return self.read_handle is not None + if name == "metadata": + return False + raise ValueError(f"Unknown field {name!r}") + + +class _Malformed(ValueError): + """Raised when the wire bytes do not look like a well-formed message.""" + + +# A 10-byte varint carries 70 bits; like upb, keep only the low 64. +_UINT64_MASK = (1 << 64) - 1 +# A field tag (``field_number << 3 | wire_type``) must fit in 32 bits, as in upb. +_UINT32_MAX = (1 << 32) - 1 + + +def _varint(buf, i, end): + """Decode the varint at ``buf[i]``; it must end at or before ``end``. + + Like upb, accept up to ten bytes and keep only the low 64 bits. + """ + shift = 0 + result = 0 + while True: + b = buf[i] + i += 1 + result |= (b & 0x7F) << shift + if b < 0x80: + break + shift += 7 + if shift > 63: + raise _Malformed("varint longer than 10 bytes") + if i > end: + raise _Malformed("varint runs past its enclosing message") + return result & _UINT64_MASK, i + + +def _tag(buf, i, end): + """Decode a field tag at ``buf[i]`` into ``(field_number, wire_type, i)``. + + upb rejects field number 0, tags above 32 bits and tags encoded in more + than five bytes. + """ + start = i + tag, i = _varint(buf, i, end) + if tag < 8 or tag > _UINT32_MAX or i - start > 5: + raise _Malformed("invalid field tag") + return tag >> 3, tag & 7, i + + +def _length_prefixed(buf, i, end): + """Return ``(start, stop)`` of the length-delimited field body at ``i``. + + upb rejects lengths encoded in more than five bytes. + """ + start = i + n, i = _varint(buf, i, end) + if i - start > 5: + raise _Malformed("field length longer than five bytes") + stop = i + n + if stop > end: + raise _Malformed("field runs past its enclosing message") + return i, stop + + +def _skip(buf, i, wt, end): + if wt == _WT_VARINT: + return _varint(buf, i, end)[1] + if wt == _WT_LEN: + return _length_prefixed(buf, i, end)[1] + if wt == _WT_I64: + i += 8 + elif wt == _WT_I32: + i += 4 + else: + raise _Malformed(f"unsupported wire type {wt}") + if i > end: + raise _Malformed("field runs past its enclosing message") + return i + + +# The sub-message parsers fill in an existing object so that a singular +# embedded message occurring more than once is merged, as the generated +# parser does, instead of replacing what the earlier occurrence set. + + +def _parse_checksummed(cd, buf, mv, i, end): + while i < end: + f, wt, i = _tag(buf, i, end) + if f == 1 and wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + # Non-empty payloads alias the wire buffer; an empty one is b"". + cd.content = mv[start:i] if mv is not None and i > start else buf[start:i] + elif f == 2 and wt == _WT_I32: + if i + 4 > end: + raise _Malformed("crc32c runs past its enclosing message") + cd.crc32c = int.from_bytes(buf[i : i + 4], "little") + cd._has_crc32c = True + i += 4 + else: + i = _skip(buf, i, wt, end) + + +def _parse_range_data(buf, mv, i, end): + rd = _ObjectRangeData() + while i < end: + f, wt, i = _tag(buf, i, end) + if wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + if f == 1: + _parse_checksummed(rd.checksummed_data, buf, mv, start, i) + rd._has_checksummed_data = True + elif f == 2: + rd.read_range.MergeFromString(buf[start:i]) + rd._has_read_range = True + elif f == 3 and wt == _WT_VARINT: + v, i = _varint(buf, i, end) + rd.range_end = bool(v) + else: + i = _skip(buf, i, wt, end) + return rd + + +_generated_deserialize = _storage_v2.BidiReadObjectResponse.deserialize + + +def deserialize(buf): + """Parse a serialized ``BidiReadObjectResponse``; chunk payloads alias ``buf``. + + Anything that is not a plain data response (object metadata present, or + bytes that do not parse cleanly) is handed to the generated parser, which + either produces the real message or raises its usual ``DecodeError``. So + is any buffer that is not an immutable ``bytes``: the chunks alias ``buf``, + which is only safe when nothing can change or recycle it. + """ + if type(buf) is not bytes: + return _generated_deserialize(buf) + try: + mv = memoryview(buf) if is_supported() else None + n = len(buf) + i = 0 + ranges = [] + handle = None + while i < n: + f, wt, i = _tag(buf, i, n) + if wt == _WT_LEN: + start, end = _length_prefixed(buf, i, n) + if f == 6: + ranges.append(_parse_range_data(buf, mv, start, end)) + elif f == 7: + if handle is None: + handle = _storage_v2.BidiReadHandle() + handle._pb.MergeFromString(buf[start:end]) + elif f == 4: + return _generated_deserialize(buf) + i = end + else: + i = _skip(buf, i, wt, n) + return _ZeroCopyBidiReadObjectResponse(ranges, handle) + except (_Malformed, IndexError, DecodeError): + return _generated_deserialize(buf) + + +_memoryview_supported = None + + +def is_supported(): + """Whether the installed ``google_crc32c`` can hash ``memoryview`` inputs. + + When true, :func:`deserialize` exposes chunk payloads as ``memoryview`` + slices of the gRPC wire buffer; when false, it slices ``bytes`` instead so + checksum verification remains compatible with older ``google_crc32c`` + releases while still avoiding the intermediate upb arena copy. Both + ``value`` and ``Checksum.update`` (used by full-object reads) must accept + a ``memoryview``. The probe runs once per process. + """ + global _memoryview_supported + if _memoryview_supported is None: + probe = memoryview(b"\0") + try: + google_crc32c.value(probe) + google_crc32c.Checksum().update(probe) + except Exception: + _memoryview_supported = False + else: + _memoryview_supported = True + return _memoryview_supported + + +def is_enabled(): + """Whether a stream should use the zero-copy parser. + + Off by default; set ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1`` (or + ``true``, ``yes``, ``on``) to turn it on. + """ + return os.environ.get(_ENV_VAR, "").strip().lower() in _ENABLED_VALUES + + +_rpcs = weakref.WeakKeyDictionary() + + +def wrapped_rpc(transport): + """Return the zero-copy ``BidiReadObject`` callable for ``transport``. + + The callable is the generated transport's stub with :func:`deserialize` + swapped in, wrapped like the transport's own ``bidi_read_object`` entry: + error mapping, no default retry or timeout, and the default + ``x-goog-api-client`` header, which is what transports built by + ``AsyncGrpcClient`` send on every RPC. Returns ``None`` when + :func:`is_enabled` is false or ``transport`` is not the generated + grpc_asyncio transport, in which case callers should use the generated + wrapped method. + """ + if not isinstance(transport, StorageGrpcAsyncIOTransport) or not is_enabled(): + return None + rpc = _rpcs.get(transport) + if rpc is None: + stub = transport._logged_channel.stream_stream( + "/google.storage.v2.Storage/BidiReadObject", + request_serializer=_storage_v2.BidiReadObjectRequest.serialize, + response_deserializer=deserialize, + ) + rpc = transport._wrap_method( + stub, default_timeout=None, client_info=DEFAULT_CLIENT_INFO + ) + _rpcs[transport] = rpc + return rpc diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py index 32757bf754b7..d747938fd15f 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py @@ -16,8 +16,7 @@ import asyncio import logging -from io import BytesIO -from typing import Any, Dict, List, Optional, Tuple +from typing import IO, Any, Dict, List, Optional, Tuple, Union from google.api_core import exceptions from google.api_core.retry_async import AsyncRetry @@ -78,6 +77,61 @@ def _is_read_retryable(exc): ) +class _BytesWriter: + """Ensures a ``write()`` destination always receives ``bytes``.""" + + __slots__ = ("_buffer",) + + def __init__(self, buffer): + self._buffer = buffer + + def write(self, data): + return self._buffer.write(data if isinstance(data, bytes) else bytes(data)) + + +class _BufferWriter: + """Copies chunks directly into a writable buffer supporting the buffer protocol.""" + + __slots__ = ("_view", "_pos") + + def __init__(self, buffer, length: int): + if length <= 0: + raise ValueError("A writable buffer destination requires bytes_to_read > 0") + view = memoryview(buffer) + try: + if view.readonly or not view.c_contiguous: + raise ValueError( + "Buffer destination must be a writable, C-contiguous buffer" + ) + nbytes = view.nbytes + if nbytes < length: + raise ValueError( + f"Buffer destination ({nbytes} bytes) is smaller than " + f"bytes_to_read ({length})" + ) + self._view = view.cast("B")[:length] + finally: + view.release() + self._pos = 0 + + def write(self, data): + n = len(data) + end = self._pos + n + self._view[self._pos : end] = data + self._pos = end + return n + + def close(self): + # Lets the caller resize the buffer again. + self._view.release() + + +def _wrap_user_buffer(user_buffer, length: int, accepts_memoryview: bool = False): + if hasattr(user_buffer, "write"): + return user_buffer if accepts_memoryview else _BytesWriter(user_buffer) + return _BufferWriter(user_buffer, length) + + class AsyncMultiRangeDownloader: """Provides an interface for downloading multiple ranges of a GCS ``Object`` concurrently. @@ -360,30 +414,48 @@ async def factory(): async def download_ranges( self, - read_ranges: List[Tuple[int, int, BytesIO]], - lock: asyncio.Lock = None, + read_ranges: List[Tuple[int, int, Union[IO[bytes], bytearray, memoryview]]], + lock: Optional[asyncio.Lock] = None, retry_policy: Optional[AsyncRetry] = None, metadata: Optional[List[Tuple[str, str]]] = None, enable_checksum: bool = True, + accepts_memoryview: bool = False, ) -> None: """Downloads multiple byte ranges from the object into the buffers provided by user with automatic retries. - :type read_ranges: List[Tuple[int, int, "BytesIO"]] + :type read_ranges: List[Tuple[int, int, Union[IO[bytes], bytearray, memoryview]]] :param read_ranges: A list of tuples, where each tuple represents a combination of byte_range and writeable buffer in format - (`start_byte`, `bytes_to_read`, `writeable_buffer`). Buffer has to be provided by the user, and user has to make sure appropriate memory is available in the application to avoid out-of-memory crash. + If ``writeable_buffer`` has a ``write()`` method (e.g. + :class:`io.BytesIO`, an open file, or a custom stream), each + chunk is passed to ``writeable_buffer.write()`` as ``bytes`` + (or as a ``memoryview`` when ``accepts_memoryview`` is True). + Otherwise, ``writeable_buffer`` must be a writable, C-contiguous + buffer supporting the buffer protocol (e.g. :class:`bytearray` + or a writable :class:`memoryview` slice) of at least + ``bytes_to_read`` bytes (with ``bytes_to_read > 0``), and the SDK + copies each chunk directly into it in place (straight from the + gRPC response buffer, without an intermediate ``bytes``, when the + zero-copy parser is enabled with + ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1`` and the installed + ``google-crc32c`` accepts ``memoryview``). If fewer than + ``bytes_to_read`` bytes are available before end of object, + :exc:`BufferError` is raised. + Special cases: if the value of `bytes_to_read` is 0, it'll be interpreted as - download all contents until the end of the file from `start_byte`. + download all contents until the end of the file from `start_byte` + (requires a ``write()`` destination). Examples: * (0, 0, buffer) : downloads 0 to end , i.e. entire object. * (100, 0, buffer) : downloads from 100 to end. - :type lock: asyncio.Lock + :type lock: Optional[asyncio.Lock] :param lock: (Deprecated) This parameter is deprecated and has no effect. :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry` @@ -395,8 +467,29 @@ async def download_ranges( :type enable_checksum: bool :param enable_checksum: (Optional) If True, checksums are verified for downloaded data. Defaults to True. + :type accepts_memoryview: bool + :param accepts_memoryview: (Optional) If True, ``write()`` destinations + may receive each chunk as a read-only :class:`memoryview` of the + received gRPC message instead of ``bytes`` (only when the zero-copy + parser is enabled with ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1``). + The view keeps that message (the whole response, which may also + carry other ranges' chunks) alive while the destination or the + SDK's most recently processed response references it; use + ``bytes(data)`` to keep the data without the message. Destinations + must still accept ``bytes``, which the generated parser produces. + Unlike buffer destinations (which raise :exc:`BufferError` on + short reads), ``write()`` destinations observe each chunk's length + and allow short reads clamped at end of object as well as + ``bytes_to_read=0``. No effect on buffer destinations. Defaults + to False. + :raises ValueError: if the underlying bidi-GRPC stream is not open. :raises ValueError: if the length of read_ranges is more than 1000. + :raises ValueError: if a buffer destination is read-only, not + C-contiguous, smaller than ``bytes_to_read``, or used with + ``bytes_to_read <= 0``. + :raises BufferError: if fewer than ``bytes_to_read`` bytes are returned + for a buffer destination. :raises DataCorruption: if a checksum mismatch is detected while reading data. """ @@ -417,38 +510,41 @@ async def download_ranges( # Initialize Global State for Retry Strategy download_states = {} - for read_range in read_ranges: - read_id = generate_random_56_bit_integer() - # Unpack tuple into self-documenting variable names to improve readability. - offset, length, user_buffer = read_range - - # Heuristic to detect full object reads: - # - Implicit full object read: start offset is 0 and length is 0 (read all). - # - Explicit full object read: start offset is 0 and length matches the exact persisted size. - is_full_object_read = (offset == 0 and length == 0) or ( - self.persisted_size is not None - and offset == 0 - and length == self.persisted_size - ) - download_states[read_id] = _DownloadState( - initial_offset=offset, - initial_length=length, - user_buffer=user_buffer, - is_full_object_read=is_full_object_read, - ) + read_ids = set() + try: + for read_range in read_ranges: + read_id = generate_random_56_bit_integer() + # Unpack tuple into self-documenting variable names to improve readability. + offset, length, user_buffer = read_range + + # Heuristic to detect full object reads: + # - Implicit full object read: start offset is 0 and length is 0 (read all). + # - Explicit full object read: start offset is 0 and length matches the exact persisted size. + is_full_object_read = (offset == 0 and length == 0) or ( + self.persisted_size is not None + and offset == 0 + and length == self.persisted_size + ) + download_states[read_id] = _DownloadState( + initial_offset=offset, + initial_length=length, + user_buffer=_wrap_user_buffer( + user_buffer, length, accepts_memoryview + ), + is_full_object_read=is_full_object_read, + ) - initial_state = { - "download_states": download_states, - "read_handle": self.read_handle, - "routing_token": None, - "enable_checksum": enable_checksum, - "full_obj_server_crc32c": self.full_obj_server_crc32c, - } + initial_state = { + "download_states": download_states, + "read_handle": self.read_handle, + "routing_token": None, + "enable_checksum": enable_checksum, + "full_obj_server_crc32c": self.full_obj_server_crc32c, + } - read_ids = set(download_states.keys()) - queue = self._multiplexer.register(read_ids) + read_ids = set(download_states.keys()) + queue = self._multiplexer.register(read_ids) - try: attempt_count = 0 last_broken_generation = None @@ -539,9 +635,23 @@ async def generator(): if initial_state.get("read_handle"): self.read_handle = initial_state["read_handle"] + + for download_state in download_states.values(): + if ( + isinstance(download_state.user_buffer, _BufferWriter) + and download_state.bytes_written != download_state.initial_length + ): + raise BufferError( + f"Buffer destination expected {download_state.initial_length} " + f"bytes at offset {download_state.initial_offset}, got " + f"{download_state.bytes_written}" + ) finally: if self._multiplexer is not None: self._multiplexer.unregister(read_ids) + for download_state in download_states.values(): + if isinstance(download_state.user_buffer, _BufferWriter): + download_state.user_buffer.close() async def close(self): """ diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index d3c0780673c7..daea4ec6b8f0 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py @@ -19,6 +19,7 @@ from google.api_core.bidi_async import AsyncBidiRpc from google.cloud import _storage_v2 +from google.cloud.storage.asyncio import _zero_copy_bidi_read from google.cloud.storage.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) @@ -76,9 +77,10 @@ def __init__( self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" - self.rpc = self.client._client._transport._wrapped_methods[ - self.client._client._transport.bidi_read_object - ] + transport = self.client._client._transport + self.rpc = _zero_copy_bidi_read.wrapped_rpc(transport) + if self.rpc is None: + self.rpc = transport._wrapped_methods[transport.bidi_read_object] self.metadata = (("x-goog-request-params", f"bucket={self._full_bucket_name}"),) self.socket_like_rpc: Optional[AsyncBidiRpc] = None self._is_stream_open: bool = False diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py index 6cf17af19089..7c1ba457481a 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py @@ -149,6 +149,18 @@ def update_state_from_response( # Update State & Write Data chunk_size = len(data) + # Reject a response longer than the range before it reaches the + # destination, which may be the caller's preallocated buffer. + if ( + read_state.initial_length != 0 + and read_state.bytes_written + chunk_size > read_state.initial_length + ): + raise DataCorruption( + response, + f"Byte count mismatch for read_id {read_id}. " + f"Expected {read_state.initial_length}, " + f"got {read_state.bytes_written + chunk_size}", + ) read_state.user_buffer.write(data) # Commit updates only after the write succeeds @@ -157,18 +169,9 @@ def update_state_from_response( read_state.bytes_written += chunk_size read_state.next_expected_offset += chunk_size - # Final Byte Count & Full Object Checksum Verification + # Full Object Checksum Verification if object_data_range.range_end: read_state.is_complete = True - if ( - read_state.initial_length != 0 - and read_state.bytes_written > read_state.initial_length - ): - raise DataCorruption( - response, - f"Byte count mismatch for read_id {read_id}. " - f"Expected {read_state.initial_length}, got {read_state.bytes_written}", - ) # Perform full-object checksum verification once the stream finishes. if ( diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py index afe1dfd221d9..a20959aeeddb 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py @@ -23,7 +23,7 @@ from google.rpc import error_details_pb2, status_pb2 from google.cloud import _storage_v2 -from google.cloud.storage.asyncio import async_read_object_stream +from google.cloud.storage.asyncio import _zero_copy_bidi_read, async_read_object_stream from google.cloud.storage.asyncio.async_multi_range_downloader import ( AsyncMultiRangeDownloader, ) @@ -36,6 +36,48 @@ _TEST_READ_HANDLE = b"test-handle" +def _accept_memoryview_in_crc32c(monkeypatch): + # Let the installed google-crc32c hash the zero-copy parser's views even if + # it predates memoryview support, so that is_supported() holds. + value, update = google_crc32c.value, google_crc32c.Checksum.update + monkeypatch.setattr(google_crc32c, "value", lambda data: value(bytes(data))) + monkeypatch.setattr( + google_crc32c.Checksum, "update", lambda self, data: update(self, bytes(data)) + ) + + +def _zero_copy_response(content, offset, read_id, range_end, crc32c=None): + return _zero_copy_bidi_read.deserialize( + _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=content, + crc32c=google_crc32c.value(content) + if crc32c is None + else crc32c, + ), + range_end=range_end, + read_range=_storage_v2.ReadRange( + read_offset=offset, + read_length=len(content), + read_id=read_id, + ), + ) + ] + ) + ) + ) + + +@pytest.fixture(autouse=True) +def _reprobe_crc32c_memoryview_support(monkeypatch): + # is_supported() caches its google_crc32c probe; tests below patch + # google_crc32c.value, so each test must probe afresh. + monkeypatch.setattr(_zero_copy_bidi_read, "_memoryview_supported", None) + + class TestAsyncMultiRangeDownloader: def create_read_ranges(self, num_ranges): ranges = [] @@ -98,6 +140,485 @@ async def test_create_mrd(self, mock_cls_async_read_object_stream): assert mrd.object_metadata == mrd.read_obj_str.object_metadata assert mrd._open_retries == 0 + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "zero_copy, accepts_memoryview, expected_type", + [ + # A write() sink sees a memoryview only when the zero-copy parser + # is enabled (GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1) AND the + # caller opted in; the generated parser always yields bytes. + (True, True, memoryview), + (True, False, bytes), + (False, True, bytes), + (False, False, bytes), + ], + ) + async def test_download_ranges_write_destination_chunk_type( + self, + mock_cls_async_read_object_stream, + mock_random_int, + monkeypatch, + zero_copy, + accepts_memoryview, + expected_type, + ): + _accept_memoryview_in_crc32c(monkeypatch) + data = b"these_are_18_chars" + crc32c_int = google_crc32c.value(data) + + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456] + + wire = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=data, crc32c=crc32c_int + ), + range_end=True, + read_range=_storage_v2.ReadRange( + read_offset=0, read_length=18, read_id=456 + ), + ) + ] + ) + ) + if zero_copy: + resp = _zero_copy_bidi_read.deserialize(wire) + assert isinstance( + resp.object_data_ranges[0].checksummed_data.content, memoryview + ) + else: + resp = _storage_v2.BidiReadObjectResponse.deserialize(wire) + + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[resp, None]) + + seen_types = [] + + class RecordingSink: + def __init__(self): + self.chunks = [] + + def write(self, chunk): + seen_types.append(type(chunk)) + self.chunks.append(chunk) + return len(chunk) + + sink = RecordingSink() + await mock_mrd.download_ranges( + [(0, 18, sink)], accepts_memoryview=accepts_memoryview + ) + + assert seen_types == [expected_type] + assert b"".join(sink.chunks) == data + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_writes_into_writable_buffer( + self, mock_cls_async_read_object_stream, mock_random_int, monkeypatch + ): + from google.cloud.storage.asyncio import _zero_copy_bidi_read + + _accept_memoryview_in_crc32c(monkeypatch) + data = b"these_are_18_chars" + crc1 = google_crc32c.value(data[:10]) + crc2 = google_crc32c.value(data[10:]) + + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456, 789] + + resp = _zero_copy_bidi_read.deserialize( + _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=data[:10], crc32c=crc1 + ), + range_end=False, + read_range=_storage_v2.ReadRange( + read_offset=0, read_length=10, read_id=456 + ), + ), + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=data[10:], crc32c=crc2 + ), + range_end=True, + read_range=_storage_v2.ReadRange( + read_offset=10, read_length=8, read_id=456 + ), + ), + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=data[10:], crc32c=crc2 + ), + range_end=True, + read_range=_storage_v2.ReadRange( + read_offset=10, read_length=8, read_id=789 + ), + ), + ] + ) + ) + ) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[resp, None]) + + raw = bytearray(18) + backing = bytearray(b"." * 12) + slice_view = memoryview(backing)[2:10] + + await mock_mrd.download_ranges([(0, 18, raw), (10, 8, slice_view)]) + + assert bytes(raw) == data + assert bytes(backing) == b".." + data[10:] + b".." + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "length, buf, match", + [ + (0, bytearray(10), "requires bytes_to_read > 0"), + (4, memoryview(b"abcd"), "writable, C-contiguous"), + (4, memoryview(bytearray(8))[::2], "writable, C-contiguous"), + (10, bytearray(4), "smaller than bytes_to_read"), + ], + ) + async def test_download_ranges_rejects_invalid_writable_buffer( + self, mock_cls_async_read_object_stream, length, buf, match + ): + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + with pytest.raises(ValueError, match=match): + await mock_mrd.download_ranges([(0, length, buf)]) + + def test_buffer_destination_rejects_chunk_past_bytes_to_read(self): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + _BufferWriter, + ) + + # A buffer may be larger than bytes_to_read; bytes past the range are + # never written, even if a chunk runs past it. + buffer = bytearray(b"." * 8) + writer = _BufferWriter(buffer, 4) + writer.write(b"ab") + with pytest.raises(ValueError): + writer.write(b"cde") + assert bytes(buffer) == b"ab......" + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + @pytest.mark.parametrize("initial", [b"\0" * 10, b"B" * 20]) + async def test_download_ranges_rejects_chunk_past_bytes_to_read( + self, + mock_cls_async_read_object_stream, + mock_random_int, + monkeypatch, + initial, + ): + # The server sends 12 bytes for a 10-byte range, in a message that is + # not the range's last one. + _accept_memoryview_in_crc32c(monkeypatch) + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456] + response = _zero_copy_response(b"A" * 12, 0, 456, False) + assert isinstance( + response.object_data_ranges[0].checksummed_data.content, memoryview + ) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.close = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[response, None]) + buffer = bytearray(initial) + + with pytest.raises(DataCorruption, match="Byte count mismatch"): + await mock_mrd.download_ranges([(0, 10, buffer)]) + + assert bytes(buffer) == initial + assert not mock_mrd.is_stream_open + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_raises_on_short_read_into_buffer( + self, mock_cls_async_read_object_stream, mock_random_int, monkeypatch + ): + # GCS clamps reads extending past EOF and closes the range after the + # remaining bytes; a fixed buffer destination has no cursor or return + # count, so unwritten tail bytes must not be left unnoticed. + _accept_memoryview_in_crc32c(monkeypatch) + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456] + response = _zero_copy_response(b"6789", 6, 456, True) + assert isinstance( + response.object_data_ranges[0].checksummed_data.content, memoryview + ) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.close = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[response, None]) + buffer = bytearray(b"STALE_TAIL") + + with pytest.raises( + BufferError, match="expected 10 bytes at offset 6, got 4" + ) as exc_info: + await mock_mrd.download_ranges([(6, 10, buffer)]) + + assert mock_mrd.is_stream_open + mock_mrd.read_obj_str.close.assert_not_awaited() + # exc_info.tb still references download_ranges' frame here. + assert exc_info.tb is not None + buffer.clear() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_allows_short_read_into_write_destination( + self, mock_cls_async_read_object_stream, mock_random_int, monkeypatch + ): + _accept_memoryview_in_crc32c(monkeypatch) + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456] + response = _zero_copy_response(b"6789", 6, 456, True) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[response, None]) + chunks = [] + + class RecordingSink: + def write(self, chunk): + chunks.append((type(chunk), bytes(chunk))) + return len(chunk) + + await mock_mrd.download_ranges( + [(6, 10, RecordingSink())], accepts_memoryview=True + ) + + assert chunks == [(memoryview, b"6789")] + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_releases_buffer_when_a_later_range_is_invalid( + self, mock_cls_async_read_object_stream + ): + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + buffer = bytearray(4) + + with pytest.raises(ValueError) as exc_info: + await mock_mrd.download_ranges([(0, 4, buffer), (0, 0, bytearray(4))]) + # exc_info.tb still references download_ranges' frame here. + assert exc_info.tb is not None + buffer.clear() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_checksum_mismatch_into_buffer( + self, mock_cls_async_read_object_stream, mock_random_int, monkeypatch + ): + _accept_memoryview_in_crc32c(monkeypatch) + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456] + response = _zero_copy_response(b"x" * 18, 0, 456, True, crc32c=12345) + assert isinstance( + response.object_data_ranges[0].checksummed_data.content, memoryview + ) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.close = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[response, None]) + buffer = bytearray(18) + + with pytest.raises(DataCorruption) as exc_info: + await mock_mrd.download_ranges([(0, 18, buffer)]) + assert bytes(buffer) == bytes(18) + assert not mock_mrd.is_stream_open + # exc_info.tb still references download_ranges' frame here. + assert exc_info.tb is not None + buffer.clear() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_resumes_into_buffer_after_retry( + self, mock_cls_async_read_object_stream, mock_random_int, monkeypatch + ): + _accept_memoryview_in_crc32c(monkeypatch) + data = b"these_are_18_chars" + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.return_value = 456 + + first_chunk = _zero_copy_response(data[:10], 0, 456, False) + second_chunk = _zero_copy_response(data[10:], 10, 456, True) + assert isinstance( + first_chunk.object_data_ranges[0].checksummed_data.content, memoryview + ) + assert isinstance( + second_chunk.object_data_ranges[0].checksummed_data.content, memoryview + ) + + first_stream = mock_mrd.read_obj_str + first_stream.send = AsyncMock() + first_stream.close = AsyncMock() + first_stream.is_stream_open = True + first_stream.recv = AsyncMock( + side_effect=[ + first_chunk, + exceptions.Aborted("Idle stream has been closed."), + ] + ) + second_stream = mock.MagicMock() + second_stream.open = AsyncMock() + second_stream.send = AsyncMock() + second_stream.close = AsyncMock() + second_stream.recv = AsyncMock(side_effect=[second_chunk]) + mock_cls_async_read_object_stream.side_effect = [second_stream] + buffer = bytearray(18) + + await mock_mrd.download_ranges([(0, 18, buffer)]) + + assert bytes(buffer) == data + second_stream.open.assert_awaited_once() + second_stream.send.assert_awaited_once_with( + _storage_v2.BidiReadObjectRequest( + read_ranges=[ + _storage_v2.ReadRange(read_offset=10, read_length=8, read_id=456) + ] + ) + ) + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "length, make_destination", + [(0, BytesIO), (18, lambda: bytearray(18))], + ) + async def test_download_ranges_full_object_read_of_zero_copy_chunks( + self, + mock_cls_async_read_object_stream, + mock_random_int, + monkeypatch, + length, + make_destination, + ): + # Full-object reads also feed every chunk to Checksum.update. + _accept_memoryview_in_crc32c(monkeypatch) + destination = make_destination() + data = b"these_are_18_chars" + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_mrd.persisted_size = len(data) + mock_mrd.full_obj_server_crc32c = google_crc32c.value(data) + mock_random_int.side_effect = [456] + response = _zero_copy_response(data, 0, 456, True) + assert isinstance( + response.object_data_ranges[0].checksummed_data.content, memoryview + ) + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[response, None]) + + await mock_mrd.download_ranges([(0, length, destination)], enable_checksum=True) + + written = ( + destination.getvalue() + if isinstance(destination, BytesIO) + else bytes(destination) + ) + assert written == data + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_verifies_parser_output_with_installed_crc32c( + self, mock_cls_async_read_object_stream, mock_random_int + ): + # Deliberately no monkeypatch of google_crc32c or is_supported: the + # chunks the parser hands out (memoryview or bytes, depending on the + # installed google-crc32c) must verify as-is, including the content-less + # final message of a range that ends on a chunk boundary. + from google.cloud.storage.asyncio import _zero_copy_bidi_read + + data = b"these_are_18_chars" + mock_mrd, _ = await self._make_mock_mrd(mock_cls_async_read_object_stream) + mock_random_int.side_effect = [456, 789] + + def response(content, offset, range_end): + return _zero_copy_bidi_read.deserialize( + _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=content, + crc32c=google_crc32c.value(content), + ), + range_end=range_end, + read_range=_storage_v2.ReadRange( + read_offset=offset, + read_length=len(content), + read_id=read_id, + ), + ) + for read_id in (456, 789) + ] + ) + ) + ) + + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock( + side_effect=[response(data, 0, False), response(b"", 18, True), None] + ) + + sink = BytesIO() + raw = bytearray(18) + await mock_mrd.download_ranges([(0, 18, sink), (0, 18, raw)]) + + assert sink.getvalue() == data + assert bytes(raw) == data + @mock.patch( "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer" ) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py index 393fa7b2b709..608ff1178a88 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py @@ -113,6 +113,59 @@ def test_init_with_bucket_object_generation(mock_client, mock_async_bidi_rpc): assert read_obj_stream.rpc == rpc_sentinel +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream._zero_copy_bidi_read.wrapped_rpc" +) +def test_init_prefers_zero_copy_rpc(mock_wrapped_rpc, mock_client, mock_async_bidi_rpc): + # Arrange + mock_wrapped_rpc.return_value = mock.sentinel.zero_copy_rpc + mock_client._client._transport._wrapped_methods = {} + + # Act + read_obj_stream = _AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + ) + + # Assert + assert read_obj_stream.rpc is mock.sentinel.zero_copy_rpc + mock_wrapped_rpc.assert_called_once_with(mock_client._client._transport) + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream._zero_copy_bidi_read.wrapped_rpc" +) +def test_init_falls_back_to_generated_rpc( + mock_wrapped_rpc, mock_client, mock_async_bidi_rpc +): + # Arrange + mock_wrapped_rpc.return_value = None + rpc_sentinel = mock.sentinel.A + mock_client._client._transport.bidi_read_object = "bidi_read_object_rpc" + mock_client._client._transport._wrapped_methods = { + "bidi_read_object_rpc": rpc_sentinel, + } + + # Act + read_obj_stream = _AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + ) + + # Assert + assert read_obj_stream.rpc is rpc_sentinel + + @mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") @mock.patch( "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py new file mode 100644 index 000000000000..d6c43fcbb43a --- /dev/null +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py @@ -0,0 +1,790 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the zero-copy ``BidiReadObjectResponse`` deserializer.""" + +import logging +from unittest import mock + +import google_crc32c +import grpc +import pytest +from google.api_core import client_info as api_core_client_info +from google.api_core.gapic_v1.client_info import ClientInfo +from google.auth.credentials import AnonymousCredentials +from google.protobuf.message import DecodeError + +from google.cloud import _storage_v2 +from google.cloud._storage_v2.services.storage.transports.base import ( + DEFAULT_CLIENT_INFO, +) +from google.cloud._storage_v2.services.storage.transports.grpc_asyncio import ( + StorageGrpcAsyncIOTransport, +) +from google.cloud.storage.asyncio import _zero_copy_bidi_read +from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient +from google.cloud.storage.asyncio.async_multi_range_downloader import ( + AsyncMultiRangeDownloader, +) +from google.cloud.storage.asyncio.async_read_object_stream import ( + _AsyncReadObjectStream, +) + +PAYLOAD = bytes(range(256)) * 8 +_orig_crc32c_value = google_crc32c.value +_orig_checksum_update = google_crc32c.Checksum.update + + +@pytest.fixture(autouse=True) +def _crc32c_accepts_memoryview(monkeypatch): + monkeypatch.setattr( + google_crc32c, "value", lambda data: _orig_crc32c_value(bytes(data)) + ) + monkeypatch.setattr( + google_crc32c.Checksum, + "update", + lambda self, chunk: _orig_checksum_update(self, bytes(chunk)), + ) + # is_supported() caches its probe; re-probe the patched functions. + monkeypatch.setattr(_zero_copy_bidi_read, "_memoryview_supported", None) + + +# Field numbers from google/storage/v2/storage.proto. +_RESPONSE_METADATA, _RESPONSE_RANGES, _RESPONSE_HANDLE = 4, 6, 7 +_RANGE_CHECKSUMMED, _RANGE_READ_RANGE, _RANGE_END = 1, 2, 3 +_CHECKSUMMED_CONTENT, _CHECKSUMMED_CRC32C = 1, 2 +_READ_RANGE_OFFSET, _READ_RANGE_LENGTH, _READ_RANGE_ID = 1, 2, 3 +_HANDLE_HANDLE = 1 +_UNKNOWN = 15 + +# --- minimal wire-format encoder ----------------------------------------------- + + +def _varint(value): + out = bytearray() + while True: + byte = value & 0x7F + value >>= 7 + if not value: + out.append(byte) + return bytes(out) + out.append(byte | 0x80) + + +def _tag(number, wire_type): + return _varint(number << 3 | wire_type) + + +def _varint_field(number, value): + return _tag(number, 0) + _varint(value) + + +def _i64_field(number, value): + return _tag(number, 1) + value.to_bytes(8, "little") + + +def _len_field(number, payload): + return _tag(number, 2) + _varint(len(payload)) + payload + + +def _i32_field(number, value): + return _tag(number, 5) + value.to_bytes(4, "little") + + +def _read_range(offset=4096, length=len(PAYLOAD), read_id=7, extra=b""): + return ( + _varint_field(_READ_RANGE_OFFSET, offset) + + _varint_field(_READ_RANGE_LENGTH, length) + + _varint_field(_READ_RANGE_ID, read_id) + + extra + ) + + +def _checksummed(content=PAYLOAD, crc32c=None, extra=b""): + body = _len_field(_CHECKSUMMED_CONTENT, content) + if crc32c is not None: + body += _i32_field(_CHECKSUMMED_CRC32C, crc32c) + return body + extra + + +def _range_data(checksummed=None, read_range=None, range_end=True, extra=b""): + body = b"" + if checksummed is not None: + body += _len_field(_RANGE_CHECKSUMMED, checksummed) + if read_range is not None: + body += _len_field(_RANGE_READ_RANGE, read_range) + if range_end: + body += _varint_field(_RANGE_END, 1) + return body + extra + + +def _response(*ranges, handle=None, extra=b""): + body = b"".join(_len_field(_RESPONSE_RANGES, r) for r in ranges) + if handle is not None: + body += _len_field(_RESPONSE_HANDLE, handle) + return body + extra + + +def _data_response(content=PAYLOAD, crc32c=None, handle=None): + return _response( + _range_data(_checksummed(content, crc32c), _read_range(length=len(content))), + handle=None if handle is None else _len_field(_HANDLE_HANDLE, handle), + ) + + +def _assert_same_as_generated(buf): + """Parse ``buf`` with both parsers and check the zero-copy one matches.""" + expected = _storage_v2.BidiReadObjectResponse.deserialize(buf) + actual = _zero_copy_bidi_read.deserialize(buf) + + assert isinstance(actual, _zero_copy_bidi_read._ZeroCopyBidiReadObjectResponse) + assert actual.metadata is None + assert not actual.HasField("metadata") + assert actual.HasField("read_handle") == expected._pb.HasField("read_handle") + if actual.HasField("read_handle"): + assert isinstance(actual.read_handle, _storage_v2.BidiReadHandle) + assert actual.read_handle.handle == expected.read_handle.handle + else: + assert actual.read_handle is None + + assert len(actual.object_data_ranges) == len(expected.object_data_ranges) + for got, want in zip(actual.object_data_ranges, expected.object_data_ranges): + assert got.HasField("read_range") == want._pb.HasField("read_range") + assert got.HasField("checksummed_data") == want._pb.HasField("checksummed_data") + assert got.read_range.read_offset == want.read_range.read_offset + assert got.read_range.read_length == want.read_range.read_length + assert got.read_range.read_id == want.read_range.read_id + assert got.range_end == want.range_end + # Non-empty payloads are views of the wire buffer, empty ones b"". + assert type(got.checksummed_data.content) is ( + memoryview if want.checksummed_data.content else bytes + ) + assert bytes(got.checksummed_data.content) == want.checksummed_data.content + assert got.checksummed_data.HasField( + "crc32c" + ) == want.checksummed_data._pb.HasField("crc32c") + assert got.checksummed_data.crc32c == want.checksummed_data.crc32c + # Code that unwraps proto-plus messages with ``._pb`` sees the same, + # including the proto3 ``optional`` field's synthetic oneof. + assert got._pb.HasField("read_range") == want._pb.HasField("read_range") + assert got.checksummed_data._pb.HasField( + "_crc32c" + ) == want.checksummed_data._pb.HasField("_crc32c") + return actual + + +# --- deserialize --------------------------------------------------------------- + + +# Per-case checks for test_deserialize, run after the equivalence assertions. + + +def _check_field_values(buf, response): + # Literal values, independent of the generated parser used as the oracle. + [range_data] = response.object_data_ranges + assert range_data.read_range.read_offset == 4096 + assert range_data.read_range.read_length == len(PAYLOAD) + assert range_data.read_range.read_id == 7 + assert range_data.range_end is True + assert bytes(range_data.checksummed_data.content) == PAYLOAD + assert range_data.checksummed_data.crc32c == google_crc32c.value(PAYLOAD) + + +def _check_content_is_a_view_of_the_wire_buffer(buf, response): + assert response.object_data_ranges[0].checksummed_data.content.obj is buf + + +def _check_read_handle_is_a_bidi_read_handle(buf, response): + assert response.read_handle.handle == b"opaque" + # The SDK feeds a captured handle straight back into the next open. + spec = _storage_v2.BidiReadObjectSpec(read_handle=response.read_handle) + assert spec.read_handle.handle == b"opaque" + + +_UNKNOWN_FIELDS = ( + _varint_field(_UNKNOWN, 1) + + _i64_field(_UNKNOWN, 2) + + _i32_field(_UNKNOWN, 3) + + _len_field(_UNKNOWN, b"ignored") +) + + +@pytest.mark.parametrize( + "buf, check", + [ + pytest.param( + _data_response(crc32c=google_crc32c.value(PAYLOAD)), + _check_field_values, + id="data", + ), + pytest.param( + _data_response(), _check_content_is_a_view_of_the_wire_buffer, id="view" + ), + pytest.param( + _data_response(handle=b"opaque"), + _check_read_handle_is_a_bidi_read_handle, + id="read_handle", + ), + pytest.param( + _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=PAYLOAD, crc32c=google_crc32c.value(PAYLOAD) + ), + read_range=_storage_v2.ReadRange( + read_offset=4096, read_length=len(PAYLOAD), read_id=7 + ), + range_end=True, + ) + ], + read_handle=_storage_v2.BidiReadHandle(handle=b"opaque"), + ) + ), + None, + id="proto_plus_serialization", + ), + pytest.param(_data_response(), None, id="without_crc32c"), + pytest.param( + _response( + _range_data( + _checksummed(b"first"), _read_range(0, 5, 1), range_end=False + ), + _range_data(_checksummed(b"second"), _read_range(5, 6, 2)), + ), + None, + id="multiple_ranges", + ), + pytest.param(b"", None, id="empty_message"), + pytest.param( + _response(_range_data(_checksummed())), None, id="range_without_read_range" + ), + pytest.param( + _response(_range_data(read_range=_read_range())), + None, + id="range_without_checksummed_data", + ), + pytest.param( + _data_response() + _UNKNOWN_FIELDS, None, id="unknown_top_level_fields" + ), + pytest.param( + _response( + _range_data( + _checksummed(extra=_UNKNOWN_FIELDS), + _read_range(extra=_UNKNOWN_FIELDS), + extra=_UNKNOWN_FIELDS, + ), + handle=_len_field(_HANDLE_HANDLE, b"opaque") + _UNKNOWN_FIELDS, + ), + None, + id="unknown_nested_fields", + ), + pytest.param( + _response( + _range_data( + read_range=_tag(_READ_RANGE_ID, 0) + (b"\xff" * 9 + b"\x7f") + ) + ), + None, + id="over_long_varint_keeps_low_64_bits", + ), + pytest.param( + _response( + _range_data( + extra=_len_field( + _RANGE_READ_RANGE, _varint_field(_READ_RANGE_OFFSET, 4096) + ) + + _len_field(_RANGE_READ_RANGE, _varint_field(_READ_RANGE_ID, 7)) + + _len_field( + _RANGE_CHECKSUMMED, + _len_field(_CHECKSUMMED_CONTENT, PAYLOAD), + ) + + _len_field( + _RANGE_CHECKSUMMED, + _i32_field(_CHECKSUMMED_CRC32C, 1234), + ) + ), + handle=_len_field(_HANDLE_HANDLE, b"opaque"), + ) + + _len_field(_RESPONSE_HANDLE, b""), + None, + id="merges_split_sub_messages", + ), + ], +) +def test_deserialize(buf, check): + response = _assert_same_as_generated(buf) + if check is not None: + check(buf, response) + + +@pytest.mark.parametrize( + "offset, length, read_id", + [ + pytest.param(-1, -1, -1, id="minus_one"), + pytest.param(-5, 0, -(1 << 55), id="negative"), + pytest.param(-(1 << 63), (1 << 63) - 1, 1, id="int64_extremes"), + ], +) +def test_deserialize_signed_int64_read_range(offset, length, read_id): + # read_offset, read_length and read_id are int64; negative values travel + # as ten-byte two's-complement varints. + message = _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + read_range=_storage_v2.ReadRange( + read_offset=offset, read_length=length, read_id=read_id + ) + ) + ] + ) + response = _assert_same_as_generated( + _storage_v2.BidiReadObjectResponse.serialize(message) + ) + read_range = response.object_data_ranges[0].read_range + assert (read_range.read_offset, read_range.read_length, read_range.read_id) == ( + offset, + length, + read_id, + ) + + +def test_deserialize_metadata_response_uses_generated_parser(): + metadata = _storage_v2.Object(name="o", bucket="projects/_/buckets/b", size=10) + buf = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse(metadata=metadata) + ) + response = _zero_copy_bidi_read.deserialize(buf) + assert isinstance(response, _storage_v2.BidiReadObjectResponse) + assert response.metadata.size == 10 + + +def test_deserialize_unknown_group_uses_generated_parser(): + # Start-group/end-group wire types are valid protobuf that the zero-copy + # path does not understand; the generated parser keeps them as unknown fields. + buf = _data_response() + _tag(_UNKNOWN, 3) + _tag(_UNKNOWN, 4) + response = _zero_copy_bidi_read.deserialize(buf) + assert isinstance(response, _storage_v2.BidiReadObjectResponse) + assert response.object_data_ranges[0].checksummed_data.content == PAYLOAD + + +@pytest.mark.parametrize( + "buf", + [ + pytest.param( + # object_data_ranges (6 bytes) whose checksummed_data claims 8 bytes, + # followed by a well-formed read_handle. + bytes([0x32, 0x06, 0x0A, 0x08]) + b"abcd" + bytes([0x3A, 0x02, 0x0A, 0x00]), + id="inner_length_past_enclosing_message", + ), + pytest.param( + bytes([0x32]) + b"\x80" * 10 + b"\x00", id="varint_longer_than_10_bytes" + ), + pytest.param(bytes([0x32, 0x10, 0x0A, 0x00]), id="truncated_message"), + pytest.param( + _response(_range_data(_checksummed()))[:-3], id="truncated_tag_varint" + ), + pytest.param( + _response(_range_data(_checksummed(extra=_tag(_CHECKSUMMED_CRC32C, 5)))), + id="crc32c_past_enclosing_message", + ), + pytest.param( + _response(_range_data(read_range=_tag(_READ_RANGE_OFFSET, 0) + b"\x80")), + id="read_range_varint_runs_past_enclosing_message", + ), + pytest.param( + _response(_range_data(range_end=False, extra=_tag(_RANGE_END, 0) + b"\x80")) + + _varint_field(_UNKNOWN, 1), + id="range_end_varint_runs_past_enclosing_message", + ), + pytest.param( + _data_response() + _tag(_UNKNOWN, 1) + b"\x00" * 3, + id="truncated_fixed64", + ), + pytest.param( + _data_response() + _tag(_UNKNOWN, 5) + b"\x00", + id="truncated_fixed32", + ), + pytest.param(b"\x00\x00", id="field_number_zero"), + pytest.param( + _response(_range_data(extra=b"\x00\x00")), + id="field_number_zero_in_range_data", + ), + pytest.param( + _response(_range_data(_checksummed(extra=b"\x00\x00"))), + id="field_number_zero_in_checksummed_data", + ), + pytest.param( + _response(_range_data(read_range=_read_range(extra=b"\x00\x00"))), + id="field_number_zero_in_read_range", + ), + pytest.param( + _response(handle=b"\x00\x00"), + id="field_number_zero_in_read_handle", + ), + pytest.param( + # Wire type 6 ensures the pure-Python protobuf backend also raises + # DecodeError after _tag() rejects the >32-bit tag value. + _tag(1 << 29, 6) + b"\x01" + _data_response(), + id="field_number_too_large", + ), + pytest.param( + # Field 6 / wire type 6 padded to a six-byte varint: _tag() rejects + # i - start > 5 before inspecting the wire type, and both upb and + # the pure-Python backend raise DecodeError on the fallback call. + b"\xb6\x80\x80\x80\x80\x00" + _varint(2) + _varint_field(_RANGE_END, 1), + id="tag_longer_than_five_bytes", + ), + ], +) +def test_deserialize_malformed_input_raises_like_generated_parser(buf): + with pytest.raises(DecodeError): + _storage_v2.BidiReadObjectResponse.deserialize(buf) + with pytest.raises(DecodeError): + _zero_copy_bidi_read.deserialize(buf) + + +@pytest.mark.parametrize( + "buf", + [ + pytest.param(b"\x32\x80\x80\x80\x80\x80\x00", id="six_byte_length"), + pytest.param(b"\x32" + b"\x80" * 9 + b"\x02", id="length_wraps_to_zero"), + pytest.param( + b"\x32\x09\x0a\x07\x0a\x80\x80\x80\x80\x80\x00", + id="nested_six_byte_length", + ), + ], +) +def test_deserialize_defers_overlong_length_prefix_to_generated_parser( + monkeypatch, buf +): + # upb in protobuf 7 rejects a length encoded in more than five bytes while + # protobuf 6 and the pure-Python backend accept it, so the zero-copy + # parser leaves the verdict to the installed generated parser. + monkeypatch.setattr( + _zero_copy_bidi_read, + "_generated_deserialize", + lambda data: mock.sentinel.generated, + ) + assert _zero_copy_bidi_read.deserialize(buf) is mock.sentinel.generated + + +def test_has_field_rejects_unknown_names(): + response = _zero_copy_bidi_read.deserialize(_data_response(crc32c=1)) + [range_data] = response.object_data_ranges + for message in (response, range_data, range_data.checksummed_data): + with pytest.raises(ValueError): + message.HasField("no_such_field") + + +def test_zero_copy_response_pb_is_itself(): + # ReadsResumptionStrategy unwraps ``response._pb`` before calling HasField. + response = _zero_copy_bidi_read.deserialize(_data_response()) + assert response._pb is response + + +# --- is_supported / is_enabled ------------------------------------------------- + + +def test_is_supported_requires_crc32c_to_accept_memoryview(monkeypatch): + probed_with = [] + + def bytes_only(data): + probed_with.append(type(data)) + raise TypeError("argument 1 must be read-only bytes-like object") + + monkeypatch.setattr(google_crc32c, "value", bytes_only) + assert _zero_copy_bidi_read.is_supported() is False + assert probed_with == [memoryview] + monkeypatch.setattr(google_crc32c, "value", lambda data: len(data)) + monkeypatch.setattr(_zero_copy_bidi_read, "_memoryview_supported", None) + assert _zero_copy_bidi_read.is_supported() is True + + +def test_is_supported_probes_crc32c_once(monkeypatch): + probes = [] + + def counting_value(data): + probes.append(type(data)) + return 0 + + monkeypatch.setattr(google_crc32c, "value", counting_value) + assert _zero_copy_bidi_read.is_supported() is True + assert _zero_copy_bidi_read.is_supported() is True + assert probes == [memoryview] + + +def test_is_supported_requires_checksum_update_to_accept_memoryview(monkeypatch): + # Full-object reads also feed each chunk to Checksum.update. + probed_with = [] + + def bytes_only(self, chunk): + probed_with.append(type(chunk)) + raise TypeError("argument 2 must be read-only bytes-like object") + + monkeypatch.setattr(google_crc32c.Checksum, "update", bytes_only) + assert _zero_copy_bidi_read.is_supported() is False + assert probed_with == [memoryview] + + +def test_deserialize_slices_bytes_when_memoryview_crc32c_not_supported(monkeypatch): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: False) + response = _zero_copy_bidi_read.deserialize(_data_response(crc32c=1)) + assert isinstance(response, _zero_copy_bidi_read._ZeroCopyBidiReadObjectResponse) + content = response.object_data_ranges[0].checksummed_data.content + assert type(content) is bytes + assert content == PAYLOAD + + +@pytest.mark.parametrize("supported", [False, True]) +def test_deserialize_empty_chunk_content_is_bytes(monkeypatch, supported): + # The last message of a read that ends on a chunk boundary carries a + # checksum and range_end but no content field; crc32c releases that do + # not take memoryview must still be able to hash it. + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: supported) + buf = _response( + _range_data( + _i32_field(_CHECKSUMMED_CRC32C, google_crc32c.value(b"")), + _read_range(length=0), + ) + ) + response = _zero_copy_bidi_read.deserialize(buf) + checksummed = response.object_data_ranges[0].checksummed_data + assert type(checksummed.content) is bytes + assert checksummed.content == b"" + assert checksummed.HasField("crc32c") + _assert_same_as_generated(buf) + + +@pytest.mark.parametrize("value", [None, "", "0", "false", "No", " off "]) +def test_is_enabled_defaults_to_off(monkeypatch, value): + assert _zero_copy_bidi_read._ENV_VAR == "GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ" + if value is None: + monkeypatch.delenv(_zero_copy_bidi_read._ENV_VAR, raising=False) + else: + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, value) + assert _zero_copy_bidi_read.is_enabled() is False + + +@pytest.mark.parametrize("value", ["1", "true", "Yes", " on "]) +def test_is_enabled_by_environment(monkeypatch, value): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, value) + assert _zero_copy_bidi_read.is_enabled() is True + + +@pytest.mark.parametrize("wrap", [bytearray, memoryview]) +def test_deserialize_only_aliases_immutable_bytes(wrap): + # A mutable or recyclable buffer must not be aliased by the chunks. + response = _zero_copy_bidi_read.deserialize(wrap(_data_response(crc32c=1))) + assert isinstance(response, _storage_v2.BidiReadObjectResponse) + assert isinstance(response.object_data_ranges[0].checksummed_data.content, bytes) + + +# --- wrapped_rpc --------------------------------------------------------------- + + +def _transport(target): + # Built like AsyncGrpcClient's transports, without client_info. + return StorageGrpcAsyncIOTransport( + channel=grpc.aio.insecure_channel(target), + credentials=AnonymousCredentials(), + ) + + +def test_wrapped_rpc_ignores_transports_it_does_not_know(monkeypatch): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1") + assert _zero_copy_bidi_read.wrapped_rpc(object()) is None + + +@pytest.mark.asyncio +async def test_wrapped_rpc_returns_none_when_not_enabled(monkeypatch): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "0") + transport = _transport("localhost:1") + try: + assert _zero_copy_bidi_read.wrapped_rpc(transport) is None + finally: + await transport.close() + + +@pytest.mark.asyncio +async def test_wrapped_rpc_is_built_once_per_transport(monkeypatch): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1") + transport = _transport("localhost:1") + other = _transport("localhost:1") + try: + rpc = _zero_copy_bidi_read.wrapped_rpc(transport) + assert rpc is not None + assert rpc is not transport._wrapped_methods[transport.bidi_read_object] + assert _zero_copy_bidi_read.wrapped_rpc(transport) is rpc + assert _zero_copy_bidi_read.wrapped_rpc(other) is not rpc + finally: + await transport.close() + await other.close() + + +# --- end to end through _AsyncReadObjectStream --------------------------------- + + +class _FakeStorageServer: + """In-process gRPC server that answers BidiReadObject with canned bytes.""" + + def __init__(self, responses): + self._responses = list(responses) + self.requests = [] + self.invocation_metadata = None + self._server = grpc.aio.server() + handler = grpc.method_handlers_generic_handler( + "google.storage.v2.Storage", + {"BidiReadObject": grpc.stream_stream_rpc_method_handler(self._handle)}, + ) + self._server.add_generic_rpc_handlers((handler,)) + self._port = self._server.add_insecure_port("localhost:0") + + @property + def target(self): + return f"localhost:{self._port}" + + async def _handle(self, request_iterator, context): + self.invocation_metadata = dict(context.invocation_metadata()) + async for request in request_iterator: + self.requests.append(_storage_v2.BidiReadObjectRequest.deserialize(request)) + yield self._responses.pop(0) + + async def __aenter__(self): + await self._server.start() + return self + + async def __aexit__(self, *exc_info): + await self._server.stop(None) + + +@pytest.fixture +def keep_google_logger_propagation(monkeypatch): + # Building a real GAPIC client runs api-core's client_logging + # initialization, which switches off propagation on the "google" logger + # and would hide log records from ``caplog`` in later test modules. + logger = logging.getLogger("google") + monkeypatch.setattr(logger, "propagate", logger.propagate) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [True, False]) +async def test_read_object_stream_end_to_end( + monkeypatch, keep_google_logger_propagation, enabled +): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1" if enabled else "0") + crc = google_crc32c.value(PAYLOAD) + first = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + metadata=_storage_v2.Object( + name="o", bucket="projects/_/buckets/b", generation=5, size=2048 + ), + read_handle=_storage_v2.BidiReadHandle(handle=b"opaque"), + ) + ) + second = _data_response(crc32c=crc) + + async with _FakeStorageServer([first, second]) as server: + transport = _transport(server.target) + client = _storage_v2.StorageAsyncClient(transport=transport) + stream = _AsyncReadObjectStream( + client=client, + bucket_name="b", + object_name="o", + ) + try: + await stream.open() + assert stream.generation_number == 5 + assert stream.persisted_size == 2048 + assert stream.read_handle.handle == b"opaque" + + await stream.send( + _storage_v2.BidiReadObjectRequest( + read_ranges=[ + _storage_v2.ReadRange( + read_offset=4096, read_length=len(PAYLOAD), read_id=7 + ) + ] + ) + ) + response = await stream.recv() + await stream.close() + finally: + await transport.close() + + [range_data] = response.object_data_ranges + content = range_data.checksummed_data.content + assert isinstance(content, memoryview if enabled else bytes) + assert bytes(content) == PAYLOAD + assert range_data.checksummed_data.crc32c == crc + assert range_data.read_range.read_id == 7 + assert ( + isinstance(response, _zero_copy_bidi_read._ZeroCopyBidiReadObjectResponse) + is enabled + ) + + assert server.requests[0].read_object_spec.object_ == "o" + assert server.requests[1].read_ranges[0].read_id == 7 + assert server.invocation_metadata["x-goog-request-params"] == ( + "bucket=projects/_/buckets/b" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [True, False]) +@pytest.mark.parametrize( + "client_info", + [ + pytest.param(None, id="default"), + pytest.param( + api_core_client_info.ClientInfo(user_agent="gcsfs/1"), id="api_core" + ), + pytest.param(ClientInfo(user_agent="gcsfs/1"), id="gapic_v1"), + ], +) +async def test_downloader_sends_same_api_client_header_with_either_parser( + monkeypatch, keep_google_logger_propagation, enabled, client_info +): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1" if enabled else "0") + first = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + metadata=_storage_v2.Object(name="o", bucket="projects/_/buckets/b") + ) + ) + + async with _FakeStorageServer([first]) as server: + transports = [] + + def create_grpc_client(**kwargs): + # Like AsyncGrpcClient, build the transport without client_info, + # after the constructor has updated the ClientInfo it was given. + transports.append(_transport(server.target)) + return _storage_v2.StorageAsyncClient(transport=transports[0]) + + with mock.patch.object( + AsyncGrpcClient, "_create_async_grpc_client", side_effect=create_grpc_client + ): + client = AsyncGrpcClient(credentials=mock.Mock(), client_info=client_info) + try: + mrd = await AsyncMultiRangeDownloader.create_mrd(client, "b", "o") + await mrd.close() + finally: + await transports[0].close() + + # The generated stub sends the transport's default header on every RPC. + header = server.invocation_metadata["x-goog-api-client"] + assert header.split() == DEFAULT_CLIENT_INFO.to_user_agent().split()