fix: add warning to retry target to avoid incorrect use

ohmayr · ohmayr · commit 87fa112abe91 · 2023-10-26T21:37:39.000Z
diff --git a/google/api_core/operations_v1/operations_async_client.py b/google/api_core/operations_v1/operations_async_client.py
@@ -26,7 +26,7 @@
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import gapic_v1, page_iterator_async
-from google.api_core import retry as retries
+from google.api_core import retry_async as retries
 from google.api_core import timeout as timeouts
 from google.longrunning import operations_pb2
 from grpc import Compression
@@ -48,7 +48,7 @@ def __init__(self, channel, client_config=None):
         # Create the gRPC client stub with gRPC AsyncIO channel.
         self.operations_stub = operations_pb2.OperationsStub(channel)
 
-        default_retry = retries.Retry(
+        default_retry = retries.AsyncRetry(
             initial=0.1,  # seconds
             maximum=60.0,  # seconds
             multiplier=1.3,
diff --git a/google/api_core/retry.py b/google/api_core/retry.py
@@ -62,6 +62,7 @@ def check_if_exists():
 import random
 import sys
 import time
+import inspect
 from typing import Any, Callable, TypeVar, TYPE_CHECKING
 
 import requests.exceptions
@@ -201,7 +202,10 @@ def retry_target(
 
     for sleep in sleep_generator:
         try:
-            return target()
+            result = target()
+            if inspect.isawaitable(result):
+                raise exceptions.GoogleAPIError("Warning: Use google.api_core.retry_async.AsyncRetry for async calls!")
+            return result
 
         # pylint: disable=broad-except
         # This function explicitly must deal with broad exceptions.
diff --git a/tests/asyncio/test_retry_async.py b/tests/asyncio/test_retry_async.py
@@ -20,6 +20,7 @@
 
 from google.api_core import exceptions
 from google.api_core import retry_async
+from google.api_core import retry
 
 
 @mock.patch("asyncio.sleep", autospec=True)
@@ -95,6 +96,24 @@ async def test_retry_target_non_retryable_error(utcnow, sleep):
     assert exc_info.value == exception
     sleep.assert_not_called()
 
+@mock.patch("asyncio.sleep", autospec=True)
+@mock.patch(
+    "google.api_core.datetime_helpers.utcnow",
+    return_value=datetime.datetime.min,
+    autospec=True,
+)
+@pytest.mark.asyncio
+async def test_retry_target_warning_for_retry(utcnow, sleep):
+    predicate = retry.if_exception_type(exceptions.GoogleAPICallError)
+    exception = exceptions.GoogleAPICallError("Warning: Use google.api_core.retry_async.AsyncRetry for async calls!")
+    target = mock.Mock(side_effect=exception)
+
+    with pytest.raises(exceptions.GoogleAPICallError) as exc_info:
+        retry.retry_target(target, predicate, range(10), None)
+
+    assert exc_info.value == exception
+    sleep.assert_not_called()
+
 
 @mock.patch("asyncio.sleep", autospec=True)
 @mock.patch("google.api_core.datetime_helpers.utcnow", autospec=True)
