From 3366eb0c3e666b053e6198ec56f82a940c97b90f Mon Sep 17 00:00:00 2001 From: Sneha Sahu Date: Thu, 16 Jul 2026 10:19:46 +0000 Subject: [PATCH] fix(storage): add retry support for finalize and close in AsyncAppendableObjectWriter Fixes: b/532527637 --- .../asyncio/async_appendable_object_writer.py | 50 +++++++++++- .../test_async_appendable_object_writer.py | 78 ++++++++++++++++++- 2 files changed, 123 insertions(+), 5 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 ce954ed3d7c2..0bcc6caca56e 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 @@ -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,7 +608,10 @@ 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() @@ -612,7 +619,9 @@ async def close( 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 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 a51b0961eb8d..a06bbcbfd88f 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 @@ -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