addressing comments:

Author: Pulkit0110

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

@@ -25,15 +25,14 @@
 class _WriteState:
     """A helper class to track the state of a single upload operation.
 
-    Attributes:
-        spec (AppendObjectSpec): The specification for the object to write.
-        chunk_size (int): The size of chunks to read from the buffer.
-        user_buffer (IO[bytes]): The data source.
-        persisted_size (int): The amount of data confirmed as persisted by the server.
-        bytes_sent (int): The amount of data currently sent in the active stream.
-        write_handle (bytes | BidiWriteHandle | None): The handle for the append session.
-        routing_token (str | None): Token for routing to the correct backend.
-        is_complete (bool): Whether the upload has finished.
+    :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__(
@@ -71,23 +70,29 @@ def generate_requests(
         if write_state.routing_token:
             write_state.spec.routing_token = write_state.routing_token
 
+        # Initial request of the stream must provide the specification.
+        # If we have a write_handle, we request a state lookup to verify persisted offset.
         do_state_lookup = write_state.write_handle is not None
-        yield storage_type.BidiWriteObjectRequest(
-            append_object_spec=write_state.spec, state_lookup=do_state_lookup
+
+        # Determine if we need to send WriteObjectSpec or AppendObjectSpec
+        initial_request = storage_type.BidiWriteObjectRequest(
+            state_lookup=do_state_lookup
         )
 
+        if isinstance(write_state.spec, storage_type.WriteObjectSpec):
+            initial_request.write_object_spec = write_state.spec
+        else:
+            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_complete:
             chunk = write_state.user_buffer.read(write_state.chunk_size)
 
             # End of File detection
             if not chunk:
-                write_state.is_complete = True
-                yield storage_type.BidiWriteObjectRequest(
-                    write_offset=write_state.bytes_sent,
-                    finish_write=True,
-                )
                 return
 
             checksummed_data = storage_type.ChecksummedData(content=chunk)
@@ -122,19 +127,25 @@ def update_state_from_response(
     async def recover_state_on_failure(
         self, error: Exception, state: Dict[str, Any]
     ) -> None:
-        """Handles errors, specifically BidiWriteObjectRedirectedError, and rewinds state."""
+        """
+        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.
+        # 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
 
-            if hasattr(cause, "write_handle") and cause.write_handle:
-                write_state.write_handle = cause.write_handle
+            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.
+        # 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_writes_resumption_strategy.py

@@ -38,72 +38,86 @@ def test_ctor(self):
         strategy = self._make_one()
         self.assertIsInstance(strategy, self._get_target_class())
 
-    def test_generate_requests_initial(self):
+    def test_generate_requests_initial_new_object(self):
         """
-        Verify the initial request sequence for a new upload.
-        - First request is AppendObjectSpec with state_lookup=False.
-        - Following requests are data chunks.
-        - Final request has finish_write=True.
+        Verify the initial request sequence for a new upload (WriteObjectSpec).
         """
         strategy = self._make_one()
         mock_buffer = io.BytesIO(b"0123456789")
-        mock_spec = storage_type.AppendObjectSpec(object_="test-object")
+        # Use WriteObjectSpec for new objects
+        mock_spec = storage_type.WriteObjectSpec(
+            resource=storage_type.Object(name="test-object")
+        )
         state = {
             "write_state": _WriteState(
                 mock_spec, chunk_size=4, user_buffer=mock_buffer
             ),
-            "first_request": True,
         }
 
         requests = list(strategy.generate_requests(state))
 
-        self.assertEqual(requests[0].append_object_spec, mock_spec)
+        # Check first request (Spec)
+        self.assertEqual(requests[0].write_object_spec, mock_spec)
         self.assertFalse(requests[0].state_lookup)
-        self.assertFalse(requests[0].append_object_spec.write_handle)
 
+        # Check data chunks
         self.assertEqual(requests[1].write_offset, 0)
         self.assertEqual(requests[1].checksummed_data.content, b"0123")
         self.assertEqual(requests[2].write_offset, 4)
         self.assertEqual(requests[2].checksummed_data.content, b"4567")
         self.assertEqual(requests[3].write_offset, 8)
         self.assertEqual(requests[3].checksummed_data.content, b"89")
 
-        self.assertEqual(requests[4].write_offset, 10)
-        self.assertTrue(requests[4].finish_write)
-
-        self.assertEqual(len(requests), 5)
+        # Total requests: 1 Spec + 3 Chunks
+        self.assertEqual(len(requests), 4)
 
