exposes metadata parameter to be added in grpc calls to pass additional client info

Author: ankitaluthra1

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:

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,
+        )

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"

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)

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

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)