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

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,14 +678,44 @@ async def finalize(
                 ),
             )
 
-        try:
+        if retry_policy is None:
+            retry_policy = AsyncRetry(predicate=_is_write_retryable)
+
+        attempt_count = 0
+
+        async def _do_finalize():
+            nonlocal attempt_count
+            attempt_count += 1
+
+            if attempt_count > 1:
+                logger.info(
+                    f"Re-opening the stream for finalize retry attempt: {attempt_count}"
+                )
+                self._is_stream_open = False
+                await self.open()
+
             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()
+            if self.write_obj_stream:
+                if self.write_obj_stream.is_stream_open:
+                    try:
+                        await self.write_obj_stream.close()
+                    except Exception as e:
+                        logger.warning(
+                            f"Error closing write stream during finalize cleanup. Got exception: {e}"
+                        )
+                else:
+                    try:
+                        await self.write_obj_stream.close()
+                    except Exception:
+                        pass
             self._is_stream_open = False
             self.offset = None
 

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,77 @@ 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