adding license

Pulkit0110 · Pulkit0110 · commit 21ef7810d7ab · 2025-12-01T04:40:11.000Z
diff --git a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py
@@ -1,3 +1,17 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import abc
 from typing import Any, Iterable
 
diff --git a/google/cloud/storage/_experimental/asyncio/retry/bidi_stream_retry_manager.py b/google/cloud/storage/_experimental/asyncio/retry/bidi_stream_retry_manager.py
@@ -1,7 +1,23 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import asyncio
+import time
 from typing import Any, AsyncIterator, Callable
 
 from google.api_core import exceptions
+from google.api_core.retry.retry_base import exponential_sleep_generator
 from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
     _BaseResumptionStrategy,
 )
@@ -14,54 +30,54 @@ def __init__(
         self,
         strategy: _BaseResumptionStrategy,
         stream_opener: Callable[..., AsyncIterator[Any]],
-        retry_policy,
     ):
-        """Initializes the retry manager."""
+        """Initializes the retry manager.
+
+        Args:
+            strategy: The strategy for managing the state of a specific
+                bidi operation (e.g., reads or writes).
+            stream_opener: An async callable that opens a new gRPC stream.
+        """
         self._strategy = strategy
         self._stream_opener = stream_opener
-        self._retry_policy = retry_policy
 
-    async def execute(self, initial_state: Any):
+    async def execute(self, initial_state: Any, retry_policy):
         """
         Executes the bidi operation with the configured retry policy.
 
         This method implements a manual retry loop that provides the necessary
-        control points to manage state between attempts, which is not possible
-        with a simple retry decorator.
+        control points to manage state between attempts.
+
+        Args:
+            initial_state: An object containing all state for the operation.
+            retry_policy: The `google.api_core.retry.AsyncRetry` object to
+                govern the retry behavior for this specific operation.
         """
         state = initial_state
-        retry_policy = self._retry_policy
+
+        deadline = time.monotonic() + retry_policy._deadline if retry_policy._deadline else 0
+
+        sleep_generator = exponential_sleep_generator(
+            retry_policy._initial, retry_policy._maximum, retry_policy._multiplier
+        )
 
         while True:
             try:
-                # 1. Generate requests based on the current state.
                 requests = self._strategy.generate_requests(state)
-
-                # 2. Open and consume the stream.
                 stream = self._stream_opener(requests, state)
                 async for response in stream:
                     self._strategy.update_state_from_response(response, state)
-
-                # 3. If the stream completes without error, exit the loop.
                 return
-
             except Exception as e:
-                # 4. If an error occurs, check if it's retriable.
-                if not retry_policy.predicate(e):
-                    # If not retriable, fail fast.
+                if not retry_policy._predicate(e):
                     raise
 
-                # 5. If retriable, allow the strategy to recover state.
-                #    This is where routing tokens are extracted or QueryWriteStatus is called.
-                try:
-                    await self._strategy.recover_state_on_failure(e, state)
-                except Exception as recovery_exc:
-                    # If state recovery itself fails, we must abort.
+                await self._strategy.recover_state_on_failure(e, state)
+
+                sleep = next(sleep_generator)
+                if deadline is not None and time.monotonic() + sleep > deadline:
                     raise exceptions.RetryError(
-                        "Failed to recover state after a transient error.",
-                        cause=recovery_exc,
-                    ) from recovery_exc
+                        f"Deadline of {retry_policy._deadline}s exceeded", cause=e
+                    ) from e
 
-                # 6. Use the policy to sleep and check for deadline expiration.
-                #    This will raise a RetryError if the deadline is exceeded.
-                await asyncio.sleep(await retry_policy.sleep(e))
+                await asyncio.sleep(sleep)
diff --git a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py
@@ -1,3 +1,17 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, List, IO
 
 from google.cloud import _storage_v2 as storage_v2
diff --git a/tests/unit/asyncio/retry/test_bidi_stream_retry_manager.py b/tests/unit/asyncio/retry/test_bidi_stream_retry_manager.py
@@ -1,49 +1,52 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
 import unittest
 from unittest import mock
 
 import pytest
 from google.api_core import exceptions
-from google.api_core.retry.retry_streaming_async import AsyncStreamingRetry
+from google.api_core.retry_async import AsyncRetry
 
-from google.cloud.storage._experimental.asyncio.retry import manager
-from google.cloud.storage._experimental.asyncio.retry import strategy
+from google.cloud.storage._experimental.asyncio.retry import bidi_stream_retry_manager as manager
+from google.cloud.storage._experimental.asyncio.retry import base_strategy
 
 
 def _is_retriable(exc):
     return isinstance(exc, exceptions.ServiceUnavailable)
 
