feat(experimental): integrate writes strategy and appendable object writer (#1695)

Integrate all the components required for bidi writes retries

Author: Pulkit0110

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

@@ -42,6 +42,7 @@ class AsyncGrpcClient:
  (Optional) Whether to attempt to use DirectPath for gRPC connections.
  Defaults to ``True``.
  """
+
 def__init__(
 self,
 credentials=None,

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

@@ -493,4 +493,4 @@ async def close(self):
 
 @property
 defis_stream_open(self) ->bool:
-returnself._is_stream_open
+returnself._is_stream_open

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

@@ -195,4 +195,4 @@ async def recv(self) -> _storage_v2.BidiReadObjectResponse:
 
 @property
 defis_stream_open(self) ->bool:
-returnself._is_stream_open
+returnself._is_stream_open

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

@@ -21,9 +21,9 @@
 if you want to use these Rapid Storage APIs.
 
 """
-fromtypingimportOptional
-from . import_utils
+fromtypingimportList, Optional, Tuple
 fromgoogle.cloudimport_storage_v2
+fromgoogle.cloud.storage._experimental.asyncioimport_utils
 fromgoogle.cloud.storage._experimental.asyncio.async_grpc_clientimportAsyncGrpcClient
 fromgoogle.cloud.storage._experimental.asyncio.async_abstract_object_streamimport (
 _AsyncAbstractObjectStream,
@@ -72,6 +72,7 @@ def __init__(
 object_name: str,
 generation_number: Optional[int] =None, # None means new object
 write_handle: Optional[_storage_v2.BidiWriteHandle] =None,
+routing_token: Optional[str] =None,
  ) ->None:
 ifclientisNone:
 raiseValueError("client must be provided")
@@ -87,6 +88,7 @@ def __init__(
  )
 self.client: AsyncGrpcClient.grpc_client=client
 self.write_handle: Optional[_storage_v2.BidiWriteHandle] =write_handle
+self.routing_token: Optional[str] =routing_token
 
 self._full_bucket_name=f"projects/_/buckets/{self.bucket_name}"
 
@@ -101,7 +103,7 @@ def __init__(
 self.persisted_size=0
 self.object_resource: Optional[_storage_v2.Object] =None
 
-asyncdefopen(self) ->None:
+asyncdefopen(self, metadata: Optional[List[Tuple[str, str]]] =None) ->None:
 """
  Opens the bidi-gRPC connection to write to the object.
 
@@ -110,7 +112,7 @@ async def open(self) -> None:
 
  :rtype: None
  :raises ValueError: If the stream is already open.
- :raises google.api_core.exceptions.FailedPrecondition:
+ :raises google.api_core.exceptions.FailedPrecondition:
  if `generation_number` is 0 and object already exists.
  """
 ifself._is_stream_open:
@@ -121,9 +123,6 @@ async def open(self) -> None:
 # Created object type would be Appendable Object.
 # if `generation_number` == 0 new object will be created only if there
 # isn't any existing object.
-is_open_via_write_handle= (
-self.write_handleisnotNoneandself.generation_number
- )
 ifself.generation_numberisNoneorself.generation_number==0:
 self.first_bidi_write_req=_storage_v2.BidiWriteObjectRequest(
 write_object_spec=_storage_v2.WriteObjectSpec(
@@ -140,44 +139,46 @@ async def open(self) -> None:
 bucket=self._full_bucket_name,
 object=self.object_name,
 generation=self.generation_number,
-write_handle=self.write_handle,
+write_handle=self.write_handleifself.write_handleelseNone,
+routing_token=self.routing_tokenifself.routing_tokenelseNone,
  ),
  )
+
+request_param_values= [f"bucket={self._full_bucket_name}"]
+final_metadata= []
+ifmetadata:
+forkey, valueinmetadata:
+ifkey=="x-goog-request-params":
+request_param_values.append(value)
+else:
+final_metadata.append((key, value))
+
+final_metadata.append(("x-goog-request-params", ",".join(request_param_values)))
+
 self.socket_like_rpc=AsyncBidiRpc(
-self.rpc, initial_request=self.first_bidi_write_req, metadata=self.metadata
+self.rpc,
+initial_request=self.first_bidi_write_req,
+metadata=final_metadata,
  )
 
 awaitself.socket_like_rpc.open() # this is actually 1 send
 response=awaitself.socket_like_rpc.recv()
 self._is_stream_open=True
-ifis_open_via_write_handle:
-# Don't use if not response.persisted_size because this will be true
-# if persisted_size==0 (0 is considered "Falsy" in Python)
-ifresponse.persisted_sizeisNone:
-raiseValueError(
-"Failed to obtain persisted_size after opening the stream via write_handle"
- )
+
+ifresponse.persisted_size:
 self.persisted_size=response.persisted_size
-else:
-ifnotresponse.resource:
-raiseValueError(
-"Failed to obtain object resource after opening the stream"
- )
-ifnotresponse.resource.generation:
-raiseValueError(
-"Failed to obtain object generation after opening the stream"
- )
+
+ifresponse.resource:
 ifnotresponse.resource.size:
 # Appending to a 0 byte appendable object.
 self.persisted_size=0
 else:
 self.persisted_size=response.resource.size
 
-ifnotresponse.write_handle:
-raiseValueError("Failed to obtain write_handle after opening the stream")
+self.generation_number=response.resource.generation
 
-self.generation_number=response.resource.generation
-self.write_handle=response.write_handle
+ifresponse.write_handle:
+self.write_handle=response.write_handle
 
 asyncdefclose(self) ->None:
 """Closes the bidi-gRPC connection."""
@@ -220,7 +221,14 @@ async def recv(self) -> _storage_v2.BidiWriteObjectResponse:
 ifnotself._is_stream_open:
 raiseValueError("Stream is not open")
 response=awaitself.socket_like_rpc.recv()
-_utils.update_write_handle_if_exists(self, response)
+# Update write_handle if present in response
+ifresponse:
+ifresponse.write_handle:
+self.write_handle=response.write_handle
+ifresponse.persisted_sizeisnotNone:
+self.persisted_size=response.persisted_size
+ifresponse.resourceandresponse.resource.size:
+self.persisted_size=response.resource.size
 returnresponse
 
 @property

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

@@ -18,12 +18,19 @@
 fromtypingimportTuple, Optional
 
 fromgoogle.api_coreimportexceptions
-fromgoogle.cloud._storage_v2.typesimportBidiReadObjectRedirectedError
+fromgoogle.cloud._storage_v2.typesimport (
+BidiReadObjectRedirectedError,
+BidiWriteObjectRedirectedError,
+)
 fromgoogle.rpcimportstatus_pb2
 
 _BIDI_READ_REDIRECTED_TYPE_URL= (
 "type.googleapis.com/google.storage.v2.BidiReadObjectRedirectedError"
 )
+_BIDI_WRITE_REDIRECTED_TYPE_URL= (
+"type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError"
+)
+logger=logging.getLogger(__name__)
 
 
 def_handle_redirect(
@@ -78,6 +85,41 @@ def _handle_redirect(
 read_handle=redirect_proto.read_handle
 break
 exceptExceptionase:
-logging.ERROR(f"Error unpacking redirect: {e}")
+logger.error(f"Error unpacking redirect: {e}")
 
 returnrouting_token, read_handle
+
+
+def_extract_bidi_writes_redirect_proto(exc: Exception):
+grpc_error=None
+ifisinstance(exc, exceptions.Aborted) andexc.errors:
+grpc_error=exc.errors[0]
+
+ifgrpc_error:
+ifisinstance(grpc_error, BidiWriteObjectRedirectedError):
+returngrpc_error
+
+ifhasattr(grpc_error, "trailing_metadata"):
+trailers=grpc_error.trailing_metadata()
+ifnottrailers:
+return
+
+status_details_bin=None
+forkey, valueintrailers:
+ifkey=="grpc-status-details-bin":
+status_details_bin=value
+break
+
+ifstatus_details_bin:
+status_proto=status_pb2.Status()
+try:
+status_proto.ParseFromString(status_details_bin)
+fordetailinstatus_proto.details:
+ifdetail.type_url==_BIDI_WRITE_REDIRECTED_TYPE_URL:
+redirect_proto=BidiWriteObjectRedirectedError.deserialize(
+detail.value
+ )
+returnredirect_proto
+exceptException:
+logger.error("Error unpacking redirect details from gRPC error.")
+pass

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

@@ -12,40 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-fromtypingimportAny, Dict, IO, Iterable, Optional, Union
+fromtypingimportAny, Dict, IO, List, Optional, Union
 
 importgoogle_crc32c
 fromgoogle.cloud._storage_v2.typesimportstorageasstorage_type
 fromgoogle.cloud._storage_v2.types.storageimportBidiWriteObjectRedirectedError
 fromgoogle.cloud.storage._experimental.asyncio.retry.base_strategyimport (
 _BaseResumptionStrategy,
 )
+fromgoogle.cloud.storage._experimental.asyncio.retry._helpersimport (
+_extract_bidi_writes_redirect_proto,
+)
 
 
 class_WriteState:
 """A helper class to track the state of a single upload operation.
 
- :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.
+
+ :type flush_interval: int
+ :param flush_interval: The flush interval at which the data is flushed.
  """
 
 def__init__(
 self,
-spec: Union[storage_type.AppendObjectSpec, storage_type.WriteObjectSpec],
 chunk_size: int,
 user_buffer: IO[bytes],
+flush_interval: int,
  ):
-self.spec=spec
 self.chunk_size=chunk_size
 self.user_buffer=user_buffer
 self.persisted_size: int=0
 self.bytes_sent: int=0
+self.bytes_since_last_flush: int=0
+self.flush_interval: int=flush_interval
 self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] =None
 self.routing_token: Optional[str] =None
 self.is_finalized: bool=False
@@ -56,39 +60,22 @@ class _WriteResumptionStrategy(_BaseResumptionStrategy):
 
 defgenerate_requests(
 self, state: Dict[str, Any]
- ) ->Iterable[storage_type.BidiWriteObjectRequest]:
+ ) ->List[storage_type.BidiWriteObjectRequest]:
 """Generates BidiWriteObjectRequests to resume or continue the upload.
 
- For Appendable Objects, every stream opening should send an
- AppendObjectSpec. If resuming, the `write_handle` is added to that spec.
-
  This method is not applicable for `open` methods.
  """
 write_state: _WriteState=state["write_state"]
 
-initial_request=storage_type.BidiWriteObjectRequest()
-
-# Determine if we need to send WriteObjectSpec or AppendObjectSpec
-ifisinstance(write_state.spec, storage_type.WriteObjectSpec):
-initial_request.write_object_spec=write_state.spec
-else:
-ifwrite_state.write_handle:
-write_state.spec.write_handle=write_state.write_handle
-
-ifwrite_state.routing_token:
-write_state.spec.routing_token=write_state.routing_token
-initial_request.append_object_spec=write_state.spec
-
-yieldinitial_request
-
+requests= []
 # The buffer should already be seeked to the correct position (persisted_size)
 # by the `recover_state_on_failure` method before this is called.
 whilenotwrite_state.is_finalized:
 chunk=write_state.user_buffer.read(write_state.chunk_size)
 
 # End of File detection
 ifnotchunk:
-return
+break
 
 checksummed_data=storage_type.ChecksummedData(content=chunk)
 checksum=google_crc32c.Checksum(chunk)
@@ -98,16 +85,25 @@ def generate_requests(
 write_offset=write_state.bytes_sent,
 checksummed_data=checksummed_data,
  )
-write_state.bytes_sent+=len(chunk)
+chunk_len=len(chunk)
+write_state.bytes_sent+=chunk_len
+write_state.bytes_since_last_flush+=chunk_len
+
+ifwrite_state.bytes_since_last_flush>=write_state.flush_interval:
+request.flush=True
+# reset counter after marking flush
+write_state.bytes_since_last_flush=0
 
-yieldrequest
+requests.append(request)
+returnrequests
 
 defupdate_state_from_response(
 self, response: storage_type.BidiWriteObjectResponse, state: Dict[str, Any]
  ) ->None:
 """Processes a server response and updates the write state."""
 write_state: _WriteState=state["write_state"]
-
+ifresponseisNone:
+return
 ifresponse.persisted_size:
 write_state.persisted_size=response.persisted_size
 
@@ -129,18 +125,23 @@ async def recover_state_on_failure(
  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 for redirection.
-ifisinstance(cause, BidiWriteObjectRedirectedError):
-ifcause.routing_token:
-write_state.routing_token=cause.routing_token
+redirect_proto=None
 
-redirect_handle=getattr(cause, "write_handle", None)
-ifredirect_handle:
-write_state.write_handle=redirect_handle
+ifisinstance(error, BidiWriteObjectRedirectedError):
+redirect_proto=error
+else:
+redirect_proto=_extract_bidi_writes_redirect_proto(error)
+
+# Extract routing token and potentially a new write handle for redirection.
+ifredirect_proto:
+ifredirect_proto.routing_token:
+write_state.routing_token=redirect_proto.routing_token
+ifredirect_proto.write_handle:
+write_state.write_handle=redirect_proto.write_handle
 
 # We must assume any data sent beyond 'persisted_size' was lost.
 # 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
+write_state.bytes_since_last_flush=0