feat(experimental): flush the last chunk in append method

google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py

@@ -189,7 +189,7 @@ async def append(self, data: bytes) -> None:
 
         This method sends the provided `data` to the GCS server in chunks.
         and persists data in GCS at every `_MAX_BUFFER_SIZE_BYTES` bytes by
-        calling `self.simple_flush`.
+        sending `flush=true`.
 
         :type data: bytes
         :param data: The bytes to append to the object.
@@ -214,21 +214,28 @@ async def append(self, data: bytes) -> None:
         while start_idx < total_bytes:
             end_idx = min(start_idx + _MAX_CHUNK_SIZE_BYTES, total_bytes)
             data_chunk = data[start_idx:end_idx]
-            await self.write_obj_stream.send(
-                _storage_v2.BidiWriteObjectRequest(
-                    write_offset=self.offset,
-                    checksummed_data=_storage_v2.ChecksummedData(
-                        content=data_chunk,
-                        crc32c=int.from_bytes(Checksum(data_chunk).digest(), "big"),
-                    ),
-                )
+            is_last_chunk = end_idx == total_bytes
+
+            request = _storage_v2.BidiWriteObjectRequest(
+                write_offset=self.offset,
+                checksummed_data=_storage_v2.ChecksummedData(
+                    content=data_chunk,
+                    crc32c=int.from_bytes(Checksum(data_chunk).digest(), "big"),
+                ),
             )
             chunk_size = end_idx - start_idx
             self.offset += chunk_size
             self.bytes_appended_since_last_flush += chunk_size
+
             if self.bytes_appended_since_last_flush >= self.flush_interval:
-                await self.simple_flush()
+                request.flush = True
                 self.bytes_appended_since_last_flush = 0
+
+            if is_last_chunk:
+                request.state_lookup = True
+                request.flush = True
+
+            await self.write_obj_stream.send(request)
             start_idx = end_idx
 
     async def simple_flush(self) -> None:
@@ -292,15 +299,13 @@ async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object]
             raise ValueError("Stream is not open. Call open() before close().")
 
         if finalize_on_close:
-            await self.finalize()
-        else:
-            await self.flush()
+            return await self.finalize()
 
         await self.write_obj_stream.close()
 
         self._is_stream_open = False
         self.offset = None
