adding unit tests

Author: Pulkit0110

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

@@ -23,12 +23,12 @@
 if you want to use these Rapid Storage APIs.
 
 """
-from io import BufferedReader, BytesIO
+from io import BufferedReader
 import asyncio
 import io
-from typing import List, Optional, Tuple, Union
+import logging
+from typing import List, Optional, Tuple
 
-from google_crc32c import Checksum
 from google.api_core import exceptions
 from google.api_core.retry_async import AsyncRetry
 from google.rpc import status_pb2
@@ -58,13 +58,12 @@
 _BIDI_WRITE_REDIRECTED_TYPE_URL = (
     "type.googleapis.com/google.storage.v2.BidiWriteObjectRedirectedError"
 )
+logger = logging.getLogger(__name__)
 
 
 def _is_write_retryable(exc):
     """Predicate to determine if a write operation should be retried."""
 
-    print("In _is_write_retryable method, exception:", exc)
-
     if isinstance(
         exc,
         (
@@ -74,6 +73,7 @@ def _is_write_retryable(exc):
             exceptions.TooManyRequests,
         ),
     ):
+        logger.info(f"Retryable write exception encountered: {exc}")
         return True
 
     grpc_error = None
@@ -97,6 +97,7 @@ 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.")
                 return False
     return False
 
@@ -201,16 +202,6 @@ def __init__(
         self._routing_token: Optional[str] = None
         self.object_resource: Optional[_storage_v2.Object] = None
 
-    def _stream_opener(self, write_handle=None):
-        """Helper to create a new _AsyncWriteObjectStream."""
-        return _AsyncWriteObjectStream(
-            client=self.client,
-            bucket_name=self.bucket_name,
-            object_name=self.object_name,
-            generation_number=self.generation,
-            write_handle=write_handle if write_handle else self.write_handle,
-        )
-
     async def state_lookup(self) -> int:
         """Returns the persisted_size
 
@@ -244,6 +235,8 @@ def _on_open_error(self, exc):
                 self._routing_token = grpc_error.routing_token
                 if grpc_error.write_handle:
                     self.write_handle = grpc_error.write_handle
+                if grpc_error.generation:
+                    self.generation = grpc_error.generation
                 return
 
             if hasattr(grpc_error, "trailing_metadata"):
@@ -272,9 +265,13 @@ def _on_open_error(self, exc):
                                     self._routing_token = redirect_proto.routing_token
                                 if redirect_proto.write_handle:
                                     self.write_handle = redirect_proto.write_handle
+                                if redirect_proto.generation:
+                                    self.generation = redirect_proto.generation
                                 break
                     except Exception:
-                        # Could not parse the error, ignore
+                        logger.error(
+                            "Error unpacking redirect details from gRPC error."
+                        )
                         pass
 
     async def open(
@@ -302,12 +299,16 @@ def combined_on_error(exc):
                 if original_on_error:
                     original_on_error(exc)
 
-            retry_policy = retry_policy.with_predicate(
-                _is_write_retryable
-            ).with_on_error(combined_on_error)
+            retry_policy = AsyncRetry(
+                predicate=_is_write_retryable,
+                initial=retry_policy._initial,
+                maximum=retry_policy._maximum,
+                multiplier=retry_policy._multiplier,
+                deadline=retry_policy._deadline,
+                on_error=combined_on_error,
+            )
 
         async def _do_open():
-            print("In _do_open method")
             current_metadata = list(metadata) if metadata else []
 
             # Cleanup stream from previous failed attempt, if any.
@@ -326,15 +327,14 @@ async def _do_open():
                 object_name=self.object_name,
                 generation_number=self.generation,
                 write_handle=self.write_handle,
+                routing_token=self._routing_token,
             )
 
             if self._routing_token:
                 current_metadata.append(
                     ("x-goog-request-params", f"routing_token={self._routing_token}")
                 )
-                self._routing_token = None
 
-            print("Current metadata in _do_open:", current_metadata)
             await self.write_obj_stream.open(
                 metadata=current_metadata if metadata else None
             )
@@ -347,11 +347,10 @@ async def _do_open():
                 self.persisted_size = self.write_obj_stream.persisted_size
 
             self._is_stream_open = True
+            self._routing_token = None
 
-        print("In open method, before retry_policy call")
         await retry_policy(_do_open)()
 
