feat(experimental): add write resumption strategy (#1663)

Adding writes resumption strategy which will be used for error handling
of bidi writes operation.

Author: Pulkit0110

google/cloud/storage/_experimental/asyncio/retry/writes_resumption_strategy.py

@@ -0,0 +1,146 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+fromtypingimportAny, Dict, IO, Iterable, Optional, Union
+
+importgoogle_crc32c
+fromgoogle.cloud._storage_v2.typesimportstorageasstorage_type
+fromgoogle.cloud._storage_v2.types.storageimportBidiWriteObjectRedirectedError
+fromgoogle.cloud.storage._experimental.asyncio.retry.base_strategyimport (
+_BaseResumptionStrategy,
+)
+
+
+class_WriteState:
+"""A helper class to track the state of a single upload operation.
+
+ :type spec: :class:`google.cloud.storage_v2.types.AppendObjectSpec`
+ :param spec: The specification for the object to write.
+
+ :type chunk_size: int
+ :param chunk_size: The size of chunks to write to the server.
+
+ :type user_buffer: IO[bytes]
+ :param user_buffer: The data source.
+ """
+
+def__init__(
+self,
+spec: Union[storage_type.AppendObjectSpec, storage_type.WriteObjectSpec],
+chunk_size: int,
+user_buffer: IO[bytes],
+ ):
+self.spec=spec
+self.chunk_size=chunk_size
+self.user_buffer=user_buffer
+self.persisted_size: int=0
+self.bytes_sent: int=0
+self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] =None
+self.routing_token: Optional[str] =None
+self.is_finalized: bool=False
+
+
+class_WriteResumptionStrategy(_BaseResumptionStrategy):
+"""The concrete resumption strategy for bidi writes."""
+
+defgenerate_requests(
+self, state: Dict[str, Any]
+ ) ->Iterable[storage_type.BidiWriteObjectRequest]:
+"""Generates BidiWriteObjectRequests to resume or continue the upload.
+
+ For Appendable Objects, every stream opening should send an
+ AppendObjectSpec. If resuming, the `write_handle` is added to that spec.
+
+ This method is not applicable for `open` methods.
+ """
+write_state: _WriteState=state["write_state"]
+
+initial_request=storage_type.BidiWriteObjectRequest()
+
+# Determine if we need to send WriteObjectSpec or AppendObjectSpec
+ifisinstance(write_state.spec, storage_type.WriteObjectSpec):
+initial_request.write_object_spec=write_state.spec
+else:
+ifwrite_state.write_handle:
+write_state.spec.write_handle=write_state.write_handle
+
+ifwrite_state.routing_token:
+write_state.spec.routing_token=write_state.routing_token
+initial_request.append_object_spec=write_state.spec
+
+yieldinitial_request
+
+# The buffer should already be seeked to the correct position (persisted_size)
+# by the `recover_state_on_failure` method before this is called.
+whilenotwrite_state.is_finalized:
+chunk=write_state.user_buffer.read(write_state.chunk_size)
+
+# End of File detection
+ifnotchunk:
+return
+
+checksummed_data=storage_type.ChecksummedData(content=chunk)
+checksum=google_crc32c.Checksum(chunk)
+checksummed_data.crc32c=int.from_bytes(checksum.digest(), "big")
+
+request=storage_type.BidiWriteObjectRequest(
+write_offset=write_state.bytes_sent,
+checksummed_data=checksummed_data,
+ )
+write_state.bytes_sent+=len(chunk)
+
+yieldrequest
+
+defupdate_state_from_response(
+self, response: storage_type.BidiWriteObjectResponse, state: Dict[str, Any]
+ ) ->None:
+"""Processes a server response and updates the write state."""
+write_state: _WriteState=state["write_state"]
+
+ifresponse.persisted_size:
+write_state.persisted_size=response.persisted_size
+
+ifresponse.write_handle:
+write_state.write_handle=response.write_handle
+
+ifresponse.resource:
+write_state.persisted_size=response.resource.size
+ifresponse.resource.finalize_time:
+write_state.is_finalized=True
+
+asyncdefrecover_state_on_failure(
+self, error: Exception, state: Dict[str, Any]
+ ) ->None:
+"""
+ Handles errors, specifically BidiWriteObjectRedirectedError, and rewinds state.
+
+ This method rewinds the user buffer and internal byte tracking to the
+ last confirmed 'persisted_size' from the server.
+ """
+write_state: _WriteState=state["write_state"]
+cause=getattr(error, "cause", error)
+
+# Extract routing token and potentially a new write handle for redirection.
+ifisinstance(cause, BidiWriteObjectRedirectedError):
+ifcause.routing_token:
+write_state.routing_token=cause.routing_token
+
+redirect_handle=getattr(cause, "write_handle", None)
+ifredirect_handle:
+write_state.write_handle=redirect_handle
+
+# We must assume any data sent beyond 'persisted_size' was lost.
+# Reset the user buffer to the last known good byte confirmed by the server.
+write_state.user_buffer.seek(write_state.persisted_size)
+write_state.bytes_sent=write_state.persisted_size

tests/unit/asyncio/retry/test_reads_resumption_strategy.py

@@ -27,7 +27,9 @@
 fromgoogle.cloud._storage_v2.types.storageimportBidiReadObjectRedirectedError
 
 _READ_ID=1
-LOGGER_NAME="google.cloud.storage._experimental.asyncio.retry.reads_resumption_strategy"
+LOGGER_NAME= (
+"google.cloud.storage._experimental.asyncio.retry.reads_resumption_strategy"
+)
 
 
 classTestDownloadState(unittest.TestCase):
@@ -309,7 +311,9 @@ def test_update_state_missing_read_range_logs_warning(self):
 withself.assertLogs(LOGGER_NAME, level="WARNING") ascm:
 self.strategy.update_state_from_response(response, self.state)
 
-self.assertTrue(any("missing read_range field"inoutputforoutputincm.output))
+self.assertTrue(
+any("missing read_range field"inoutputforoutputincm.output)
+ )
 
 deftest_update_state_unknown_id_logs_warning(self):
 """Verify we log a warning and continue when read_id is unknown."""
@@ -320,8 +324,12 @@ def test_update_state_unknown_id_logs_warning(self):
 withself.assertLogs(LOGGER_NAME, level="WARNING") ascm:
 self.strategy.update_state_from_response(response, self.state)
 
-self.assertTrue(any(f"unknown or stale read_id {unknown_id}"inoutputforoutputincm.output))
-
+self.assertTrue(
+any(
+f"unknown or stale read_id {unknown_id}"inoutput
+foroutputincm.output
+ )
+ )
 
 # --- Recovery Tests ---