-        return self.object_resource if finalize_on_close else self.persisted_size
+        return self.persisted_size
 
     async def finalize(self) -> _storage_v2.Object:
         """Finalizes the Appendable Object.
@@ -322,6 +327,10 @@ async def finalize(self) -> _storage_v2.Object:
         response = await self.write_obj_stream.recv()
         self.object_resource = response.resource
         self.persisted_size = self.object_resource.size
+        await self.write_obj_stream.close()
+
+        self._is_stream_open = False
+        self.offset = None
         return self.object_resource
 
     # helper methods.

tests/system/test_zonal.py

@@ -307,3 +307,51 @@ async def _read_and_verify(expected_content, generation=None):
     del mrd
     del mrd_2
     gc.collect()
+
+
+@pytest.mark.asyncio
+async def test_append_flushes_and_state_lookup(storage_client, blobs_to_delete):
+    """
+    System test for AsyncAppendableObjectWriter, verifying flushing behavior
+    for both small and large appends.
+    """
+    object_name = f"test-append-flush-varied-size-{uuid.uuid4()}"
+    grpc_client = AsyncGrpcClient().grpc_client
+    writer = AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name)
+
+    # Schedule for cleanup
+    blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+    # --- Part 1: Test with small data ---
+    small_data = b"small data"
+
+    await writer.open()
+    assert writer._is_stream_open
+
+    await writer.append(small_data)
+    persisted_size = await writer.state_lookup()
+    assert persisted_size == len(small_data)
+
+    # --- Part 2: Test with large data ---
+    large_data = os.urandom(38 * 1024 * 1024)
+
+    # Append data larger than the default flush interval (16 MiB).
+    # This should trigger the interval-based flushing logic.
+    await writer.append(large_data)
+
+    # Verify the total data has been persisted.
+    total_size = len(small_data) + len(large_data)
+    persisted_size = await writer.state_lookup()
+    assert persisted_size == total_size
+
+    # --- Part 3: Finalize and verify ---
+    final_object = await writer.close(finalize_on_close=True)
+
+    assert not writer._is_stream_open
+    assert final_object.size == total_size
+
+    # Verify the full content of the object.
+    full_data = small_data + large_data
+    blob = storage_client.bucket(_ZONAL_BUCKET).get_blob(object_name)
+    content = blob.download_as_bytes()
+    assert content == full_data

tests/unit/asyncio/retry/test_writes_resumption_strategy.py

@@ -206,7 +206,9 @@ def test_update_state_from_response(self):
         strategy.update_state_from_response(response2, state)
         self.assertEqual(write_state.persisted_size, 1024)
 
-        final_resource = storage_type.Object(name="test-object", bucket="b", size=2048, finalize_time=datetime.now())
+        final_resource = storage_type.Object(
+            name="test-object", bucket="b", size=2048, finalize_time=datetime.now()
+        )
         response3 = storage_type.BidiWriteObjectResponse(resource=final_resource)
         strategy.update_state_from_response(response3, state)
         self.assertEqual(write_state.persisted_size, 2048)

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -364,6 +364,7 @@ async def test_close(mock_write_object_stream, mock_client):
     writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
     writer._is_stream_open = True
     writer.offset = 1024
+    writer.persisted_size = 1024
     mock_stream = mock_write_object_stream.return_value
     mock_stream.send = mock.AsyncMock()
     mock_stream.recv = mock.AsyncMock(
@@ -435,13 +436,15 @@ async def test_finalize(mock_write_object_stream, mock_client):
     mock_stream.recv = mock.AsyncMock(
         return_value=_storage_v2.BidiWriteObjectResponse(resource=mock_resource)
     )
+    mock_stream.close = mock.AsyncMock()
 
     gcs_object = await writer.finalize()
 
     mock_stream.send.assert_awaited_once_with(
         _storage_v2.BidiWriteObjectRequest(finish_write=True)
     )
     mock_stream.recv.assert_awaited_once()
+    mock_stream.close.assert_awaited_once()
     assert writer.object_resource == mock_resource
     assert writer.persisted_size == 123
     assert gcs_object == mock_resource
@@ -501,30 +504,33 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client
     writer.persisted_size = 100
     mock_stream = mock_write_object_stream.return_value
     mock_stream.send = mock.AsyncMock()
-    writer.simple_flush = mock.AsyncMock()
 
     data = b"a" * (_MAX_CHUNK_SIZE_BYTES + 1)
     await writer.append(data)
 
     assert mock_stream.send.await_count == 2
-    first_call = mock_stream.send.await_args_list[0]
-    second_call = mock_stream.send.await_args_list[1]
+    first_request = mock_stream.send.await_args_list[0].args[0]
+    second_request = mock_stream.send.await_args_list[1].args[0]
 
     # First chunk
-    assert first_call[0][0].write_offset == 100
-    assert len(first_call[0][0].checksummed_data.content) == _MAX_CHUNK_SIZE_BYTES
-    assert first_call[0][0].checksummed_data.crc32c == int.from_bytes(
+    assert first_request.write_offset == 100
+    assert len(first_request.checksummed_data.content) == _MAX_CHUNK_SIZE_BYTES
+    assert first_request.checksummed_data.crc32c == int.from_bytes(
         Checksum(data[:_MAX_CHUNK_SIZE_BYTES]).digest(), byteorder="big"
     )
-    # Second chunk
-    assert second_call[0][0].write_offset == 100 + _MAX_CHUNK_SIZE_BYTES
-    assert len(second_call[0][0].checksummed_data.content) == 1
-    assert second_call[0][0].checksummed_data.crc32c == int.from_bytes(
+    assert not first_request.flush
+    assert not first_request.state_lookup
+
+    # Second chunk (last chunk)
+    assert second_request.write_offset == 100 + _MAX_CHUNK_SIZE_BYTES
+    assert len(second_request.checksummed_data.content) == 1
+    assert second_request.checksummed_data.crc32c == int.from_bytes(
         Checksum(data[_MAX_CHUNK_SIZE_BYTES:]).digest(), byteorder="big"
     )
+    assert second_request.flush
+    assert second_request.state_lookup
 
     assert writer.offset == 100 + len(data)
-    writer.simple_flush.assert_not_awaited()
 
 
 @pytest.mark.asyncio
@@ -541,12 +547,24 @@ async def test_append_flushes_when_buffer_is_full(
     writer.persisted_size = 0
     mock_stream = mock_write_object_stream.return_value
     mock_stream.send = mock.AsyncMock()
-    writer.simple_flush = mock.AsyncMock()
 
     data = b"a" * _DEFAULT_FLUSH_INTERVAL_BYTES
     await writer.append(data)
 
-    writer.simple_flush.assert_awaited_once()
+    num_chunks = _DEFAULT_FLUSH_INTERVAL_BYTES // _MAX_CHUNK_SIZE_BYTES
+    assert mock_stream.send.await_count == num_chunks
+
+    # All but the last request should not have flush or state_lookup set.
+    for i in range(num_chunks - 1):
+        request = mock_stream.send.await_args_list[i].args[0]
+        assert not request.flush
+        assert not request.state_lookup
+
+    # The last request should have flush and state_lookup set.
+    last_request = mock_stream.send.await_args_list[-1].args[0]
+    assert last_request.flush
+    assert last_request.state_lookup
+    assert writer.bytes_appended_since_last_flush == 0
 
 
 @pytest.mark.asyncio
@@ -561,12 +579,17 @@ async def test_append_handles_large_data(mock_write_object_stream, mock_client):
     writer.persisted_size = 0
     mock_stream = mock_write_object_stream.return_value
     mock_stream.send = mock.AsyncMock()
-    writer.simple_flush = mock.AsyncMock()
 
     data = b"a" * (_DEFAULT_FLUSH_INTERVAL_BYTES * 2 + 1)
     await writer.append(data)
 
-    assert writer.simple_flush.await_count == 2
+    flushed_requests = [
+        call.args[0] for call in mock_stream.send.await_args_list if call.args[0].flush
+    ]
+    assert len(flushed_requests) == 3
+
+    last_request = mock_stream.send.await_args_list[-1].args[0]
+    assert last_request.state_lookup
 
 
 @pytest.mark.asyncio
@@ -584,17 +607,25 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client):
     writer.persisted_size = 0
     mock_stream = mock_write_object_stream.return_value
     mock_stream.send = mock.AsyncMock()
-    writer.simple_flush = mock.AsyncMock()
 
     data1 = b"a" * (_MAX_CHUNK_SIZE_BYTES + 10)
     await writer.append(data1)
 
+    assert mock_stream.send.await_count == 2
+    last_request_data1 = mock_stream.send.await_args_list[-1].args[0]
+    assert last_request_data1.flush
+    assert last_request_data1.state_lookup
+
     data2 = b"b" * (_MAX_CHUNK_SIZE_BYTES + 20)
     await writer.append(data2)
 
+    assert mock_stream.send.await_count == 4
+    last_request_data2 = mock_stream.send.await_args_list[-1].args[0]
+    assert last_request_data2.flush
+    assert last_request_data2.state_lookup
+
     total_data_length = len(data1) + len(data2)
     assert writer.offset == total_data_length
-    assert writer.simple_flush.await_count == 0
 
 
 @pytest.mark.asyncio