From 12649ffeca8952bfbfd15430ae97bf09d6539d23 Mon Sep 17 00:00:00 2001 From: Ankita Luthra Date: Mon, 6 Jul 2026 18:05:11 +0000 Subject: [PATCH 1/4] exposes metadata parameter to be added in grpc calls to pass additional client info --- .../asyncio/async_appendable_object_writer.py | 8 +- .../storage/asyncio/async_grpc_client.py | 20 +++- .../asyncio/async_multi_range_downloader.py | 8 +- .../test_async_appendable_object_writer.py | 43 +++++++++ .../unit/asyncio/test_async_grpc_client.py | 22 +++++ .../test_async_multi_range_downloader.py | 92 +++++++++++++++++++ 6 files changed, 189 insertions(+), 4 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 1e31b15faf0c..0b3165be94a9 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -213,6 +213,7 @@ def __init__( self.object_resource: Optional[_storage_v2.Object] = None self._flush_count = 0 self.blob: Optional[Blob] = None + self.metadata: Optional[List[Tuple[str, str]]] = None @classmethod def from_blob( @@ -312,6 +313,8 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") + self.metadata = metadata + if retry_policy is None: retry_policy = AsyncRetry( predicate=_is_write_retryable, on_error=self._on_open_error @@ -334,7 +337,7 @@ def combined_on_error(exc): ) async def _do_open(): - current_metadata = list(metadata) if metadata else [] + current_metadata = list(self.metadata) if self.metadata else [] # Cleanup stream from previous failed attempt, if any. if self.write_obj_stream: @@ -408,6 +411,9 @@ async def append( :raises ValueError: If the stream is not open. """ + if metadata is None: + metadata = self.metadata + if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before append().") if not data: 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 d48fa07c00d9..3a024544a1ae 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 @@ -154,6 +154,9 @@ async def delete_object( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + metadata=(), + timeout=None, + retry=None, **kwargs, ): """Deletes an object and its metadata. @@ -195,7 +198,12 @@ async def delete_object( if_metageneration_not_match=if_metageneration_not_match, **kwargs, ) - await self._grpc_client.delete_object(request=request) + await self._grpc_client.delete_object( + request=request, + metadata=metadata, + timeout=timeout, + retry=retry, + ) async def get_object( self, @@ -207,6 +215,9 @@ async def get_object( if_metageneration_match=None, if_metageneration_not_match=None, soft_deleted=None, + metadata=(), + timeout=None, + retry=None, **kwargs, ): """Retrieves an object's metadata. @@ -258,4 +269,9 @@ async def get_object( ) # Calls the underlying GAPIC StorageAsyncClient.get_object method - return await self._grpc_client.get_object(request=request) + return await self._grpc_client.get_object( + request=request, + metadata=metadata, + timeout=timeout, + retry=retry, + ) 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 6d3f5e2fab4b..fcca8488688b 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 @@ -233,6 +233,7 @@ def __init__( self._open_retries: int = 0 self.is_finalized: bool = False self.full_obj_server_crc32c: Optional[int] = None + self.metadata: Optional[List[Tuple[str, str]]] = None async def __aenter__(self): """Opens the underlying bidi-gRPC connection to read from the object.""" @@ -262,6 +263,8 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") + self.metadata = metadata + if retry_policy is None: def on_error_wrapper(exc): @@ -290,7 +293,7 @@ def combined_on_error(exc): ) async def _do_open(): - current_metadata = list(metadata) if metadata else [] + current_metadata = list(self.metadata) if self.metadata else [] # Cleanup stream from previous failed attempt, if any. if self.read_obj_str: @@ -413,6 +416,9 @@ async def download_ranges( """ + if metadata is None: + metadata = self.metadata + if len(read_ranges) > 1000: raise ValueError( "Invalid input - length of read_ranges cannot be more than 1000" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index c0e62eb54242..abad0a2c203d 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -503,3 +503,46 @@ async def test_methods_require_open_stream_raises(self, mock_appendable_writer): for coro in methods: with pytest.raises(ValueError, match="Stream is not open"): await coro + + @pytest.mark.asyncio + async def test_append_persists_metadata_on_resumption(self, mock_appendable_writer): + # Arrange + mock_client = mock_appendable_writer["mock_client"] + mock_stream = mock_appendable_writer["mock_stream"] + + test_metadata = [("custom-key", "custom-value")] + writer = self._make_one(mock_client) + + # Act - Open with metadata + await writer.open(metadata=test_metadata) + + # Assert first open used metadata + mock_stream.open.assert_called_once_with(metadata=test_metadata) + assert writer.metadata == test_metadata + + # Setup resumption trigger + retryable_exc = exceptions.ServiceUnavailable("Retry me") + mock_stream.send.side_effect = retryable_exc + + # Reset mock_stream.open call count to verify it is called again + mock_stream.open.reset_mock() + + # Setup a fast retry policy to fail quickly in test + from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( + predicate=lambda e: True, + initial=0.01, + maximum=0.01, + multiplier=1.0, + deadline=0.1 + ) + + # Act - append (should trigger retry and use stored metadata) + from google.api_core.exceptions import RetryError + try: + await writer.append(b"data", retry_policy=fast_retry) + except RetryError: + pass + + # Assert second open (during retry) used the same test_metadata + mock_stream.open.assert_called_with(metadata=test_metadata) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py index cea0c785788f..6c132efcb90b 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py @@ -285,6 +285,11 @@ async def test_delete_object(self, mock_async_storage_client): if_metageneration_match = 111 if_metageneration_not_match = 222 + # New parameters + metadata = (("x-goog-api-client", "test-ua"),) + timeout = 10.0 + retry = mock.Mock() + # Act await client.delete_object( bucket_name, @@ -294,6 +299,9 @@ async def test_delete_object(self, mock_async_storage_client): if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + metadata=metadata, + timeout=timeout, + retry=retry, ) # Assert @@ -306,6 +314,9 @@ async def test_delete_object(self, mock_async_storage_client): assert request.if_generation_not_match == if_generation_not_match assert request.if_metageneration_match == if_metageneration_match assert request.if_metageneration_not_match == if_metageneration_not_match + assert call_kwargs["metadata"] == metadata + assert call_kwargs["timeout"] == timeout + assert call_kwargs["retry"] == retry @mock.patch("google.cloud._storage_v2.StorageAsyncClient") @pytest.mark.asyncio @@ -354,6 +365,11 @@ async def test_get_object_with_all_parameters(self, mock_async_storage_client): if_metageneration_not_match = 222 soft_deleted = True + # New parameters + metadata = (("x-goog-api-client", "test-ua"),) + timeout = 10.0 + retry = mock.Mock() + # Act await client.get_object( bucket_name, @@ -364,6 +380,9 @@ async def test_get_object_with_all_parameters(self, mock_async_storage_client): if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, soft_deleted=soft_deleted, + metadata=metadata, + timeout=timeout, + retry=retry, ) # Assert @@ -377,3 +396,6 @@ async def test_get_object_with_all_parameters(self, mock_async_storage_client): assert request.if_metageneration_match == if_metageneration_match assert request.if_metageneration_not_match == if_metageneration_not_match assert request.soft_deleted is True + assert call_kwargs["metadata"] == metadata + assert call_kwargs["timeout"] == timeout + assert call_kwargs["retry"] == retry 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 6ead8d8964e9..c3c44bf657a6 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 @@ -707,3 +707,95 @@ async def test_download_ranges_closes_on_datacorruption( await mrd.download_ranges([(0, 0, BytesIO())]) mrd.close.assert_called_once() + + @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_persists_metadata_on_resumption( + self, mock_cls_async_read_object_stream, mock_random_int + ): + # Arrange + mock_client = mock.MagicMock() + mock_client.grpc_client = mock.AsyncMock() + + mock_stream_1 = mock.MagicMock() + mock_stream_1.open = AsyncMock() + mock_stream_1.generation_number = _TEST_GENERATION_NUMBER + mock_stream_1.persisted_size = _TEST_OBJECT_SIZE + mock_stream_1.read_handle = _TEST_READ_HANDLE + mock_stream_1.is_stream_open = True + + mock_stream_2 = mock.MagicMock() + mock_stream_2.open = AsyncMock() + mock_stream_2.generation_number = _TEST_GENERATION_NUMBER + mock_stream_2.persisted_size = _TEST_OBJECT_SIZE + mock_stream_2.read_handle = _TEST_READ_HANDLE + mock_stream_2.is_stream_open = True + + mock_cls_async_read_object_stream.side_effect = [mock_stream_1, mock_stream_2] + + test_metadata = [("custom-key", "custom-value")] + + # Act - Create MRD with metadata + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, + metadata=test_metadata + ) + + # Assert first open used metadata + mock_stream_1.open.assert_called_once_with(metadata=test_metadata) + assert mrd.metadata == test_metadata + + # Setup resumption trigger + from google.api_core import exceptions as core_exceptions + retryable_exc = core_exceptions.ServiceUnavailable("Retry me") + + # Mock multiplexer to trigger retry + mrd._multiplexer = mock.MagicMock() + mrd._multiplexer.stream_generation = 1 + + # Configure reopen_stream to execute the factory + async def fake_reopen_stream(broken_gen, stream_factory): + await stream_factory() + mrd._multiplexer.reopen_stream = AsyncMock(side_effect=fake_reopen_stream) + + mrd._multiplexer.send = AsyncMock(side_effect=retryable_exc) + + # We need mock_multiplexer.register to return a queue + mock_queue = mock.AsyncMock() + mrd._multiplexer.register.return_value = mock_queue + + # We need mock_queue.get to return a StreamError to propagate exception to retry manager + from google.cloud.storage.asyncio._stream_multiplexer import _StreamError + mock_queue.get.return_value = _StreamError(retryable_exc, generation=1) + + mock_random_int.return_value = 123 + + # Setup a fast retry policy to fail quickly in test + from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( + predicate=lambda e: True, + initial=0.01, + maximum=0.01, + multiplier=1.0, + deadline=0.1 + ) + + # Act - download ranges (should trigger retry and use stored metadata) + buffer = BytesIO() + from google.api_core.exceptions import RetryError + try: + await mrd.download_ranges([(0, 10, buffer)], retry_policy=fast_retry) + except RetryError: + pass + + # Assert that reopen_stream was called. + mrd._multiplexer.reopen_stream.assert_called() + + # Verify mock_stream_2 was opened with test_metadata + # (It should have been opened during the retry attempts) + mock_stream_2.open.assert_called_with(metadata=test_metadata) From b1c2b3b6d0c16b65575a771789cdd17383948a43 Mon Sep 17 00:00:00 2001 From: Ankita Luthra Date: Mon, 6 Jul 2026 18:20:56 +0000 Subject: [PATCH 2/4] lint fix --- .../asyncio/test_async_appendable_object_writer.py | 4 +++- .../unit/asyncio/test_async_multi_range_downloader.py | 10 +++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index abad0a2c203d..335822651e4b 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -529,16 +529,18 @@ async def test_append_persists_metadata_on_resumption(self, mock_appendable_writ # Setup a fast retry policy to fail quickly in test from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( predicate=lambda e: True, initial=0.01, maximum=0.01, multiplier=1.0, - deadline=0.1 + deadline=0.1, ) # Act - append (should trigger retry and use stored metadata) from google.api_core.exceptions import RetryError + try: await writer.append(b"data", retry_policy=fast_retry) except RetryError: 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 c3c44bf657a6..d102d6e09275 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 @@ -742,8 +742,7 @@ async def test_download_ranges_persists_metadata_on_resumption( # Act - Create MRD with metadata mrd = await AsyncMultiRangeDownloader.create_mrd( - mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, - metadata=test_metadata + mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, metadata=test_metadata ) # Assert first open used metadata @@ -752,6 +751,7 @@ async def test_download_ranges_persists_metadata_on_resumption( # Setup resumption trigger from google.api_core import exceptions as core_exceptions + retryable_exc = core_exceptions.ServiceUnavailable("Retry me") # Mock multiplexer to trigger retry @@ -761,6 +761,7 @@ async def test_download_ranges_persists_metadata_on_resumption( # Configure reopen_stream to execute the factory async def fake_reopen_stream(broken_gen, stream_factory): await stream_factory() + mrd._multiplexer.reopen_stream = AsyncMock(side_effect=fake_reopen_stream) mrd._multiplexer.send = AsyncMock(side_effect=retryable_exc) @@ -771,23 +772,26 @@ async def fake_reopen_stream(broken_gen, stream_factory): # We need mock_queue.get to return a StreamError to propagate exception to retry manager from google.cloud.storage.asyncio._stream_multiplexer import _StreamError + mock_queue.get.return_value = _StreamError(retryable_exc, generation=1) mock_random_int.return_value = 123 # Setup a fast retry policy to fail quickly in test from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( predicate=lambda e: True, initial=0.01, maximum=0.01, multiplier=1.0, - deadline=0.1 + deadline=0.1, ) # Act - download ranges (should trigger retry and use stored metadata) buffer = BytesIO() from google.api_core.exceptions import RetryError + try: await mrd.download_ranges([(0, 10, buffer)], retry_policy=fast_retry) except RetryError: From 10af2a572d11df5730e88e02e877aa5754755527 Mon Sep 17 00:00:00 2001 From: Ankita Luthra Date: Tue, 7 Jul 2026 07:19:45 +0000 Subject: [PATCH 3/4] persist custom metadata on append and download_ranges calls for resume --- .../asyncio/async_appendable_object_writer.py | 4 +- .../asyncio/async_multi_range_downloader.py | 4 +- .../test_async_appendable_object_writer.py | 42 ++++++++++ .../test_async_multi_range_downloader.py | 79 +++++++++++++++++++ 4 files changed, 127 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 0b3165be94a9..0eb18df61896 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -411,7 +411,9 @@ async def append( :raises ValueError: If the stream is not open. """ - if metadata is None: + if metadata is not None: + self.metadata = metadata + else: metadata = self.metadata if not self._is_stream_open: 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 fcca8488688b..df19f482889b 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 @@ -416,7 +416,9 @@ async def download_ranges( """ - if metadata is None: + if metadata is not None: + self.metadata = metadata + else: metadata = self.metadata if len(read_ranges) > 1000: diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 335822651e4b..5c1dde670219 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -548,3 +548,45 @@ async def test_append_persists_metadata_on_resumption(self, mock_appendable_writ # Assert second open (during retry) used the same test_metadata mock_stream.open.assert_called_with(metadata=test_metadata) + + @pytest.mark.asyncio + async def test_append_updates_and_persists_metadata_on_resumption( + self, mock_appendable_writer + ): + # Arrange + mock_client = mock_appendable_writer["mock_client"] + mock_stream = mock_appendable_writer["mock_stream"] + + initial_metadata = [("initial-key", "initial-value")] + updated_metadata = [("updated-key", "updated-value")] + writer = self._make_one(mock_client) + + # Act - Open with initial metadata + await writer.open(metadata=initial_metadata) + assert writer.metadata == initial_metadata + + # Setup resumption trigger when append is called + retryable_exc = exceptions.ServiceUnavailable("Retry me") + mock_stream.send.side_effect = retryable_exc + mock_stream.open.reset_mock() + + from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( + predicate=lambda e: True, + initial=0.01, + maximum=0.01, + multiplier=1.0, + deadline=0.1, + ) + + from google.api_core.exceptions import RetryError + try: + await writer.append( + b"data", retry_policy=fast_retry, metadata=updated_metadata + ) + except RetryError: + pass + + # Assert writer.metadata was updated and used during stream reopening + assert writer.metadata == updated_metadata + mock_stream.open.assert_called_with(metadata=updated_metadata) 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 d102d6e09275..4537846f2bf4 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 @@ -803,3 +803,82 @@ async def fake_reopen_stream(broken_gen, stream_factory): # Verify mock_stream_2 was opened with test_metadata # (It should have been opened during the retry attempts) mock_stream_2.open.assert_called_with(metadata=test_metadata) + + @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_updates_and_persists_metadata_on_resumption( + self, mock_cls_async_read_object_stream, mock_random_int + ): + # Arrange + mock_client = mock.MagicMock() + mock_client.grpc_client = mock.AsyncMock() + + mock_stream_1 = mock.MagicMock() + mock_stream_1.open = AsyncMock() + mock_stream_1.generation_number = _TEST_GENERATION_NUMBER + mock_stream_1.persisted_size = _TEST_OBJECT_SIZE + mock_stream_1.read_handle = _TEST_READ_HANDLE + mock_stream_1.is_stream_open = True + + mock_stream_2 = mock.MagicMock() + mock_stream_2.open = AsyncMock() + mock_stream_2.generation_number = _TEST_GENERATION_NUMBER + mock_stream_2.persisted_size = _TEST_OBJECT_SIZE + mock_stream_2.read_handle = _TEST_READ_HANDLE + mock_stream_2.is_stream_open = True + + mock_cls_async_read_object_stream.side_effect = [mock_stream_1, mock_stream_2] + + initial_metadata = [("initial-key", "initial-value")] + updated_metadata = [("updated-key", "updated-value")] + + # Act - Create MRD with initial metadata + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, metadata=initial_metadata + ) + assert mrd.metadata == initial_metadata + + from google.api_core import exceptions as core_exceptions + retryable_exc = core_exceptions.ServiceUnavailable("Retry me") + + mrd._multiplexer = mock.MagicMock() + mrd._multiplexer.stream_generation = 1 + + async def fake_reopen_stream(broken_gen, stream_factory): + await stream_factory() + + mrd._multiplexer.reopen_stream = AsyncMock(side_effect=fake_reopen_stream) + mrd._multiplexer.send = AsyncMock(side_effect=retryable_exc) + + mock_queue = mock.AsyncMock() + mrd._multiplexer.register.return_value = mock_queue + from google.cloud.storage.asyncio._stream_multiplexer import _StreamError + mock_queue.get.return_value = _StreamError(retryable_exc, generation=1) + + mock_random_int.return_value = 123 + + from google.api_core.retry_async import AsyncRetry + fast_retry = AsyncRetry( + predicate=lambda e: True, + initial=0.01, + maximum=0.01, + multiplier=1.0, + deadline=0.1, + ) + + buffer = BytesIO() + from google.api_core.exceptions import RetryError + try: + await mrd.download_ranges( + [(0, 10, buffer)], retry_policy=fast_retry, metadata=updated_metadata + ) + except RetryError: + pass + + assert mrd.metadata == updated_metadata + mock_stream_2.open.assert_called_with(metadata=updated_metadata) From ef11aeb15c4140ece8086e2980c9756877242e3e Mon Sep 17 00:00:00 2001 From: Ankita Luthra Date: Tue, 7 Jul 2026 15:47:38 +0000 Subject: [PATCH 4/4] removes metadata from reader and writer | updates docstring --- .../asyncio/async_appendable_object_writer.py | 10 +- .../storage/asyncio/async_grpc_client.py | 38 ++++ .../asyncio/async_multi_range_downloader.py | 10 +- .../test_async_appendable_object_writer.py | 87 --------- .../unit/asyncio/test_async_grpc_client.py | 32 ++++ .../test_async_multi_range_downloader.py | 175 ------------------ 6 files changed, 72 insertions(+), 280 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 0eb18df61896..1e31b15faf0c 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -213,7 +213,6 @@ def __init__( self.object_resource: Optional[_storage_v2.Object] = None self._flush_count = 0 self.blob: Optional[Blob] = None - self.metadata: Optional[List[Tuple[str, str]]] = None @classmethod def from_blob( @@ -313,8 +312,6 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") - self.metadata = metadata - if retry_policy is None: retry_policy = AsyncRetry( predicate=_is_write_retryable, on_error=self._on_open_error @@ -337,7 +334,7 @@ def combined_on_error(exc): ) async def _do_open(): - current_metadata = list(self.metadata) if self.metadata else [] + current_metadata = list(metadata) if metadata else [] # Cleanup stream from previous failed attempt, if any. if self.write_obj_stream: @@ -411,11 +408,6 @@ async def append( :raises ValueError: If the stream is not open. """ - if metadata is not None: - self.metadata = metadata - else: - metadata = self.metadata - if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before append().") if not data: 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 3a024544a1ae..717a8724de47 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 @@ -26,6 +26,22 @@ _DEFAULT_HOST = "storage.googleapis.com" +def _validate_metadata(metadata): + """Validates that metadata is a sequence of (key, value) pairs.""" + if metadata is None: + return + if not isinstance(metadata, (list, tuple)): + raise TypeError("metadata must be a list or tuple of (key, value) pairs.") + for item in metadata: + if not isinstance(item, (list, tuple)) or len(item) != 2: + raise ValueError( + "Each element in metadata must be a list or tuple of exactly 2 strings (key, value)." + ) + key, value = item + if not isinstance(key, str) or not isinstance(value, str): + raise TypeError("Both key and value in metadata pairs must be strings.") + + class AsyncGrpcClient: """An asynchronous client for interacting with Google Cloud Storage using the gRPC API. @@ -184,8 +200,18 @@ async def delete_object( :type if_metageneration_not_match: int :param if_metageneration_not_match: (Optional) + :type metadata: Sequence[Tuple[str, str]] + :param metadata: (Optional) Additional metadata that is provided to the method. + + :type timeout: float or None + :param timeout: + (Optional) The amount of time, in seconds, to wait for the request to + complete. + :type retry: :class:`~google.api_core.retry_async.AsyncRetry` or :class:`~google.api_core.retry.Retry` or None + :param retry: (Optional) Designation of what errors, if any, should be retried. """ + _validate_metadata(metadata) # The gRPC API requires the bucket name to be in the format "projects/_/buckets/bucket_name" bucket_path = f"projects/_/buckets/{bucket_name}" request = storage_v2.DeleteObjectRequest( @@ -251,9 +277,21 @@ async def get_object( :param soft_deleted: (Optional) If True, return the soft-deleted version of this object. + :type metadata: Sequence[Tuple[str, str]] + :param metadata: (Optional) Additional metadata that is provided to the method. + + :type timeout: float or None + :param timeout: + (Optional) The amount of time, in seconds, to wait for the request to + complete. + + :type retry: :class:`~google.api_core.retry_async.AsyncRetry` or :class:`~google.api_core.retry.Retry` or None + :param retry: (Optional) Designation of what errors, if any, should be retried. + :rtype: :class:`google.cloud._storage_v2.types.Object` :returns: The object metadata resource. """ + _validate_metadata(metadata) bucket_path = f"projects/_/buckets/{bucket_name}" request = storage_v2.GetObjectRequest( 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 df19f482889b..6d3f5e2fab4b 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 @@ -233,7 +233,6 @@ def __init__( self._open_retries: int = 0 self.is_finalized: bool = False self.full_obj_server_crc32c: Optional[int] = None - self.metadata: Optional[List[Tuple[str, str]]] = None async def __aenter__(self): """Opens the underlying bidi-gRPC connection to read from the object.""" @@ -263,8 +262,6 @@ async def open( if self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is already open") - self.metadata = metadata - if retry_policy is None: def on_error_wrapper(exc): @@ -293,7 +290,7 @@ def combined_on_error(exc): ) async def _do_open(): - current_metadata = list(self.metadata) if self.metadata else [] + current_metadata = list(metadata) if metadata else [] # Cleanup stream from previous failed attempt, if any. if self.read_obj_str: @@ -416,11 +413,6 @@ async def download_ranges( """ - if metadata is not None: - self.metadata = metadata - else: - metadata = self.metadata - if len(read_ranges) > 1000: raise ValueError( "Invalid input - length of read_ranges cannot be more than 1000" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index 5c1dde670219..c0e62eb54242 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -503,90 +503,3 @@ async def test_methods_require_open_stream_raises(self, mock_appendable_writer): for coro in methods: with pytest.raises(ValueError, match="Stream is not open"): await coro - - @pytest.mark.asyncio - async def test_append_persists_metadata_on_resumption(self, mock_appendable_writer): - # Arrange - mock_client = mock_appendable_writer["mock_client"] - mock_stream = mock_appendable_writer["mock_stream"] - - test_metadata = [("custom-key", "custom-value")] - writer = self._make_one(mock_client) - - # Act - Open with metadata - await writer.open(metadata=test_metadata) - - # Assert first open used metadata - mock_stream.open.assert_called_once_with(metadata=test_metadata) - assert writer.metadata == test_metadata - - # Setup resumption trigger - retryable_exc = exceptions.ServiceUnavailable("Retry me") - mock_stream.send.side_effect = retryable_exc - - # Reset mock_stream.open call count to verify it is called again - mock_stream.open.reset_mock() - - # Setup a fast retry policy to fail quickly in test - from google.api_core.retry_async import AsyncRetry - - fast_retry = AsyncRetry( - predicate=lambda e: True, - initial=0.01, - maximum=0.01, - multiplier=1.0, - deadline=0.1, - ) - - # Act - append (should trigger retry and use stored metadata) - from google.api_core.exceptions import RetryError - - try: - await writer.append(b"data", retry_policy=fast_retry) - except RetryError: - pass - - # Assert second open (during retry) used the same test_metadata - mock_stream.open.assert_called_with(metadata=test_metadata) - - @pytest.mark.asyncio - async def test_append_updates_and_persists_metadata_on_resumption( - self, mock_appendable_writer - ): - # Arrange - mock_client = mock_appendable_writer["mock_client"] - mock_stream = mock_appendable_writer["mock_stream"] - - initial_metadata = [("initial-key", "initial-value")] - updated_metadata = [("updated-key", "updated-value")] - writer = self._make_one(mock_client) - - # Act - Open with initial metadata - await writer.open(metadata=initial_metadata) - assert writer.metadata == initial_metadata - - # Setup resumption trigger when append is called - retryable_exc = exceptions.ServiceUnavailable("Retry me") - mock_stream.send.side_effect = retryable_exc - mock_stream.open.reset_mock() - - from google.api_core.retry_async import AsyncRetry - fast_retry = AsyncRetry( - predicate=lambda e: True, - initial=0.01, - maximum=0.01, - multiplier=1.0, - deadline=0.1, - ) - - from google.api_core.exceptions import RetryError - try: - await writer.append( - b"data", retry_policy=fast_retry, metadata=updated_metadata - ) - except RetryError: - pass - - # Assert writer.metadata was updated and used during stream reopening - assert writer.metadata == updated_metadata - mock_stream.open.assert_called_with(metadata=updated_metadata) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py index 6c132efcb90b..24168dc137d4 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_grpc_client.py @@ -399,3 +399,35 @@ async def test_get_object_with_all_parameters(self, mock_async_storage_client): assert call_kwargs["metadata"] == metadata assert call_kwargs["timeout"] == timeout assert call_kwargs["retry"] == retry + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "invalid_metadata, expected_exc", + [ + ("not-a-sequence", TypeError), + ([("key_only",)], ValueError), + ([("too", "many", "items")], ValueError), + ([(123, "val")], TypeError), + ([("key", 456)], TypeError), + ], + ) + async def test_delete_object_invalid_metadata(self, invalid_metadata, expected_exc): + client = async_grpc_client.AsyncGrpcClient(credentials=_make_credentials()) + with pytest.raises(expected_exc): + await client.delete_object("bucket", "object", metadata=invalid_metadata) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "invalid_metadata, expected_exc", + [ + ("not-a-sequence", TypeError), + ([("key_only",)], ValueError), + ([("too", "many", "items")], ValueError), + ([(123, "val")], TypeError), + ([("key", 456)], TypeError), + ], + ) + async def test_get_object_invalid_metadata(self, invalid_metadata, expected_exc): + client = async_grpc_client.AsyncGrpcClient(credentials=_make_credentials()) + with pytest.raises(expected_exc): + await client.get_object("bucket", "object", metadata=invalid_metadata) 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 4537846f2bf4..6ead8d8964e9 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 @@ -707,178 +707,3 @@ async def test_download_ranges_closes_on_datacorruption( await mrd.download_ranges([(0, 0, BytesIO())]) mrd.close.assert_called_once() - - @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_persists_metadata_on_resumption( - self, mock_cls_async_read_object_stream, mock_random_int - ): - # Arrange - mock_client = mock.MagicMock() - mock_client.grpc_client = mock.AsyncMock() - - mock_stream_1 = mock.MagicMock() - mock_stream_1.open = AsyncMock() - mock_stream_1.generation_number = _TEST_GENERATION_NUMBER - mock_stream_1.persisted_size = _TEST_OBJECT_SIZE - mock_stream_1.read_handle = _TEST_READ_HANDLE - mock_stream_1.is_stream_open = True - - mock_stream_2 = mock.MagicMock() - mock_stream_2.open = AsyncMock() - mock_stream_2.generation_number = _TEST_GENERATION_NUMBER - mock_stream_2.persisted_size = _TEST_OBJECT_SIZE - mock_stream_2.read_handle = _TEST_READ_HANDLE - mock_stream_2.is_stream_open = True - - mock_cls_async_read_object_stream.side_effect = [mock_stream_1, mock_stream_2] - - test_metadata = [("custom-key", "custom-value")] - - # Act - Create MRD with metadata - mrd = await AsyncMultiRangeDownloader.create_mrd( - mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, metadata=test_metadata - ) - - # Assert first open used metadata - mock_stream_1.open.assert_called_once_with(metadata=test_metadata) - assert mrd.metadata == test_metadata - - # Setup resumption trigger - from google.api_core import exceptions as core_exceptions - - retryable_exc = core_exceptions.ServiceUnavailable("Retry me") - - # Mock multiplexer to trigger retry - mrd._multiplexer = mock.MagicMock() - mrd._multiplexer.stream_generation = 1 - - # Configure reopen_stream to execute the factory - async def fake_reopen_stream(broken_gen, stream_factory): - await stream_factory() - - mrd._multiplexer.reopen_stream = AsyncMock(side_effect=fake_reopen_stream) - - mrd._multiplexer.send = AsyncMock(side_effect=retryable_exc) - - # We need mock_multiplexer.register to return a queue - mock_queue = mock.AsyncMock() - mrd._multiplexer.register.return_value = mock_queue - - # We need mock_queue.get to return a StreamError to propagate exception to retry manager - from google.cloud.storage.asyncio._stream_multiplexer import _StreamError - - mock_queue.get.return_value = _StreamError(retryable_exc, generation=1) - - mock_random_int.return_value = 123 - - # Setup a fast retry policy to fail quickly in test - from google.api_core.retry_async import AsyncRetry - - fast_retry = AsyncRetry( - predicate=lambda e: True, - initial=0.01, - maximum=0.01, - multiplier=1.0, - deadline=0.1, - ) - - # Act - download ranges (should trigger retry and use stored metadata) - buffer = BytesIO() - from google.api_core.exceptions import RetryError - - try: - await mrd.download_ranges([(0, 10, buffer)], retry_policy=fast_retry) - except RetryError: - pass - - # Assert that reopen_stream was called. - mrd._multiplexer.reopen_stream.assert_called() - - # Verify mock_stream_2 was opened with test_metadata - # (It should have been opened during the retry attempts) - mock_stream_2.open.assert_called_with(metadata=test_metadata) - - @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_updates_and_persists_metadata_on_resumption( - self, mock_cls_async_read_object_stream, mock_random_int - ): - # Arrange - mock_client = mock.MagicMock() - mock_client.grpc_client = mock.AsyncMock() - - mock_stream_1 = mock.MagicMock() - mock_stream_1.open = AsyncMock() - mock_stream_1.generation_number = _TEST_GENERATION_NUMBER - mock_stream_1.persisted_size = _TEST_OBJECT_SIZE - mock_stream_1.read_handle = _TEST_READ_HANDLE - mock_stream_1.is_stream_open = True - - mock_stream_2 = mock.MagicMock() - mock_stream_2.open = AsyncMock() - mock_stream_2.generation_number = _TEST_GENERATION_NUMBER - mock_stream_2.persisted_size = _TEST_OBJECT_SIZE - mock_stream_2.read_handle = _TEST_READ_HANDLE - mock_stream_2.is_stream_open = True - - mock_cls_async_read_object_stream.side_effect = [mock_stream_1, mock_stream_2] - - initial_metadata = [("initial-key", "initial-value")] - updated_metadata = [("updated-key", "updated-value")] - - # Act - Create MRD with initial metadata - mrd = await AsyncMultiRangeDownloader.create_mrd( - mock_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME, metadata=initial_metadata - ) - assert mrd.metadata == initial_metadata - - from google.api_core import exceptions as core_exceptions - retryable_exc = core_exceptions.ServiceUnavailable("Retry me") - - mrd._multiplexer = mock.MagicMock() - mrd._multiplexer.stream_generation = 1 - - async def fake_reopen_stream(broken_gen, stream_factory): - await stream_factory() - - mrd._multiplexer.reopen_stream = AsyncMock(side_effect=fake_reopen_stream) - mrd._multiplexer.send = AsyncMock(side_effect=retryable_exc) - - mock_queue = mock.AsyncMock() - mrd._multiplexer.register.return_value = mock_queue - from google.cloud.storage.asyncio._stream_multiplexer import _StreamError - mock_queue.get.return_value = _StreamError(retryable_exc, generation=1) - - mock_random_int.return_value = 123 - - from google.api_core.retry_async import AsyncRetry - fast_retry = AsyncRetry( - predicate=lambda e: True, - initial=0.01, - maximum=0.01, - multiplier=1.0, - deadline=0.1, - ) - - buffer = BytesIO() - from google.api_core.exceptions import RetryError - try: - await mrd.download_ranges( - [(0, 10, buffer)], retry_policy=fast_retry, metadata=updated_metadata - ) - except RetryError: - pass - - assert mrd.metadata == updated_metadata - mock_stream_2.open.assert_called_with(metadata=updated_metadata)