-    def test_generate_requests_empty_file(self):
+    def test_generate_requests_initial_existing_object(self):
         """
-        Verify the request sequence for an empty file upload.
-        - First request is AppendObjectSpec.
-        - Second and final request has finish_write=True.
+        Verify the initial request sequence for appending to an existing object (AppendObjectSpec).
         """
         strategy = self._make_one()
-        mock_buffer = io.BytesIO(b"")
-        mock_spec = storage_type.AppendObjectSpec(object_="test-object")
+        mock_buffer = io.BytesIO(b"0123")
+        # Use AppendObjectSpec for existing objects
+        mock_spec = storage_type.AppendObjectSpec(
+            object_="test-object", bucket="test-bucket"
+        )
         state = {
             "write_state": _WriteState(
                 mock_spec, chunk_size=4, user_buffer=mock_buffer
             ),
-            "first_request": True,
         }
 
         requests = list(strategy.generate_requests(state))
 
+        # Check first request (Spec)
         self.assertEqual(requests[0].append_object_spec, mock_spec)
         self.assertFalse(requests[0].state_lookup)
 
+        # Check data chunk
         self.assertEqual(requests[1].write_offset, 0)
-        self.assertTrue(requests[1].finish_write)
+        self.assertEqual(requests[1].checksummed_data.content, b"0123")
+
+    def test_generate_requests_empty_file(self):
+        """
+        Verify request sequence for an empty file. Should just be the Spec.
+        """
+        strategy = self._make_one()
+        mock_buffer = io.BytesIO(b"")
+        mock_spec = storage_type.AppendObjectSpec(object_="test-object")
+        state = {
+            "write_state": _WriteState(
+                mock_spec, chunk_size=4, user_buffer=mock_buffer
+            ),
+        }
+
+        requests = list(strategy.generate_requests(state))
 
-        self.assertEqual(len(requests), 2)
+        self.assertEqual(len(requests), 1)
+        self.assertEqual(requests[0].append_object_spec, mock_spec)
 
     def test_generate_requests_resumption(self):
         """
         Verify request sequence when resuming an upload.
-        - First request is AppendObjectSpec with write_handle and state_lookup=True.
-        - Data streaming starts from the persisted_size.
         """
         strategy = self._make_one()
         mock_buffer = io.BytesIO(b"0123456789")
@@ -115,34 +129,26 @@ def test_generate_requests_resumption(self):
         write_state.write_handle = storage_type.BidiWriteHandle(handle=b"test-handle")
         mock_buffer.seek(4)
 
-        state = {"write_state": write_state, "first_request": True}
+        state = {"write_state": write_state}
 
         requests = list(strategy.generate_requests(state))
 
+        # Check first request has handle and lookup
         self.assertEqual(
             requests[0].append_object_spec.write_handle.handle, b"test-handle"
         )
         self.assertTrue(requests[0].state_lookup)
 
+        # Check data starts from offset 4
         self.assertEqual(requests[1].write_offset, 4)
         self.assertEqual(requests[1].checksummed_data.content, b"4567")
         self.assertEqual(requests[2].write_offset, 8)
         self.assertEqual(requests[2].checksummed_data.content, b"89")
 
-        self.assertEqual(requests[3].write_offset, 10)
-        self.assertTrue(requests[3].finish_write)
-
-        self.assertEqual(len(requests), 4)
-
     @pytest.mark.asyncio
     async def test_generate_requests_after_failure_and_recovery(self):
         """
-        Verify a complex scenario:
-        1. Start upload.
-        2. Receive a persisted_size update.
-        3. Encounter an error.
-        4. Recover state.
-        5. Generate new requests for resumption.
+        Verify recovery and resumption flow.
         """
         strategy = self._make_one()
         mock_buffer = io.BytesIO(b"0123456789abcdef")
@@ -157,12 +163,11 @@ async def test_generate_requests_after_failure_and_recovery(self):
 
         strategy.update_state_from_response(
             storage_type.BidiWriteObjectResponse(
-                persisted_size=4, write_handle=b"handle-1"
+                persisted_size=4,
+                write_handle=storage_type.BidiWriteHandle(handle=b"handle-1"),
             ),
             state,
         )
-        self.assertEqual(write_state.persisted_size, 4)
-        self.assertEqual(write_state.write_handle, b"handle-1")
 
         await strategy.recover_state_on_failure(Exception("network error"), state)
 
@@ -172,23 +177,15 @@ async def test_generate_requests_after_failure_and_recovery(self):
         requests = list(strategy.generate_requests(state))
 
         self.assertTrue(requests[0].state_lookup)
