address comments and refactor other files

Author: Pulkit0110

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

@@ -34,6 +34,7 @@ def raise_if_no_fast_crc32c():
             "For more information, see https://github.com/googleapis/python-crc32c."
         )
 
+
 def update_write_handle_if_exists(obj, response):
     """Update the write_handle attribute of an object if it exists in the response."""
     if hasattr(response, "write_handle") and response.write_handle is not None:

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

@@ -102,7 +102,10 @@ def _is_write_retryable(exc):
                     if detail.type_url == _BIDI_WRITE_REDIRECTED_TYPE_URL:
                         return True
             except Exception:
-                logger.error("Error unpacking redirect details from gRPC error. Exception: ", {exc})
+                logger.error(
+                    "Error unpacking redirect details from gRPC error. Exception: ",
+                    {exc},
+                )
                 return False
     return False
 
@@ -296,7 +299,10 @@ async def _do_open():
                     try:
                         await self.write_obj_stream.close()
                     except Exception as e:
-                        logger.warning("Error closing previous write stream during open retry. Got exception: ", {e})
+                        logger.warning(
+                            "Error closing previous write stream during open retry. Got exception: ",
+                            {e},
+                        )
                 self.write_obj_stream = None
                 self._is_stream_open = False
 
@@ -567,7 +573,6 @@ async def finalize(self) -> _storage_v2.Object:
     def is_stream_open(self) -> bool:
         return self._is_stream_open
 
-
     # helper methods.
     async def append_from_string(self, data: str):
         """

google/cloud/storage/_experimental/asyncio/async_grpc_client.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_write_object_stream.py

@@ -23,6 +23,7 @@
 """
 from typing import List, Optional, Tuple
 from google.cloud import _storage_v2
