feat(zb-experimental): implement append (#1620)

feat(zb-experimental): implement append

Author: chandra-siri

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

@@ -21,7 +21,7 @@
 if you want to use these Rapid Storage APIs.
 
 """
-fromtypingimportOptional
+fromtypingimportOptional, Union
 fromgoogle.cloudimport_storage_v2
 fromgoogle.cloud.storage._experimental.asyncio.async_grpc_clientimport (
 AsyncGrpcClient,
@@ -31,6 +31,10 @@
 )
 
 
+_MAX_CHUNK_SIZE_BYTES=2*1024*1024# 2 MiB
+_MAX_BUFFER_SIZE_BYTES=16*1024*1024# 16 MiB
+
+
 classAsyncAppendableObjectWriter:
 """Class for appending data to a GCS Appendable Object asynchronously."""
 
@@ -118,7 +122,13 @@ async def state_lookup(self) -> int:
 
  :rtype: int
  :returns: persisted size.
+
+ :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 state_lookup().")
+
 awaitself.write_obj_stream.send(
 _storage_v2.BidiWriteObjectRequest(
 state_lookup=True,
@@ -129,7 +139,11 @@ async def state_lookup(self) -> int:
 returnself.persisted_size
 
 asyncdefopen(self) ->None:
-"""Opens the underlying bidi-gRPC stream."""
+"""Opens the underlying bidi-gRPC stream.
+
+ :raises ValueError: If the stream is already open.
+
+ """
 ifself._is_stream_open:
 raiseValueError("Underlying bidi-gRPC stream is already open")
 
@@ -142,15 +156,65 @@ async def open(self) -> None:
 # Update self.persisted_size
 _=awaitself.state_lookup()
 
-asyncdefappend(self, data: bytes):
-raiseNotImplementedError("append is not implemented yet.")
+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.
+
+ :type data: bytes
+ :param data: The bytes to append to the object.
+
+ :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 append().")
+total_bytes=len(data)
+iftotal_bytes==0:
+# TODO: add warning.
+return
+ifself.offsetisNone:
+assertself.persisted_sizeisnotNone
+self.offset=self.persisted_size
+
+start_idx=0
+bytes_to_flush=0
+whilestart_idx<total_bytes:
+end_idx=min(start_idx+_MAX_CHUNK_SIZE_BYTES, total_bytes)
+awaitself.write_obj_stream.send(
+_storage_v2.BidiWriteObjectRequest(
+write_offset=self.offset,
+checksummed_data=_storage_v2.ChecksummedData(
+content=data[start_idx:end_idx]
+ ),
+ )
+ )
+chunk_size=end_idx-start_idx
+self.offset+=chunk_size
+bytes_to_flush+=chunk_size
+ifbytes_to_flush>=_MAX_BUFFER_SIZE_BYTES:
+awaitself.flush()
+bytes_to_flush=0
+start_idx=end_idx
 
 asyncdefflush(self) ->int:
 """Flushes the data to the server.
 
  :rtype: int
  :returns: The persisted size after flush.
+
+ :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 flush().")
+
 awaitself.write_obj_stream.send(
 _storage_v2.BidiWriteObjectRequest(
 flush=True,
@@ -162,14 +226,34 @@ async def flush(self) -> int:
 self.offset=self.persisted_size
 returnself.persisted_size
 
-asyncdefclose(self, finalize_on_close=False) ->int:
-"""Returns persisted_size"""
+asyncdefclose(self, finalize_on_close=False) ->Union[int, _storage_v2.Object]:
+"""Closes the underlying bidi-gRPC stream.
+
+ :type finalize_on_close: bool
+ :param finalize_on_close: Finalizes the Appendable Object. No more data
+ can be appended.
+
+ rtype: Union[int, _storage_v2.Object]
+ returns: Updated `self.persisted_size` by default after closing the
+ bidi-gRPC stream. However, if `finalize_on_close=True` is passed,
+ returns the finalized object resource.
+
+ :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 close().")
+
 iffinalize_on_close:
 awaitself.finalize()
+else:
+awaitself.flush()
+awaitself.write_obj_stream.close()
 
-awaitself.write_obj_stream.close()
 self._is_stream_open=False
 self.offset=None
+returnself.object_resourceiffinalize_on_closeelseself.persisted_size
 
 asyncdeffinalize(self) ->_storage_v2.Object:
 """Finalizes the Appendable Object.
@@ -178,12 +262,20 @@ async def finalize(self) -> _storage_v2.Object:
 
  rtype: google.cloud.storage_v2.types.Object
  returns: The finalized object resource.
+
+ :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 finalize().")
+
 awaitself.write_obj_stream.send(
 _storage_v2.BidiWriteObjectRequest(finish_write=True)
  )
 response=awaitself.write_obj_stream.recv()
 self.object_resource=response.resource
+self.persisted_size=self.object_resource.size
+returnself.object_resource
 
 # helper methods.
 asyncdefappend_from_string(self, data: str):