-
     async def append(
         self,
         data: bytes,
@@ -387,47 +386,62 @@ async def append(
         if retry_policy is None:
             retry_policy = AsyncRetry(predicate=_is_write_retryable)
 
+        strategy = _WriteResumptionStrategy()
         buffer = io.BytesIO(data)
-        target_persisted_size = self.persisted_size + len(data)
         attempt_count = 0
 
-        print("In append method")
-
-        def send_and_recv_generator(requests: List[BidiWriteObjectRequest], state: dict[str, _WriteState], metadata: Optional[List[Tuple[str, str]]] = None):
+        def send_and_recv_generator(
+            requests: List[BidiWriteObjectRequest],
+            state: dict[str, _WriteState],
+            metadata: Optional[List[Tuple[str, str]]] = None,
+        ):
             async def generator():
-                print("In send_and_recv_generator")
                 nonlocal attempt_count
+                nonlocal requests
                 attempt_count += 1
                 resp = None
                 async with self._lock:
                     write_state = state["write_state"]
                     # If this is a retry or redirect, we must re-open the stream
                     if attempt_count > 1 or write_state.routing_token:
-                        print("Re-opening the stream inside send_and_recv_generator with attempt_count:", attempt_count)
-                        if self.write_obj_stream and self.write_obj_stream.is_stream_open:
+                        logger.info(
+                            f"Re-opening the stream with attempt_count: {attempt_count}"
+                        )
+                        if (
+                            self.write_obj_stream
+                            and self.write_obj_stream.is_stream_open
+                        ):
                             await self.write_obj_stream.close()
 
-                        self.write_obj_stream = self._stream_opener(write_handle=write_state.write_handle)
                         current_metadata = list(metadata) if metadata else []
                         if write_state.routing_token:
-                            current_metadata.append(("x-goog-request-params", f"routing_token={write_state.routing_token}"))
-                        await self.write_obj_stream.open(metadata=current_metadata if current_metadata else None)
+                            current_metadata.append(
+                                (
+                                    "x-goog-request-params",
+                                    f"routing_token={write_state.routing_token}",
+                                )
+                            )
+                            self._routing_token = write_state.routing_token
+
+                        self._is_stream_open = False
+                        await self.open(metadata=current_metadata)
 
-                        self._is_stream_open = True
                         write_state.persisted_size = self.persisted_size
                         write_state.write_handle = self.write_handle
+                        write_state.routing_token = None
 
-                    print("Sending requests in send_and_recv_generator")
-                    # req_iter = iter(requests)
+                        write_state.user_buffer.seek(write_state.persisted_size)
+                        write_state.bytes_sent = write_state.persisted_size
+                        write_state.bytes_since_last_flush = 0
 
-                    print("Starting to send requests")
+                        requests = strategy.generate_requests(state)
+
+                    num_requests = len(requests)
                     for i, chunk_req in enumerate(requests):
-                        if i == len(requests) - 1:
+                        if i == num_requests - 1:
                             chunk_req.state_lookup = True
-                        print("Sending chunk request")
+                            chunk_req.flush = True
                         await self.write_obj_stream.send(chunk_req)
-                        print("Waiting to receive response")
-                        print("Current persisted_size:", state["write_state"].persisted_size, "Target persisted_size:", target_persisted_size)
 
                     resp = await self.write_obj_stream.recv()
                     if resp:
@@ -437,38 +451,28 @@ async def generator():
                         if resp.write_handle:
                             self.write_handle = resp.write_handle
                             state["write_state"].write_handle = resp.write_handle
-                        print("Received response in send_and_recv_generator", resp)
 
                 yield resp
 
-                    # while state["write_state"].persisted_size < target_persisted_size:
-                    #     print("Waiting to receive response")
-                    #     print("Current persisted_size:", state["write_state"].persisted_size, "Target persisted_size:", target_persisted_size)
-                    #     resp = await self.write_obj_stream.recv()
-                    #     print("Received response in send_and_recv_generator", resp)
-                    #     if resp is None:
-                    #         break
-                    #     yield resp
             return generator()
 
         # State initialization
-        spec = _storage_v2.AppendObjectSpec(
-            bucket=f"projects/_/buckets/{self.bucket_name}", object=self.object_name, generation=self.generation
-        )
-        write_state = _WriteState(spec, _MAX_CHUNK_SIZE_BYTES, buffer)
+        write_state = _WriteState(_MAX_CHUNK_SIZE_BYTES, buffer, self.flush_interval)
         write_state.write_handle = self.write_handle
         write_state.persisted_size = self.persisted_size
         write_state.bytes_sent = self.persisted_size
+        write_state.bytes_since_last_flush = self.bytes_appended_since_last_flush
 
-        print("Before creating retry manager")
-        retry_manager = _BidiStreamRetryManager(_WriteResumptionStrategy(),
-                                                lambda r, s: send_and_recv_generator(r, s, metadata))
+        retry_manager = _BidiStreamRetryManager(
+            _WriteResumptionStrategy(),
+            lambda r, s: send_and_recv_generator(r, s, metadata),
+        )
         await retry_manager.execute({"write_state": write_state}, retry_policy)
 
         # Sync local markers
         self.write_obj_stream.persisted_size = write_state.persisted_size
         self.write_obj_stream.write_handle = write_state.write_handle
-
+        self.bytes_appended_since_last_flush = write_state.bytes_since_last_flush
 
     async def simple_flush(self) -> None:
         """Flushes the data to the server.
@@ -561,16 +565,10 @@ async def finalize(self) -> _storage_v2.Object:
         if not self._is_stream_open:
             raise ValueError("Stream is not open. Call open() before finalize().")
 
-        print("In finalize method")
-
-        # async with self._lock:
-        print("Sending finish_write request")
         await self.write_obj_stream.send(
             _storage_v2.BidiWriteObjectRequest(finish_write=True)
         )
-        print("Waiting to receive response for finalize")
         response = await self.write_obj_stream.recv()
-        print("Received response for finalize:")
         self.object_resource = response.resource
         self.persisted_size = self.object_resource.size
         await self.write_obj_stream.close()

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

@@ -62,6 +62,7 @@ def __init__(
         object_name: str,
         generation_number: Optional[int] = None,  # None means new object
         write_handle: Optional[bytes] = None,
+        routing_token: Optional[str] = None,
     ) -> None:
         if client is None:
             raise ValueError("client must be provided")
@@ -77,6 +78,7 @@ def __init__(
         )
         self.client: AsyncGrpcClient.grpc_client = client
         self.write_handle: Optional[bytes] = write_handle
+        self.routing_token: Optional[str] = routing_token
 
         self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}"
 
@@ -119,6 +121,7 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
                     object=self.object_name,
                     generation=self.generation_number,
                     write_handle=write_handle,
+                    routing_token=self.routing_token if self.routing_token else None,
                 ),
             )
 
@@ -134,37 +137,44 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None:
         current_metadata = other_metadata
         current_metadata.append(("x-goog-request-params", ",".join(request_params)))
 
-        print("Before sending first_bidi_write_req in open:", self.first_bidi_write_req)
-
         self.socket_like_rpc = AsyncBidiRpc(
-            self.rpc, initial_request=self.first_bidi_write_req, metadata=current_metadata
+            self.rpc,
+            initial_request=self.first_bidi_write_req,
+            metadata=current_metadata,
         )
 
         await self.socket_like_rpc.open()  # this is actually 1 send
         response = await self.socket_like_rpc.recv()
-        print("Received response on open")
         self._is_stream_open = True
 
-        if not response.resource:
-            raise ValueError(
-                "Failed to obtain object resource after opening the stream"
-            )
-        if not response.resource.generation:
-            raise ValueError(
-                "Failed to obtain object generation after opening the stream"
-            )
-
-        if not response.write_handle:
-            raise ValueError("Failed to obtain write_handle after opening the stream")
+        if response.persisted_size >= 0:
+            self.persisted_size = response.persisted_size
 
-        if not response.resource.size:
-            # Appending to a 0 byte appendable object.
-            self.persisted_size = 0
-        else:
-            self.persisted_size = response.resource.size
+            if response.write_handle:
+                self.write_handle = response.write_handle
+            # return
+
+        # if not response.resource:
+        #     raise ValueError(
+        #         "Failed to obtain object resource after opening the stream"
+        #     )
+        # if not response.resource.generation:
+        #     raise ValueError(
+        #         "Failed to obtain object generation after opening the stream"
+        #     )
+
+        # if not response.write_handle:
+        #     raise ValueError("Failed to obtain write_handle after opening the stream")
+
+        if response.resource:
+            if not response.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
+            self.generation_number = response.resource.generation
+            self.write_handle = response.write_handle
 
     async def close(self) -> None:
         """Closes the bidi-gRPC connection."""

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

@@ -49,14 +49,10 @@ async def execute(self, initial_state: Any, retry_policy):
         """
         state = initial_state
 
-        print("Starting retry manager execute")
-
         async def attempt():
             requests = self._strategy.generate_requests(state)
-            print("Generated requests:", len(requests))
             stream = self._send_and_recv(requests, state)
             try:
-                print("Starting to receive responses in attempt")
                 async for response in stream:
                     self._strategy.update_state_from_response(response, state)
                 return