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..4db72b6f3bd8 --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py @@ -0,0 +1,332 @@ +# 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, so callers never see a difference +other than the type of ``content``. + +: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.api_core import gapic_v1 +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, +) + +# Set ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=0`` (or ``false``, ``no``, ``off``) +# to disable the wire parser and use the generated parser instead. +_ENV_VAR = "GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ" +_DISABLED_VALUES = ("0", "false", "no", "off") + +# 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 + + def HasField(self, name): + if name == "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 + + 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``.""" + n, i = _varint(buf, i, end) + 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) + + +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. + """ + try: + google_crc32c.value(memoryview(b"\0")) + except Exception: + return False + return True + + +def is_enabled(): + """Whether a stream should use the zero-copy parser. + + Enabled by default; set ``GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=0`` + (or ``false``, ``no``, ``off``) to keep the generated parser. + """ + return os.environ.get(_ENV_VAR, "").strip().lower() not in _DISABLED_VALUES + + +_rpcs = weakref.WeakKeyDictionary() + + +def wrapped_rpc(transport, client_info=None): + """Return the zero-copy ``BidiReadObject`` callable for ``transport``. + + The callable is the generated transport's stub with :func:`deserialize` + swapped in, wrapped exactly like ``transport._wrapped_methods`` entries + (error mapping, no default retry or timeout, the ``x-goog-api-client`` + header from ``client_info``, the transport's own by default). 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, + ) + if not hasattr(client_info, "to_grpc_metadata"): + client_info = DEFAULT_CLIENT_INFO + rpc = gapic_v1.method_async.wrap_method( + stub, default_timeout=None, client_info=client_info + ) + _rpcs[transport] = rpc + return rpc diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py index 18ee348401c7..19ea106d4bc7 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py @@ -91,6 +91,7 @@ def __init__( agent_version = f"gcloud-python/{__version__}" if agent_version not in client_info.user_agent: client_info.user_agent += f" {agent_version} " + self._client_info = client_info self._grpc_client = self._create_async_grpc_client( credentials=credentials, 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..3a90a123ede7 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 @@ -78,6 +78,56 @@ 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) + if view.readonly or not view.c_contiguous: + view.release() + raise ValueError( + "Buffer destination must be a writable, C-contiguous buffer" + ) + nbytes = view.nbytes + if nbytes < length: + view.release() + raise ValueError( + f"Buffer destination ({nbytes} bytes) is smaller than " + f"bytes_to_read ({length})" + ) + self._view = view.cast("B") + 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 _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. @@ -292,6 +342,7 @@ async def _do_open(): object_name=self.object_name, generation_number=self.generation, read_handle=self.read_handle, + client_info=getattr(self.client, "_client_info", None), ) if self._routing_token: @@ -331,6 +382,7 @@ async def factory(): object_name=self.object_name, generation_number=self.generation, read_handle=current_handle, + client_info=getattr(self.client, "_client_info", None), ) current_metadata = list(metadata) if metadata else [] @@ -365,6 +417,7 @@ async def download_ranges( 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. @@ -376,9 +429,21 @@ async def download_ranges( 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 from the gRPC + response buffer without an intermediate ``bytes`` allocation. + 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. @@ -395,6 +460,16 @@ 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``. The SDK drops its + reference after ``write()`` returns; the message (the chunk plus + a few bytes of framing) is freed when the destination drops the + view. Destinations must still accept ``bytes``, which the generated + parser produces. 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 DataCorruption: if a checksum mismatch is detected while reading data. @@ -433,7 +508,7 @@ async def download_ranges( download_states[read_id] = _DownloadState( initial_offset=offset, initial_length=length, - user_buffer=user_buffer, + user_buffer=_wrap_user_buffer(user_buffer, length, accepts_memoryview), is_full_object_read=is_full_object_read, ) 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..fb4ac0b1f360 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, ) @@ -58,6 +59,7 @@ def __init__( object_name: str, generation_number: Optional[int] = None, read_handle: Optional[_storage_v2.BidiReadHandle] = None, + client_info=None, ) -> None: if client is None: raise ValueError("client must be provided") @@ -76,9 +78,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, client_info) + 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/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..9a8ab4934059 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 @@ -84,6 +84,7 @@ async def test_create_mrd(self, mock_cls_async_read_object_stream): object_name=_TEST_OBJECT_NAME, generation_number=_TEST_GENERATION_NUMBER, read_handle=_TEST_READ_HANDLE, + client_info=mock_client._client_info, ) mrd.read_obj_str.open.assert_called_once() @@ -98,6 +99,168 @@ 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( + "accepts_memoryview, expected_type", [(False, bytes), (True, memoryview)] + ) + async def test_download_ranges_write_destination_chunk_type( + self, + mock_cls_async_read_object_stream, + mock_random_int, + monkeypatch, + accepts_memoryview, + expected_type, + ): + from google.cloud.storage.asyncio import _zero_copy_bidi_read + + orig_crc32c = google_crc32c.value + monkeypatch.setattr( + google_crc32c, "value", lambda data: orig_crc32c(bytes(data)) + ) + 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 + ), + ) + ] + ) + ) + zero_copy_resp = _zero_copy_bidi_read.deserialize(wire) + assert isinstance( + zero_copy_resp.object_data_ranges[0].checksummed_data.content, memoryview + ) + + mock_mrd.read_obj_str.send = AsyncMock() + mock_mrd.read_obj_str.recv = AsyncMock(side_effect=[zero_copy_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 + + orig_crc32c = google_crc32c.value + monkeypatch.setattr( + google_crc32c, "value", lambda data: orig_crc32c(bytes(data)) + ) + 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)]) + @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..fc2ada8b7df0 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, None) + + +@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..fc9edea94f61 --- /dev/null +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py @@ -0,0 +1,703 @@ +# 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 + +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_read_object_stream import ( + _AsyncReadObjectStream, +) + +PAYLOAD = bytes(range(256)) * 8 +_orig_crc32c_value = google_crc32c.value + + +@pytest.fixture(autouse=True) +def _crc32c_accepts_memoryview(monkeypatch): + monkeypatch.setattr( + google_crc32c, "value", lambda data: _orig_crc32c_value(bytes(data)) + ) + + +# 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 + return actual + + +# --- deserialize --------------------------------------------------------------- + + +def test_deserialize_data_response_round_trip(): + crc = google_crc32c.value(PAYLOAD) + response = _assert_same_as_generated(_data_response(crc32c=crc)) + + [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 == crc + + +def test_deserialize_content_is_a_view_of_the_wire_buffer(): + buf = _data_response() + content = _zero_copy_bidi_read.deserialize(buf).object_data_ranges[0] + assert content.checksummed_data.content.obj is buf + + +def test_deserialize_read_handle_is_a_bidi_read_handle(): + response = _assert_same_as_generated(_data_response(handle=b"opaque")) + + 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", + [ + 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"), + ) + ), + id="proto_plus_serialization", + ), + pytest.param(_data_response(), 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)), + ), + id="multiple_ranges", + ), + pytest.param(b"", id="empty_message"), + pytest.param( + _response(_range_data(_checksummed())), id="range_without_read_range" + ), + pytest.param( + _response(_range_data(read_range=_read_range())), + id="range_without_checksummed_data", + ), + pytest.param(_data_response() + _UNKNOWN_FIELDS, 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, + ), + id="unknown_nested_fields", + ), + pytest.param( + _response( + _range_data( + read_range=_tag(_READ_RANGE_ID, 0) + (b"\xff" * 9 + b"\x7f") + ) + ), + 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""), + id="merges_split_sub_messages", + ), + ], +) +def test_deserialize_matches_generated_parser(buf): + _assert_same_as_generated(buf) + + +@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) + + +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)) + assert _zero_copy_bidi_read.is_supported() is True + + +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, "", "1", "true", "Yes", " on "]) +def test_is_enabled_defaults_to_on(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 True + + +@pytest.mark.parametrize("value", ["0", "false", "No", " off "]) +def test_is_enabled_disabled_by_environment(monkeypatch, value): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, value) + assert _zero_copy_bidi_read.is_enabled() is False + + +@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, client_info=None): + return StorageGrpcAsyncIOTransport( + channel=grpc.aio.insecure_channel(target), + credentials=AnonymousCredentials(), + client_info=client_info or DEFAULT_CLIENT_INFO, + ) + + +def test_wrapped_rpc_ignores_transports_it_does_not_know(monkeypatch): + monkeypatch.delenv(_zero_copy_bidi_read._ENV_VAR, raising=False) + 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.delenv(_zero_copy_bidi_read._ENV_VAR, raising=False) + 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) + + client_info = ClientInfo(user_agent="custom-agent/1", gapic_version="9.9.9") + async with _FakeStorageServer([first, second]) as server: + transport = _transport(server.target, client_info) + client = _storage_v2.StorageAsyncClient(transport=transport) + stream = _AsyncReadObjectStream( + client=client, + bucket_name="b", + object_name="o", + client_info=client_info, + ) + 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" + ) + # The header the caller's ClientInfo produces, on both parser paths. + assert "gapic/9.9.9" in server.invocation_metadata["x-goog-api-client"] + assert "custom-agent/1" in server.invocation_metadata["x-goog-api-client"] + + +@pytest.mark.asyncio +async def test_read_object_stream_with_base_client_info( + monkeypatch, keep_google_logger_propagation +): + # AsyncGrpcClient accepts a google.api_core.client_info.ClientInfo, which + # (unlike the gapic_v1 subclass) cannot produce the x-goog-api-client + # header; the stream must still open and send the transport's default. + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) + monkeypatch.delenv(_zero_copy_bidi_read._ENV_VAR, raising=False) + base_client_info = api_core_client_info.ClientInfo(user_agent="gcsfs/1") + first = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + metadata=_storage_v2.Object(name="o", bucket="projects/_/buckets/b") + ) + ) + + async with _FakeStorageServer([first, _data_response()]) as server: + transport = _transport(server.target) + client = _storage_v2.StorageAsyncClient(transport=transport) + stream = _AsyncReadObjectStream( + client=client, + bucket_name="b", + object_name="o", + client_info=base_client_info, + ) + try: + await stream.open() + await stream.send( + _storage_v2.BidiReadObjectRequest( + read_ranges=[_storage_v2.ReadRange(read_length=len(PAYLOAD))] + ) + ) + response = await stream.recv() + await stream.close() + finally: + await transport.close() + + [range_data] = response.object_data_ranges + assert isinstance(range_data.checksummed_data.content, memoryview) + header = server.invocation_metadata["x-goog-api-client"] + # Token comparison: AsyncGrpcClient mutates DEFAULT_CLIENT_INFO.user_agent + # in place in other test modules, and gRPC collapses the extra spaces. + assert header.split() == DEFAULT_CLIENT_INFO.to_user_agent().split() + assert "gcsfs/1" not in header