-
-DEFAULT_TEST_RETRY = AsyncStreamingRetry(predicate=_is_retriable, deadline=1)
+DEFAULT_TEST_RETRY = AsyncRetry(predicate=_is_retriable, deadline=1)
 
 
 class TestBidiStreamRetryManager(unittest.IsolatedAsyncioTestCase):
     async def test_execute_success_on_first_try(self):
-        """Verify the manager correctly handles a stream that succeeds immediately."""
-        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
-
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
         async def mock_stream_opener(*args, **kwargs):
             yield "response_1"
 
         retry_manager = manager._BidiStreamRetryManager(
-            strategy=mock_strategy,
-            stream_opener=mock_stream_opener,
-            retry_policy=DEFAULT_TEST_RETRY,
+            strategy=mock_strategy, stream_opener=mock_stream_opener
         )
-
-        await retry_manager.execute(initial_state={})
-
+        await retry_manager.execute(initial_state={}, retry_policy=DEFAULT_TEST_RETRY)
         mock_strategy.generate_requests.assert_called_once()
-        mock_strategy.update_state_from_response.assert_called_once_with(
-            "response_1", {}
-        )
+        mock_strategy.update_state_from_response.assert_called_once_with("response_1", {})
         mock_strategy.recover_state_on_failure.assert_not_called()
 
     async def test_execute_retries_and_succeeds(self):
-        """Verify the manager retries on a transient error and then succeeds."""
-        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
-
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
         attempt_count = 0
-
         async def mock_stream_opener(*args, **kwargs):
             nonlocal attempt_count
             attempt_count += 1
@@ -53,59 +56,43 @@ async def mock_stream_opener(*args, **kwargs):
                 yield "response_2"
 
         retry_manager = manager._BidiStreamRetryManager(
-            strategy=mock_strategy,
-            stream_opener=mock_stream_opener,
-            retry_policy=AsyncStreamingRetry(predicate=_is_retriable, initial=0.01),
+            strategy=mock_strategy, stream_opener=mock_stream_opener
         )
-
-        await retry_manager.execute(initial_state={})
+        retry_policy = AsyncRetry(predicate=_is_retriable, initial=0.01)
+        await retry_manager.execute(initial_state={}, retry_policy=retry_policy)
 
         self.assertEqual(attempt_count, 2)
         self.assertEqual(mock_strategy.generate_requests.call_count, 2)
         mock_strategy.recover_state_on_failure.assert_called_once()
-        mock_strategy.update_state_from_response.assert_called_once_with(
-            "response_2", {}
-        )
+        mock_strategy.update_state_from_response.assert_called_once_with("response_2", {})
 
     async def test_execute_fails_after_deadline_exceeded(self):
-        """Verify the manager raises RetryError if the deadline is exceeded."""
-        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
-
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
         async def mock_stream_opener(*args, **kwargs):
+            if False:
+                yield
             raise exceptions.ServiceUnavailable("Service is always down")
 
-        # Use a very short deadline to make the test fast.
-        fast_retry = AsyncStreamingRetry(
-            predicate=_is_retriable, deadline=0.1, initial=0.05
-        )
+        fast_retry = AsyncRetry(predicate=_is_retriable, deadline=0.01, initial=0.02)
         retry_manager = manager._BidiStreamRetryManager(
-            strategy=mock_strategy,
-            stream_opener=mock_stream_opener,
-            retry_policy=fast_retry,
+            strategy=mock_strategy, stream_opener=mock_stream_opener
         )
+        with pytest.raises(exceptions.RetryError, match="Deadline of 0.01s exceeded"):
+            await retry_manager.execute(initial_state={}, retry_policy=fast_retry)
 
-        with pytest.raises(exceptions.RetryError, match="Deadline of 0.1s exceeded"):
-            await retry_manager.execute(initial_state={})
-
-        # Verify it attempted to recover state after each failure.
-        self.assertGreater(mock_strategy.recover_state_on_failure.call_count, 1)
+        self.assertGreater(mock_strategy.recover_state_on_failure.call_count, 0)
 
     async def test_execute_fails_immediately_on_non_retriable_error(self):
-        """Verify the manager aborts immediately on a non-retriable error."""
-        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
-
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
         async def mock_stream_opener(*args, **kwargs):
+            if False:
+                yield
             raise exceptions.PermissionDenied("Auth error")
 
         retry_manager = manager._BidiStreamRetryManager(
-            strategy=mock_strategy,
-            stream_opener=mock_stream_opener,
-            retry_policy=DEFAULT_TEST_RETRY,
+            strategy=mock_strategy, stream_opener=mock_stream_opener
         )
-
         with pytest.raises(exceptions.PermissionDenied):
-            await retry_manager.execute(initial_state={})
+            await retry_manager.execute(initial_state={}, retry_policy=DEFAULT_TEST_RETRY)
 
-        # Verify that it did not try to recover or update state.
         mock_strategy.recover_state_on_failure.assert_not_called()
-        mock_strategy.update_state_from_response.assert_not_called()
