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..8924b02f1125 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 @@ -19,8 +19,6 @@ from google.api_core import exceptions from google.api_core.retry_async import AsyncRetry -from google.rpc import status_pb2 - from google.cloud import _storage_v2 from google.cloud._storage_v2.types import BidiWriteObjectRedirectedError from google.cloud._storage_v2.types.storage import BidiWriteObjectRequest @@ -41,6 +39,7 @@ _WriteResumptionStrategy, _WriteState, ) +from google.rpc import status_pb2 from . import _utils @@ -299,21 +298,11 @@ def _on_open_error(self, exc): if redirect_proto.generation: self.generation = redirect_proto.generation - async def open( - self, - retry_policy: Optional[AsyncRetry] = None, - metadata: Optional[List[Tuple[str, str]]] = None, - ) -> None: - """Opens the underlying bidi-gRPC stream. - - :raises ValueError: If the stream is already open. - - """ - if self._is_stream_open: - raise ValueError("Underlying bidi-gRPC stream is already open") - + def _merge_retry_policy( + self, retry_policy: Optional[AsyncRetry] = None + ) -> AsyncRetry: if retry_policy is None: - retry_policy = AsyncRetry( + return AsyncRetry( predicate=_is_write_retryable, on_error=self._on_open_error ) else: @@ -324,7 +313,7 @@ def combined_on_error(exc): if original_on_error: original_on_error(exc) - retry_policy = AsyncRetry( + return AsyncRetry( predicate=_is_write_retryable, initial=retry_policy._initial, maximum=retry_policy._maximum, @@ -333,6 +322,21 @@ def combined_on_error(exc): on_error=combined_on_error, ) + async def open( + self, + retry_policy: Optional[AsyncRetry] = None, + metadata: Optional[List[Tuple[str, str]]] = None, + ) -> None: + """Opens the underlying bidi-gRPC stream. + + :raises ValueError: If the stream is already open. + + """ + if self._is_stream_open: + raise ValueError("Underlying bidi-gRPC stream is already open") + + retry_policy = self._merge_retry_policy(retry_policy) + async def _do_open(): current_metadata = list(metadata) if metadata else [] @@ -560,6 +564,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 +586,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 +612,47 @@ 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() + retry_policy = self._merge_retry_policy(retry_policy) - self._is_stream_open = False - return self.persisted_size + attempt_count = 0 + + async def _do_close(): + nonlocal attempt_count + attempt_count += 1 + + if attempt_count > 1: + logger.info( + f"Re-opening the stream for close retry attempt: {attempt_count}" + ) + expected_offset = self.offset + self._is_stream_open = False + await self.open() + if ( + self.offset is not None + and expected_offset is not None + and self.offset < expected_offset + ): + raise exceptions.InternalServerError( + f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + ) + + await self.write_obj_stream.close() + return self.persisted_size + + try: + return await retry_policy(_do_close)() + finally: + self._is_stream_open = False 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 +678,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 +709,46 @@ async def finalize( ), ) - try: + retry_policy = self._merge_retry_policy(retry_policy) + + 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}" + ) + expected_offset = self.offset + self._is_stream_open = False + await self.open() + if ( + self.offset is not None + and expected_offset is not None + and self.offset < expected_offset + ): + raise exceptions.InternalServerError( + f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}" + ) + 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: + try: + await self.write_obj_stream.close() + except Exception as e: + logger.debug( + f"Stream close during finalize cleanup resulted in: {e}" + ) self._is_stream_open = False self.offset = None diff --git a/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py b/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py index b98d8e221048..3b9f5f136493 100644 --- a/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py +++ b/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py @@ -6,10 +6,10 @@ import grpc import pytest import requests + from google.api_core import client_options, exceptions from google.api_core.retry_async import AsyncRetry from google.auth import credentials as auth_credentials - from google.cloud import _storage_v2 as storage_v2 from google.cloud.storage.asyncio.async_appendable_object_writer import ( AsyncAppendableObjectWriter, @@ -136,7 +136,8 @@ def on_retry_error(exc): CONTENT, metadata=fault_injection_metadata, retry_policy=policy_to_pass ) # await writer.finalize() - await writer.close(finalize_on_close=True) + f_o_c = scenario.get("finalize_on_close", True) + await writer.close(finalize_on_close=f_o_c, retry_policy=policy_to_pass) # If an exception was expected, this line should not be reached. if scenario["expected_error"] is not None: @@ -144,16 +145,18 @@ def on_retry_error(exc): f"Expected exception {scenario['expected_error']} was not raised." ) - # 4. Verify the object content. - read_request = storage_v2.ReadObjectRequest( - bucket=f"projects/_/buckets/{bucket_name}", - object=object_name, - ) - read_stream = await gapic_client.read_object(request=read_request) - data = b"" - async for chunk in read_stream: - data += chunk.checksummed_data.content - assert data == CONTENT + # 4. Verify the object content if applicable. + if not scenario.get("skip_verification"): + read_request = storage_v2.ReadObjectRequest( + bucket=f"projects/_/buckets/{bucket_name}", + object=object_name, + ) + read_stream = await gapic_client.read_object(request=read_request) + data = b"" + async for chunk in read_stream: + data += chunk.checksummed_data.content + assert data == CONTENT + if scenario["expected_error"] is None: # Scenarios like 503, 500, smarter resumption, and redirects # SHOULD trigger at least one retry attempt. @@ -235,6 +238,26 @@ async def test_bidi_writes(testbench): "instruction": "redirect-send-handle-and-token-tokenval", "expected_error": None, }, + { + "name": "Retry exactly on finalize/close (Redirect Error)", + "method": "storage.objects.insert", + "instruction": "redirect-send-handle-and-token-mytoken-on-finish-write", + "expected_error": None, + }, + { + "name": "Retry exactly on finalize/close (503)", + "method": "storage.objects.insert", + "instruction": "return-503-on-finish-write", + "expected_error": None, + }, + { + "name": "Retry exactly on close (finalize_on_close=False) (503)", + "method": "storage.objects.insert", + "instruction": "return-503-on-half-close", + "expected_error": None, + "finalize_on_close": False, + "skip_verification": True, + }, ] try: 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..b60801675d7f 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 @@ -17,9 +17,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from google.api_core import exceptions -from google.rpc import status_pb2 +from google.api_core import exceptions from google.cloud._storage_v2.types import storage as storage_type from google.cloud._storage_v2.types.storage import BidiWriteObjectRedirectedError from google.cloud.storage import Blob @@ -29,6 +28,7 @@ AsyncAppendableObjectWriter, _is_write_retryable, ) +from google.rpc import status_pb2 # Constants BUCKET = "test-bucket" @@ -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,134 @@ 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 + + @pytest.mark.asyncio + async def test_finalize_retry_on_redirect_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"] + + redirect = BidiWriteObjectRedirectedError( + routing_token="rt1", + write_handle=storage_type.BidiWriteHandle(handle=b"h1"), + ) + exc = exceptions.Aborted("aborted", errors=[redirect]) + + resource = storage_type.Object(size=999) + mock_appendable_writer["mock_stream"].recv.side_effect = [ + exc, + storage_type.BidiWriteObjectResponse(resource=resource), + ] + + writer.open = mock.AsyncMock() + + res = await writer.finalize() + + assert res == resource + assert writer.persisted_size == 999 + assert mock_appendable_writer["mock_stream"].send.await_count == 2 + assert writer._routing_token == "rt1" + assert writer.write_handle.handle == b"h1" + writer.open.assert_awaited_once() + + @pytest.mark.asyncio + async def test_close_retry_on_redirect_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"] + + redirect = BidiWriteObjectRedirectedError( + routing_token="rt2", + write_handle=storage_type.BidiWriteHandle(handle=b"h2"), + ) + exc = exceptions.Aborted("aborted", errors=[redirect]) + + mock_appendable_writer["mock_stream"].close.side_effect = [ + exc, + None, + ] + + writer.open = mock.AsyncMock() + writer.persisted_size = 999 + + res = await writer.close() + + assert res == 999 + assert mock_appendable_writer["mock_stream"].close.await_count == 2 + assert writer._routing_token == "rt2" + assert writer.write_handle.handle == b"h2" + writer.open.assert_awaited_once()