fix: Add simple flush for optimized writes (#1633)

fix: Add simple flush for optimized writes

Author: chandra-siri

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

@@ -199,10 +199,28 @@ async def append(self, data: bytes) -> None:
 self.offset+=chunk_size
 bytes_to_flush+=chunk_size
 ifbytes_to_flush>=_MAX_BUFFER_SIZE_BYTES:
-awaitself.flush()
+awaitself.simple_flush()
 bytes_to_flush=0
 start_idx=end_idx
 
+asyncdefsimple_flush(self) ->None:
+"""Flushes the data to the server.
+ Please note: Unlike `flush` it does not do `state_lookup`
+
+ :rtype: None
+
+ :raises ValueError: If the stream is not open (i.e., `open()` has not
+ been called).
+ """
+ifnotself._is_stream_open:
+raiseValueError("Stream is not open. Call open() before simple_flush().")
+
+awaitself.write_obj_stream.send(
+_storage_v2.BidiWriteObjectRequest(
+flush=True,
+ )
+ )
+
 asyncdefflush(self) ->int:
 """Flushes the data to the server.
 

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -224,6 +224,38 @@ async def test_flush_without_open_raises_value_error(mock_client):
 awaitwriter.flush()
 
 
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_simple_flush(mock_write_object_stream, mock_client):
+"""Test that flush sends the correct request and updates state."""
+# Arrange
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+writer._is_stream_open=True
+mock_stream=mock_write_object_stream.return_value
+mock_stream.send=mock.AsyncMock()
+
+# Act
+awaitwriter.simple_flush()
+
+# Assert
+mock_stream.send.assert_awaited_once_with(
+_storage_v2.BidiWriteObjectRequest(flush=True)
+ )
+
+
+@pytest.mark.asyncio
+asyncdeftest_simple_flush_without_open_raises_value_error(mock_client):
+"""Test that flush raises an error if the stream is not open."""
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+withpytest.raises(
+ValueError,
+match="Stream is not open. Call open\\(\\) before simple_flush\\(\\).",
+ ):
+awaitwriter.simple_flush()
+
+
 @pytest.mark.asyncio
 @mock.patch(
 "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
@@ -369,7 +401,7 @@ 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.flush=mock.AsyncMock()
+writer.simple_flush=mock.AsyncMock()
 
 data=b"a"* (_MAX_CHUNK_SIZE_BYTES+1)
 awaitwriter.append(data)
@@ -387,7 +419,7 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client
 assertlen(second_call[0][0].checksummed_data.content) ==1
 
 assertwriter.offset==100+len(data)
-writer.flush.assert_not_awaited()
+writer.simple_flush.assert_not_awaited()
 
 
 @pytest.mark.asyncio
@@ -407,12 +439,12 @@ 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.flush=mock.AsyncMock()
+writer.simple_flush=mock.AsyncMock()
 
 data=b"a"*_MAX_BUFFER_SIZE_BYTES
 awaitwriter.append(data)
 
-writer.flush.assert_awaited_once()
+writer.simple_flush.assert_awaited_once()
 
 
 @pytest.mark.asyncio
@@ -430,12 +462,12 @@ 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.flush=mock.AsyncMock()
+writer.simple_flush=mock.AsyncMock()
 
 data=b"a"* (_MAX_BUFFER_SIZE_BYTES*2+1)
 awaitwriter.append(data)
 
-assertwriter.flush.await_count==2
+assertwriter.simple_flush.await_count==2
 
 
 @pytest.mark.asyncio
@@ -453,7 +485,7 @@ 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.flush=mock.AsyncMock()
+writer.simple_flush=mock.AsyncMock()
 
 data1=b"a"* (_MAX_CHUNK_SIZE_BYTES+10)
 awaitwriter.append(data1)
@@ -463,4 +495,4 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client):
 
 total_data_length=len(data1) +len(data2)
 assertwriter.offset==total_data_length
-assertwriter.flush.await_count==0
+assertwriter.simple_flush.await_count==0