+from google.cloud.storage._experimental.asyncio import _utils
 from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
 from google.cloud.storage._experimental.asyncio.async_abstract_object_stream import (
     _AsyncAbstractObjectStream,
@@ -59,7 +60,7 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
         same name already exists, it will be overwritten the moment
         `writer.open()` is called.
 
-    :type write_handle: bytes
+    :type write_handle: _storage_v2.BidiWriteHandle
     :param write_handle: (Optional) An existing handle for writing the object.
                         If provided, opening the bidi-gRPC connection will be faster.
     """
@@ -70,7 +71,7 @@ def __init__(
         bucket_name: str,
         object_name: str,
         generation_number: Optional[int] = None,  # None means new object
-        write_handle: Optional[bytes] = None,
+        write_handle: Optional[_storage_v2.BidiWriteHandle] = None,
         routing_token: Optional[str] = None,
     ) -> None:
         if client is None:
@@ -86,7 +87,7 @@ def __init__(
             generation_number=generation_number,
         )
         self.client: AsyncGrpcClient.grpc_client = client
-        self.write_handle: Optional[bytes] = write_handle
+        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}"
@@ -117,8 +118,6 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
         if self._is_stream_open:
             raise ValueError("Stream is already open")
 
-        write_handle = self.write_handle if self.write_handle else None
-
         # Create a new object or overwrite existing one if generation_number
         # is None. This makes it consistent with GCS JSON API behavior.
         # Created object type would be Appendable Object.
@@ -140,27 +139,26 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
                     bucket=self._full_bucket_name,
                     object=self.object_name,
                     generation=self.generation_number,
-                    write_handle=write_handle,
+                    write_handle=self.write_handle if self.write_handle else None,
                     routing_token=self.routing_token if self.routing_token else None,
                 ),
             )
 
-        request_params = [f"bucket={self._full_bucket_name}"]
-        other_metadata = []
+        request_param_values = [f"bucket={self._full_bucket_name}"]
+        final_metadata = []
         if metadata:
             for key, value in metadata:
                 if key == "x-goog-request-params":
-                    request_params.append(value)
+                    request_param_values.append(value)
                 else:
-                    other_metadata.append((key, value))
+                    final_metadata.append((key, value))
 
-        current_metadata = other_metadata
-        current_metadata.append(("x-goog-request-params", ",".join(request_params)))
+        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=current_metadata,
+            metadata=final_metadata,
         )
 
         await self.socket_like_rpc.open()  # this is actually 1 send
@@ -194,7 +192,7 @@ async def requests_done(self):
         """Signals that all requests have been sent."""
 
         await self.socket_like_rpc.send(None)
-        await self.socket_like_rpc.recv()
+        _utils.update_write_handle_if_exists(self, await self.socket_like_rpc.recv())
 
     async def send(
         self, bidi_write_object_request: _storage_v2.BidiWriteObjectRequest
@@ -236,4 +234,3 @@ async def recv(self) -> _storage_v2.BidiWriteObjectResponse:
     @property
     def is_stream_open(self) -> bool:
         return self._is_stream_open
-

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

@@ -25,11 +25,6 @@
 )
 
 
-_BIDI_WRITE_REDIRECTED_TYPE_URL = (
-    "type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError"
-)
-
-
 class _WriteState:
     """A helper class to track the state of a single upload operation.
 
@@ -39,22 +34,22 @@ class _WriteState:
     :type user_buffer: IO[bytes]
     :param user_buffer: The data source.
 
-    :type flush_interval: Optional[int]
+    :type flush_interval: int
     :param flush_interval: The flush interval at which the data is flushed.
     """
 
     def __init__(
         self,
         chunk_size: int,
         user_buffer: IO[bytes],
-        flush_interval: Optional[int] = None,
+        flush_interval: int,
     ):
         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: Optional[int] = flush_interval
+        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
@@ -94,10 +89,7 @@ def generate_requests(
             write_state.bytes_sent += chunk_len
             write_state.bytes_since_last_flush += chunk_len
 
-            if (
-                write_state.flush_interval
-                and write_state.bytes_since_last_flush >= write_state.flush_interval
-            ):
+            if write_state.bytes_since_last_flush >= write_state.flush_interval:
                 request.flush = True
                 # reset counter after marking flush
                 write_state.bytes_since_last_flush = 0

tests/perf/microbenchmarks/_utils.py

@@ -43,7 +43,7 @@ def publish_benchmark_extra_info(
 
     object_size = params.file_size_bytes
     num_files = params.num_files
-    total_uploaded_mib = (object_size / (1024 * 1024) * num_files)
+    total_uploaded_mib = object_size / (1024 * 1024) * num_files
     min_throughput = total_uploaded_mib / benchmark.stats["max"]
     max_throughput = total_uploaded_mib / benchmark.stats["min"]
     mean_throughput = total_uploaded_mib / benchmark.stats["mean"]
@@ -80,8 +80,8 @@ def publish_benchmark_extra_info(
 
     # Get benchmark name, rounds, and iterations
     name = benchmark.name
-    rounds = benchmark.stats['rounds']
-    iterations = benchmark.stats['iterations']
+    rounds = benchmark.stats["rounds"]
+    iterations = benchmark.stats["iterations"]
 
     # Header for throughput table
     header = "\n\n" + "-" * 125 + "\n"
@@ -98,13 +98,15 @@ def publish_benchmark_extra_info(
     print(row)
     print("-" * 125)
 
+
 class RandomBytesIO(io.RawIOBase):
     """
     A file-like object that generates random bytes using os.urandom.
     It enforces a fixed size and an upper safety cap.
     """
+
     # 10 GiB default safety cap
-    DEFAULT_CAP = 10 * 1024 * 1024 * 1024 
+    DEFAULT_CAP = 10 * 1024 * 1024 * 1024
 
     def __init__(self, size, max_size=DEFAULT_CAP):
         """
@@ -114,9 +116,11 @@ def __init__(self, size, max_size=DEFAULT_CAP):
         """
         if size is None:
             raise ValueError("Size must be defined (cannot be infinite).")
-        
+
         if size > max_size:
-            raise ValueError(f"Requested size {size} exceeds the maximum limit of {max_size} bytes (10 GiB).")
+            raise ValueError(
+                f"Requested size {size} exceeds the maximum limit of {max_size} bytes (10 GiB)."
+            )
 
         self._size = size
         self._pos = 0
@@ -125,15 +129,15 @@ def read(self, n=-1):
         # 1. Handle "read all" (n=-1)
         if n is None or n < 0:
             n = self._size - self._pos
-        
+
         # 2. Handle EOF (End of File)
         if self._pos >= self._size:
             return b""
-            
+
         # 3. Clamp read amount to remaining size
         # This ensures we stop exactly at `size` bytes.
         n = min(n, self._size - self._pos)
-        
+
         # 4. Generate data
         data = os.urandom(n)
         self._pos += len(data)

tests/perf/microbenchmarks/conftest.py

@@ -49,6 +49,7 @@ def storage_client():
     with contextlib.closing(client):
         yield client
 
+
 @pytest.fixture
 def monitor():
     """
@@ -57,6 +58,7 @@ def monitor():
     """
     return ResourceMonitor
 
+
 def publish_resource_metrics(benchmark: Any, monitor: ResourceMonitor) -> None:
     """
     Helper function to publish resource monitor results to the extra_info property.
@@ -76,7 +78,10 @@ async def upload_appendable_object(bucket_name, object_name, object_size, chunk_
     # this method is to write "appendable" objects which will be used for
     # benchmarking reads, hence not concerned performance of writes here.
     writer = AsyncAppendableObjectWriter(
-        AsyncGrpcClient().grpc_client, bucket_name, object_name, writer_options={"FLUSH_INTERVAL_BYTES": 1026 * 1024 ** 2}
+        AsyncGrpcClient().grpc_client,
+        bucket_name,
+        object_name,
+        writer_options={"FLUSH_INTERVAL_BYTES": 1026 * 1024**2},
     )
     await writer.open()
     uploaded_bytes = 0
@@ -106,11 +111,15 @@ def _upload_worker(args):
             upload_appendable_object(bucket_name, object_name, object_size, chunk_size)
         )
     else:
-        uploaded_bytes = upload_simple_object(bucket_name, object_name, object_size, chunk_size)
+        uploaded_bytes = upload_simple_object(
+            bucket_name, object_name, object_size, chunk_size
+        )
     return object_name, uploaded_bytes
 
 
-def _create_files(num_files, bucket_name, bucket_type, object_size, chunk_size=1024 * 1024 * 1024):
+def _create_files(
+    num_files, bucket_name, bucket_type, object_size, chunk_size=1024 * 1024 * 1024
+):
     """
     Create/Upload objects for benchmarking and return a list of their names.
     """

tests/perf/microbenchmarks/reads/config.py

@@ -54,8 +54,13 @@ def _get_params() -> Dict[str, List[ReadParameters]]:
     rounds = common_params["rounds"]
 
     bucket_map = {
-        "zonal": os.environ.get("DEFAULT_RAPID_ZONAL_BUCKET", config['defaults']['DEFAULT_RAPID_ZONAL_BUCKET']),
-        "regional": os.environ.get("DEFAULT_STANDARD_BUCKET", config['defaults']['DEFAULT_STANDARD_BUCKET'])
+        "zonal": os.environ.get(
+            "DEFAULT_RAPID_ZONAL_BUCKET",
+            config["defaults"]["DEFAULT_RAPID_ZONAL_BUCKET"],
+        ),
+        "regional": os.environ.get(
+            "DEFAULT_STANDARD_BUCKET", config["defaults"]["DEFAULT_STANDARD_BUCKET"]
+        ),
     }
 
     for workload in config["workload"]:

tests/perf/microbenchmarks/reads/test_reads.py

@@ -180,6 +180,7 @@ def download_files_using_mrd_multi_coro(loop, client, files, other_params, chunk
     Returns:
         float: The maximum latency (in seconds) among all coroutines.
     """
+
     async def main():
         if len(files) == 1:
             result = await download_chunks_using_mrd_async(
@@ -356,8 +357,7 @@ def download_files_mp_mc_wrapper(files_names, params, chunks, bucket_type):
 
 @pytest.mark.parametrize(
     "workload_params",
-    all_params["read_seq_multi_process"] + 
-    all_params["read_rand_multi_process"],
+    all_params["read_seq_multi_process"] + all_params["read_rand_multi_process"],
     indirect=True,
     ids=lambda p: p.name,
 )
@@ -412,4 +412,4 @@ def target_wrapper(*args, **kwargs):
 
     blobs_to_delete.extend(
         storage_client.bucket(params.bucket_name).blob(f) for f in files_names
-    )
+    )

tests/perf/microbenchmarks/resource_monitor.py

@@ -96,4 +96,4 @@ def throughput_mb_s(self):
         """Calculates combined network throughput."""
         if self.duration <= 0:
             return 0.0
-        return (self.net_sent_mb + self.net_recv_mb) / self.duration
+        return (self.net_sent_mb + self.net_recv_mb) / self.duration