fix(experimental): no state lookup while opening bidi-write stream (#1636)

fix(experimental): no state lookup while opening bidi-write stream

Author: chandra-siri

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

@@ -114,7 +114,11 @@ def __init__(
 write_handle=self.write_handle,
  )
 self._is_stream_open: bool=False
+# `offset` is the latest size of the object without staleless.
 self.offset: Optional[int] =None
+# `persisted_size` is the total_bytes persisted in the GCS server.
+# Please note: `offset` and `persisted_size` are same when the stream is
+# opened.
 self.persisted_size: Optional[int] =None
 
 asyncdefstate_lookup(self) ->int:
@@ -152,17 +156,17 @@ async def open(self) -> None:
 ifself.generationisNone:
 self.generation=self.write_obj_stream.generation_number
 self.write_handle=self.write_obj_stream.write_handle
-
-# Update self.persisted_size
-_=awaitself.state_lookup()
+self.persisted_size=self.write_obj_stream.persisted_size
 
 asyncdefappend(self, data: bytes) ->None:
 """Appends data to the Appendable object.
 
- This method sends the provided data to the GCS server in chunks. It
- maintains an internal threshold `_MAX_BUFFER_SIZE_BYTES` and will
- automatically flush the data to make it visible to readers when that
- threshold has reached.
+ calling `self.append` will append bytes at the end of the current size
+ ie. `self.offset` bytes relative to the begining of the object.
+
+ 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`.
 
  :type data: bytes
  :param data: The bytes to append to the object.

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

@@ -216,7 +216,8 @@ async def download_ranges(
 
  :type read_ranges: List[Tuple[int, int, "BytesIO"]]
  :param read_ranges: A list of tuples, where each tuple represents a
- byte range (start_byte, bytes_to_read, writeable_buffer). Buffer has
+ combintaion of byte_range and writeable buffer in format -
+ (`start_byte`, `bytes_to_read`, `writeable_buffer`). Buffer has
  to be provided by the user, and user has to make sure appropriate
  memory is available in the application to avoid out-of-memory crash.
 

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

@@ -117,7 +117,6 @@ async def open(self) -> None:
 object=self.object_name,
 generation=self.generation_number,
  ),
-state_lookup=True,
  )
 
 self.socket_like_rpc=AsyncBidiRpc(
@@ -136,11 +135,17 @@ async def open(self) -> None:
 raiseValueError(
 "Failed to obtain object generation after opening the stream"
  )
-self.generation_number=response.resource.generation
 
 ifnotresponse.write_handle:
 raiseValueError("Failed to obtain write_handle after opening the stream")
 
+ifnotresponse.resource.size:
+# Appending to a 0 byte appendable object.
+self.persisted_size=0
+else:
+self.persisted_size=response.resource.size
+
+self.generation_number=response.resource.generation
 self.write_handle=response.write_handle
 
 asyncdefclose(self) ->None:

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -133,15 +133,10 @@ async def test_open_appendable_object_writer(mock_write_object_stream, mock_clie
 writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
 mock_stream=mock_write_object_stream.return_value
 mock_stream.open=mock.AsyncMock()
-mock_stream.send=mock.AsyncMock()
-mock_stream.recv=mock.AsyncMock()
-
-mock_state_response=mock.MagicMock()
-mock_state_response.persisted_size=1024
-mock_stream.recv.return_value=mock_state_response
 
 mock_stream.generation_number=GENERATION
 mock_stream.write_handle=WRITE_HANDLE
+mock_stream.persisted_size=0
 
 # Act
 awaitwriter.open()
@@ -151,11 +146,37 @@ async def test_open_appendable_object_writer(mock_write_object_stream, mock_clie
 assertwriter._is_stream_open
 assertwriter.generation==GENERATION
 assertwriter.write_handle==WRITE_HANDLE
+assertwriter.persisted_size==0
 
-expected_request=_storage_v2.BidiWriteObjectRequest(state_lookup=True)
-mock_stream.send.assert_awaited_once_with(expected_request)
-mock_stream.recv.assert_awaited_once()
-assertwriter.persisted_size==1024
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_open_appendable_object_writer_existing_object(
+mock_write_object_stream, mock_client
+):
+"""Test the open method."""
+# Arrange
+writer=AsyncAppendableObjectWriter(
+mock_client, BUCKET, OBJECT, generation=GENERATION
+ )
+mock_stream=mock_write_object_stream.return_value
+mock_stream.open=mock.AsyncMock()
+
+mock_stream.generation_number=GENERATION
+mock_stream.write_handle=WRITE_HANDLE
+mock_stream.persisted_size=PERSISTED_SIZE
+
+# Act
+awaitwriter.open()
+
+# Assert
+mock_stream.open.assert_awaited_once()
+assertwriter._is_stream_open
+assertwriter.generation==GENERATION
+assertwriter.write_handle==WRITE_HANDLE
+assertwriter.persisted_size==PERSISTED_SIZE
 
 
 @pytest.mark.asyncio

tests/unit/asyncio/test_async_write_object_stream.py

@@ -55,6 +55,7 @@ async def instantiate_write_obj_stream(mock_client, mock_cls_async_bidi_rpc, ope
 mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
 mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
 mock_response.resource.generation=GENERATION
+mock_response.resource.size=0
 mock_response.write_handle=WRITE_HANDLE
 socket_like_rpc.recv=AsyncMock(return_value=mock_response)
 
@@ -129,6 +130,7 @@ async def test_open_for_new_object(mock_async_bidi_rpc, mock_client):
 mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
 mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
 mock_response.resource.generation=GENERATION
+mock_response.resource.size=0
 mock_response.write_handle=WRITE_HANDLE
 socket_like_rpc.recv=mock.AsyncMock(return_value=mock_response)
 
@@ -143,6 +145,7 @@ async def test_open_for_new_object(mock_async_bidi_rpc, mock_client):
 socket_like_rpc.recv.assert_called_once()
 assertstream.generation_number==GENERATION
 assertstream.write_handle==WRITE_HANDLE
+assertstream.persisted_size==0
 
 
 @pytest.mark.asyncio
@@ -158,6 +161,7 @@ async def test_open_for_existing_object(mock_async_bidi_rpc, mock_client):
 
 mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
 mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
+mock_response.resource.size=1024
 mock_response.resource.generation=GENERATION
 mock_response.write_handle=WRITE_HANDLE
 socket_like_rpc.recv=mock.AsyncMock(return_value=mock_response)
@@ -175,6 +179,7 @@ async def test_open_for_existing_object(mock_async_bidi_rpc, mock_client):
 socket_like_rpc.recv.assert_called_once()
 assertstream.generation_number==GENERATION
 assertstream.write_handle==WRITE_HANDLE
+assertstream.persisted_size==1024
 
 
 @pytest.mark.asyncio
@@ -191,6 +196,7 @@ async def test_open_when_already_open_raises_error(mock_async_bidi_rpc, mock_cli
 mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
 mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
 mock_response.resource.generation=GENERATION
+mock_response.resource.size=0
 mock_response.write_handle=WRITE_HANDLE
 socket_like_rpc.recv=mock.AsyncMock(return_value=mock_response)