adding more changes

Author: Pulkit0110

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

@@ -33,9 +33,3 @@ def raise_if_no_fast_crc32c():
             "C extension is required for faster data integrity checks."
             "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:
-        obj.write_handle = response.write_handle

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

@@ -32,7 +32,7 @@ class _AsyncAbstractObjectStream(abc.ABC):
     :param generation_number: (Optional) If present, selects a specific revision of
                               this object.
 
-    :type handle: Any
+    :type handle: bytes
     :param handle: (Optional) The handle for the object, could be read_handle or
                    write_handle, based on how the stream is used.
     """
@@ -42,13 +42,13 @@ def __init__(
         bucket_name: str,
         object_name: str,
         generation_number: Optional[int] = None,
-        handle: Optional[Any] = None,
+        handle: Optional[bytes] = None,
     ) -> None:
         super().__init__()
         self.bucket_name: str = bucket_name
         self.object_name: str = object_name
         self.generation_number: Optional[int] = generation_number
-        self.handle: Optional[Any] = handle
+        self.handle: Optional[bytes] = handle
 
     @abc.abstractmethod
     async def open(self) -> None:

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

@@ -36,7 +36,7 @@
 from google.cloud._storage_v2.types.storage import BidiWriteObjectRequest
 
 
-from . import _utils
+from ._utils import raise_if_no_fast_crc32c
 from google.cloud import _storage_v2
 from google.cloud.storage._experimental.asyncio.async_grpc_client import (
     AsyncGrpcClient,
@@ -117,8 +117,8 @@ def __init__(
         client: AsyncGrpcClient.grpc_client,
         bucket_name: str,
         object_name: str,
-        generation: Optional[int] = None,
-        write_handle: Optional[_storage_v2.BidiWriteHandle] = None,
+        generation=None,
+        write_handle=None,
         writer_options: Optional[dict] = None,
     ):
         """
