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

Pulkit0110 · web-flow · commit a57ea0ec786a · 2026-01-02T16:18:37.000+05:30
Adding writes resumption strategy which will be used for error handling
of bidi writes operation.
diff --git a/google/cloud/storage/_experimental/asyncio/retry/writes_resumption_strategy.py b/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.
+
+from typing import Any, Dict, IO, Iterable, Optional, Union
+
+import google_crc32c
+from google.cloud._storage_v2.types import storage as storage_type
+from google.cloud._storage_v2.types.storage import BidiWriteObjectRedirectedError
+from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
+    _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."""
+
+    def generate_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
+        if isinstance(write_state.spec, storage_type.WriteObjectSpec):
+            initial_request.write_object_spec = write_state.spec
+        else:
+            if write_state.write_handle:
+                write_state.spec.write_handle = write_state.write_handle
+
+            if write_state.routing_token:
+                write_state.spec.routing_token = write_state.routing_token
+            initial_request.append_object_spec = write_state.spec
+
+        yield initial_request
+
+        # The buffer should already be seeked to the correct position (persisted_size)
+        # by the `recover_state_on_failure` method before this is called.
+        while not write_state.is_finalized:
+            chunk = write_state.user_buffer.read(write_state.chunk_size)
+
+            # End of File detection
+            if not chunk:
+                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)
+
+            yield request
+
+    def update_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"]
+
+        if response.persisted_size:
+            write_state.persisted_size = response.persisted_size
+
+        if response.write_handle:
+            write_state.write_handle = response.write_handle
+
+        if response.resource:
+            write_state.persisted_size = response.resource.size
+            if response.resource.finalize_time:
+                write_state.is_finalized = True
+
+    async def recover_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.
+        if isinstance(cause, BidiWriteObjectRedirectedError):
+            if cause.routing_token:
+                write_state.routing_token = cause.routing_token
+
+            redirect_handle = getattr(cause, "write_handle", None)
+            if redirect_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
diff --git a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py
@@ -27,7 +27,9 @@
 from google.cloud._storage_v2.types.storage import BidiReadObjectRedirectedError
 
 _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"
+)
 
 
 class TestDownloadState(unittest.TestCase):
@@ -309,7 +311,9 @@ def test_update_state_missing_read_range_logs_warning(self):
         with self.assertLogs(LOGGER_NAME, level="WARNING") as cm:
             self.strategy.update_state_from_response(response, self.state)
 
-        self.assertTrue(any("missing read_range field" in output for output in cm.output))
+        self.assertTrue(
+            any("missing read_range field" in output for output in cm.output)
+        )
 
     def test_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):
         with self.assertLogs(LOGGER_NAME, level="WARNING") as cm:
             self.strategy.update_state_from_response(response, self.state)
 
-        self.assertTrue(any(f"unknown or stale read_id {unknown_id}" in output for output in cm.output))
-
+        self.assertTrue(
+            any(
+                f"unknown or stale read_id {unknown_id}" in output
+                for output in cm.output
+            )
+        )
 
     # --- Recovery Tests ---
 
diff --git a/tests/unit/asyncio/retry/test_writes_resumption_strategy.py b/tests/unit/asyncio/retry/test_writes_resumption_strategy.py