-        self.assertEqual(requests[0].append_object_spec.write_handle, b"handle-1")
+        self.assertEqual(
+            requests[0].append_object_spec.write_handle.handle, b"handle-1"
+        )
 
         self.assertEqual(requests[1].write_offset, 4)
         self.assertEqual(requests[1].checksummed_data.content, b"4567")
-        self.assertEqual(requests[2].write_offset, 8)
-        self.assertEqual(requests[2].checksummed_data.content, b"89ab")
-        self.assertEqual(requests[3].write_offset, 12)
-        self.assertEqual(requests[3].checksummed_data.content, b"cdef")
-
-        self.assertEqual(requests[4].write_offset, 16)
-        self.assertTrue(requests[4].finish_write)
-        self.assertEqual(len(requests), 5)
 
     def test_update_state_from_response(self):
-        """
-        Verify that the write state is correctly updated based on server responses.
-        """
+        """Verify state updates from server responses."""
         strategy = self._make_one()
         mock_buffer = io.BytesIO(b"0123456789")
         mock_spec = storage_type.AppendObjectSpec(object_="test-object")
@@ -232,30 +229,10 @@ def test_update_state_from_response_ignores_smaller_persisted_size(self):
 
         self.assertEqual(write_state.persisted_size, 2048)
 
-    @pytest.mark.asyncio
-    async def test_recover_state_on_failure_rewinds_state(self):
-        """
-        Verify that on failure, the buffer is seeked to persisted_size
-        and bytes_sent is reset.
-        """
-        strategy = self._make_one()
-        mock_buffer = mock.MagicMock(spec=io.BytesIO)
-        mock_spec = storage_type.AppendObjectSpec(object_="test-object")
-
-        write_state = _WriteState(mock_spec, chunk_size=4, user_buffer=mock_buffer)
-        write_state.persisted_size = 100
-        write_state.bytes_sent = 200
-        state = {"write_state": write_state}
-
-        await strategy.recover_state_on_failure(Exception("any error"), state)
-
-        mock_buffer.seek.assert_called_once_with(100)
-        self.assertEqual(write_state.bytes_sent, 100)
-
     @pytest.mark.asyncio
     async def test_recover_state_on_failure_handles_redirect(self):
         """
-        Verify that on a redirect error, the routing_token is extracted and stored.
+        Verify redirection error handling.
         """
         strategy = self._make_one()
         mock_buffer = mock.MagicMock(spec=io.BytesIO)
@@ -271,11 +248,11 @@ async def test_recover_state_on_failure_handles_redirect(self):
         await strategy.recover_state_on_failure(wrapped_error, state)
 
         self.assertEqual(write_state.routing_token, "new-token-123")
-        mock_buffer.seek.assert_called_once_with(0)
-        self.assertEqual(write_state.bytes_sent, 0)
+        mock_buffer.seek.assert_called_with(0)
 
     @pytest.mark.asyncio
     async def test_recover_state_on_failure_handles_redirect_with_handle(self):
+        """Verify redirection that includes a write handle."""
         strategy = self._make_one()
         mock_buffer = mock.MagicMock(spec=io.BytesIO)
         mock_spec = storage_type.AppendObjectSpec(object_="test-object")
@@ -294,8 +271,7 @@ async def test_recover_state_on_failure_handles_redirect_with_handle(self):
         self.assertEqual(write_state.routing_token, "new-token-456")
         self.assertEqual(write_state.write_handle, b"redirect-handle")
 
-        mock_buffer.seek.assert_called_once_with(0)
-        self.assertEqual(write_state.bytes_sent, 0)
+        mock_buffer.seek.assert_called_with(0)
 
     def test_generate_requests_sends_crc32c_checksum(self):
         strategy = self._make_one()
@@ -305,13 +281,10 @@ def test_generate_requests_sends_crc32c_checksum(self):
             "write_state": _WriteState(
                 mock_spec, chunk_size=4, user_buffer=mock_buffer
             ),
-            "first_request": True,
         }
 
         requests = list(strategy.generate_requests(state))
 
-        self.assertEqual(len(requests), 3)
-
         expected_crc = google_crc32c.Checksum(b"0123")
         expected_crc_int = int.from_bytes(expected_crc.digest(), "big")
         self.assertEqual(requests[1].checksummed_data.crc32c, expected_crc_int)
@@ -323,7 +296,7 @@ def test_generate_requests_with_routing_token(self):
 
         write_state = _WriteState(mock_spec, chunk_size=4, user_buffer=mock_buffer)
         write_state.routing_token = "redirected-token"
-        state = {"write_state": write_state, "first_request": True}
+        state = {"write_state": write_state}
 
         requests = list(strategy.generate_requests(state))