@@ -164,7 +164,7 @@ def __init__(
         :type object_name: str
         :param object_name: The name of the GCS Appendable Object to be written.
 
-        :type generation: Optional[int]
+        :type generation: int
         :param generation: (Optional) If present, creates writer for that
             specific revision of that object. Use this to append data to an
             existing Appendable Object.
@@ -174,10 +174,10 @@ def __init__(
             overwriting existing objects).
 
             Warning: If `None`, a new object is created. If an object with the
-            same name already exists, it will be overwritten the moment
+            same name already exists, it will be overwritten the moment 
             `writer.open()` is called.
 
-        :type write_handle: _storage_v2.BidiWriteHandle
+        :type write_handle: bytes
         :param write_handle: (Optional) An handle for writing the object.
             If provided, opening the bidi-gRPC connection will be faster.
 
@@ -189,7 +189,7 @@ def __init__(
                 servers. Default is `_DEFAULT_FLUSH_INTERVAL_BYTES`.
                 Must be a multiple of `_MAX_CHUNK_SIZE_BYTES`.
         """
-        _utils.raise_if_no_fast_crc32c()
+        raise_if_no_fast_crc32c()
         self.client = client
         self.bucket_name = bucket_name
         self.object_name = object_name
@@ -372,12 +372,13 @@ async def append(
         attempt_count = 0
 
         def send_and_recv_generator(
-            requests_generator,
+            requests: List[BidiWriteObjectRequest],
             state: dict[str, _WriteState],
             metadata: Optional[List[Tuple[str, str]]] = None,
         ):
             async def generator():
                 nonlocal attempt_count
+                nonlocal requests
                 attempt_count += 1
                 resp = None
                 async with self._lock:
@@ -414,33 +415,16 @@ async def generator():
                         write_state.bytes_sent = write_state.persisted_size
                         write_state.bytes_since_last_flush = 0
 
-                    # Process requests from the generator
-                    # Strategy handles state_lookup and flush on the last request,
-                    # so we just stream requests directly
-                    for chunk_req in requests_generator:
-                        # Check if this is an open/state-lookup request (no checksummed_data)
-                        if chunk_req.state_lookup and not chunk_req.checksummed_data:
-                            # This is an open request - send it and get response
-                            await self.write_obj_stream.send(chunk_req)
-                            resp = await self.write_obj_stream.recv()
-
-                            # Update state from open response
-                            if resp:
-                                if resp.persisted_size is not None:
-                                    self.persisted_size = resp.persisted_size
-                                    write_state.persisted_size = resp.persisted_size
-                                    self.offset = self.persisted_size
-                                if resp.write_handle:
-                                    self.write_handle = resp.write_handle
-                                    write_state.write_handle = resp.write_handle
-                            continue
-
-                        # This is a data request - send it
+                        requests = strategy.generate_requests(state)
+
+                    num_requests = len(requests)
+                    for i, chunk_req in enumerate(requests):
+                        if i == num_requests - 1:
+                            chunk_req.state_lookup = True
+                            chunk_req.flush = True
                         await self.write_obj_stream.send(chunk_req)
 
-                    # Get final response from the last request (which has state_lookup=True)
                     resp = await self.write_obj_stream.recv()
-
                     if resp:
                         if resp.persisted_size is not None:
                             self.persisted_size = resp.persisted_size
@@ -584,6 +568,7 @@ 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_multi_range_downloader.py

@@ -129,7 +129,7 @@ async def create_mrd(
         bucket_name: str,
         object_name: str,
         generation_number: Optional[int] = None,
-        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
+        read_handle: Optional[bytes] = None,
         retry_policy: Optional[AsyncRetry] = None,
         metadata: Optional[List[Tuple[str, str]]] = None,
     ) -> AsyncMultiRangeDownloader:
@@ -149,7 +149,7 @@ async def create_mrd(
         :param generation_number: (Optional) If present, selects a specific
                                   revision of this object.
 
-        :type read_handle: _storage_v2.BidiReadHandle
+        :type read_handle: bytes
         :param read_handle: (Optional) An existing handle for reading the object.
                             If provided, opening the bidi-gRPC connection will be faster.
 
@@ -172,7 +172,7 @@ def __init__(
         bucket_name: str,
         object_name: str,
         generation_number: Optional[int] = None,
-        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
+        read_handle: Optional[bytes] = None,
     ) -> None:
         """Constructor for AsyncMultiRangeDownloader, clients are not adviced to
          use it directly. Instead it's adviced to use the classmethod `create_mrd`.
@@ -190,7 +190,7 @@ def __init__(
         :param generation_number: (Optional) If present, selects a specific revision of
                                   this object.
 
-        :type read_handle: _storage_v2.BidiReadHandle
+        :type read_handle: bytes
         :param read_handle: (Optional) An existing read handle.
         """
 
@@ -200,7 +200,7 @@ def __init__(
         self.bucket_name = bucket_name
         self.object_name = object_name
         self.generation_number = generation_number
-        self.read_handle: Optional[_storage_v2.BidiReadHandle] = read_handle
+        self.read_handle = read_handle
         self.read_obj_str: Optional[_AsyncReadObjectStream] = None
         self._is_stream_open: bool = False
         self._routing_token: Optional[str] = None
@@ -385,7 +385,7 @@ async def download_ranges(
         attempt_count = 0
 
         def send_ranges_and_get_bytes(
-            requests_generator,
+            requests: List[_storage_v2.ReadRange],
             state: Dict[str, Any],
             metadata: Optional[List[Tuple[str, str]]] = None,
         ):
@@ -395,7 +395,7 @@ async def generator():
 
                 if attempt_count > 1:
                     logger.info(
-                        f"Resuming download (attempt {attempt_count - 1})."
+                        f"Resuming download (attempt {attempt_count - 1}) for {len(requests)} ranges."
                     )
 
                 async with lock:
@@ -444,28 +444,17 @@ async def generator():
                         )
                         self._is_stream_open = True
 
-                    # Stream requests directly without materializing
-                    pending_read_ids = set()
-                    current_batch = []
-
-                    for read_range in requests_generator:
-                        pending_read_ids.add(read_range.read_id)
-                        current_batch.append(read_range)
-                        
-                        # Send batch when it reaches max size
-                        if len(current_batch) >= _MAX_READ_RANGES_PER_BIDI_READ_REQUEST:
-                            await self.read_obj_str.send(
-                                _storage_v2.BidiReadObjectRequest(read_ranges=current_batch)
-                            )
-                            current_batch = []
+                    pending_read_ids = {r.read_id for r in requests}
 
-                    # Send remaining partial batch
-                    if current_batch:
+                    # Send Requests
+                    for i in range(
+                        0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST
+                    ):
+                        batch = requests[i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST]
                         await self.read_obj_str.send(
-                            _storage_v2.BidiReadObjectRequest(read_ranges=current_batch)
+                            _storage_v2.BidiReadObjectRequest(read_ranges=batch)
                         )
 
-                    # Receive responses
                     while pending_read_ids:
                         response = await self.read_obj_str.recv()
                         if response is None:
@@ -504,4 +493,4 @@ async def close(self):
 
     @property
     def is_stream_open(self) -> bool:
-        return self._is_stream_open
+        return self._is_stream_open

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

@@ -51,7 +51,7 @@ class _AsyncReadObjectStream(_AsyncAbstractObjectStream):
     :param generation_number: (Optional) If present, selects a specific revision of
                               this object.
 
-    :type read_handle: _storage_v2.BidiReadHandle
+    :type read_handle: bytes
     :param read_handle: (Optional) An existing handle for reading the object.
                         If provided, opening the bidi-gRPC connection will be faster.
     """
@@ -62,7 +62,7 @@ def __init__(
         bucket_name: str,
         object_name: str,
         generation_number: Optional[int] = None,
-        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
+        read_handle: Optional[bytes] = None,
     ) -> None:
         if client is None:
             raise ValueError("client must be provided")
@@ -77,7 +77,7 @@ def __init__(
             generation_number=generation_number,
         )
         self.client: AsyncGrpcClient.grpc_client = client
-        self.read_handle: Optional[_storage_v2.BidiReadHandle] = read_handle
+        self.read_handle: Optional[bytes] = read_handle
 
         self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}"
 
@@ -195,4 +195,4 @@ async def recv(self) -> _storage_v2.BidiReadObjectResponse:
 
     @property
     def is_stream_open(self) -> bool:
-        return self._is_stream_open
+        return self._is_stream_open

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

@@ -22,7 +22,6 @@
 
 """
 from typing import List, Optional, Tuple
-from . import _utils
 from google.cloud import _storage_v2
 from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
 from google.cloud.storage._experimental.asyncio.async_abstract_object_stream import (
@@ -60,7 +59,7 @@ class _AsyncWriteObjectStream(_AsyncAbstractObjectStream):
         same name already exists, it will be overwritten the moment
         `writer.open()` is called.
 
-    :type write_handle: _storage_v2.BidiWriteHandle
+    :type write_handle: bytes
     :param write_handle: (Optional) An existing handle for writing the object.
                         If provided, opening the bidi-gRPC connection will be faster.
     """
@@ -125,9 +124,6 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> 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_handle is not None and self.generation_number
-        )
         if self.generation_number is None or self.generation_number == 0:
             self.first_bidi_write_req = _storage_v2.BidiWriteObjectRequest(
                 write_object_spec=_storage_v2.WriteObjectSpec(
@@ -196,8 +192,9 @@ async def close(self) -> None:
 
     async def requests_done(self):
         """Signals that all requests have been sent."""
+
         await self.socket_like_rpc.send(None)
-        _utils.update_write_handle_if_exists(self, await self.socket_like_rpc.recv())
+        await self.socket_like_rpc.recv()
 
     async def send(
         self, bidi_write_object_request: _storage_v2.BidiWriteObjectRequest
@@ -239,3 +236,4 @@ 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/base_strategy.py

@@ -27,20 +27,16 @@ class _BaseResumptionStrategy(abc.ABC):
     """
 
     @abc.abstractmethod
-    def generate_requests(self, state: Any):
-        """Generates requests based on the current state as a generator.
+    def generate_requests(self, state: Any) -> Iterable[Any]:
+        """Generates the next batch of requests based on the current state.
 
         This method is called at the beginning of each retry attempt. It should
-        inspect the provided state object and yield request protos to send to
-        the server. For example, a read strategy would use this to implement
-        "Smarter Resumption" by creating smaller `ReadRange` requests for
-        partially downloaded ranges. For bidi-writes, it will set the
-        `write_offset` field to the persisted size received from the server
-        in the next request.
-
-        This is a generator that yields requests incrementally rather than
-        returning them all at once, allowing for better memory efficiency
-        and on-demand generation.
+        inspect the provided state object and generate the appropriate list of
+        request protos to send to the server. For example, a read strategy
+        would use this to implement "Smarter Resumption" by creating smaller
+        `ReadRange` requests for partially downloaded ranges. For bidi-writes,
+        it will set the `write_offset` field to the persisted size received
+        from the server in the next request.
 
         :type state: Any
         :param state: An object containing all the state needed for the

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

@@ -50,8 +50,8 @@ async def execute(self, initial_state: Any, retry_policy):
         state = initial_state
 
         async def attempt():
-            requests_generator = self._strategy.generate_requests(state)
-            stream = self._send_and_recv(requests_generator, state)
+            requests = self._strategy.generate_requests(state)
+            stream = self._send_and_recv(requests, state)
             try:
                 async for response in stream:
                     self._strategy.update_state_from_response(response, state)