persist custom metadata on append and download_ranges calls for resume

Author: ankitaluthra1

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:

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:

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)

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)