fix(storage): add retry support for finalize and close in AsyncAppendableObjectWriter

Fixes: b/532527637

Author: sahusneha2004

packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py

@@ -560,6 +560,7 @@ async def close(
         self,
         finalize_on_close=False,
         full_object_checksum: Optional[int] = None,
+        retry_policy: Optional[AsyncRetry] = None,
     ) -> Union[int, _storage_v2.Object]:
         """Closes the underlying bidi-gRPC stream.
 
@@ -581,6 +582,9 @@ async def close(
                 crc32c_int = google_crc32c.value(data)
                 print(crc32c_int)
 
+        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
+        :param retry_policy: (Optional) The retry policy to use for the operation.
+
         rtype: Union[int, _storage_v2.Object]
         returns: Updated `self.persisted_size` by default after closing the
             bidi-gRPC stream. However, if `finalize_on_close=True` is passed,
@@ -604,15 +608,20 @@ async def close(
             )
 
         if finalize_on_close:
-            return await self.finalize(full_object_checksum=full_object_checksum)
+            return await self.finalize(
+                full_object_checksum=full_object_checksum,
+                retry_policy=retry_policy,
+            )
 
         await self.write_obj_stream.close()
 
         self._is_stream_open = False
         return self.persisted_size
 
     async def finalize(
-        self, full_object_checksum: Optional[int] = None
+        self,
+        full_object_checksum: Optional[int] = None,
+        retry_policy: Optional[AsyncRetry] = None,
     ) -> _storage_v2.Object:
         """Finalizes the Appendable Object.
 
@@ -638,6 +647,9 @@ async def finalize(
                 crc32c_int = google_crc32c.value(data)
                 print(crc32c_int)
 
+        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
+        :param retry_policy: (Optional) The retry policy to use for the operation.
+
         rtype: google.cloud.storage_v2.types.Object
         returns: The finalized object resource.
 
@@ -666,12 +678,18 @@ async def finalize(
                 ),
             )
 
-        try:
+        if retry_policy is None:
+            retry_policy = AsyncRetry(predicate=_is_write_retryable)
+
+        async def _do_finalize():
             await self.write_obj_stream.send(finalize_req)
             response = await self.write_obj_stream.recv()
             self.object_resource = response.resource
             self.persisted_size = self.object_resource.size
             return self.object_resource
+
+        try:
+            return await retry_policy(_do_finalize)()
         finally:
             await self.write_obj_stream.close()
             self._is_stream_open = False

packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -560,7 +560,9 @@ async def test_close_with_checksum_and_finalize(self, mock_appendable_writer):
 
         checksum = 12345678
         await writer.close(finalize_on_close=True, full_object_checksum=checksum)
-        writer.finalize.assert_awaited_once_with(full_object_checksum=checksum)
+        writer.finalize.assert_awaited_once_with(
+            full_object_checksum=checksum, retry_policy=None
+        )
 
     @pytest.mark.asyncio
     async def test_close_with_checksum_without_finalize_raises(
@@ -625,3 +627,76 @@ async def test_finalize_mismatch_closes_stream(self, mock_appendable_writer):
         # Assert stream was closed and local state reset despite exception
         mock_appendable_writer["mock_stream"].close.assert_awaited()
         assert not writer._is_stream_open
+
+    @pytest.mark.asyncio
+    async def test_finalize_retry_on_transient_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.side_effect = [
+            exceptions.InternalServerError("500 Transient Error"),
+            storage_type.BidiWriteObjectResponse(resource=resource),
+        ]
+
+        res = await writer.finalize()
+
+        assert res == resource
+        assert writer.persisted_size == 999
+        assert mock_appendable_writer["mock_stream"].send.await_count == 2
+        assert not writer._is_stream_open
+
+    @pytest.mark.asyncio
+    async def test_finalize_custom_retry_policy(self, mock_appendable_writer):
+        from google.api_core.retry_async import AsyncRetry
+
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        custom_policy = AsyncRetry(predicate=lambda exc: isinstance(exc, exceptions.InternalServerError))
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.return_value = (
+            storage_type.BidiWriteObjectResponse(resource=resource)
+        )
+
+        res = await writer.finalize(retry_policy=custom_policy)
+        assert res == resource
+
+    @pytest.mark.asyncio
+    async def test_close_with_finalize_and_custom_retry_policy(self, mock_appendable_writer):
+        from google.api_core.retry_async import AsyncRetry
+
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.finalize = AsyncMock()
+
+        custom_policy = AsyncRetry(predicate=lambda exc: False)
+        await writer.close(finalize_on_close=True, retry_policy=custom_policy)
+        writer.finalize.assert_awaited_once_with(
+            full_object_checksum=None,
+            retry_policy=custom_policy,
+        )
+
+    @pytest.mark.asyncio
+    async def test_close_retry_on_transient_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.side_effect = [
+            exceptions.InternalServerError("500 Transient Error"),
+            storage_type.BidiWriteObjectResponse(resource=resource),
+        ]
+
+        res = await writer.close(finalize_on_close=True)
+
+        assert res == resource
+        assert writer.persisted_size == 999
+        assert mock_appendable_writer["mock_stream"].send.await_count == 2
+        assert not writer._is_stream_open
+
+
+