Renamed and cleaned up predicate

Author: gkevinzheng

google/cloud/bigtable/data/_async/_read_rows.py

@@ -29,10 +29,8 @@
 from google.cloud.bigtable.data.exceptions import _ResetRow
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._helpers import _read_rows_predicate_with_exceptions
+from google.cloud.bigtable.data._helpers import _rst_stream_aware_predicate
 
-from google.api_core import retry as retries
-from google.api_core import exceptions as core_exceptions
 from google.api_core.retry import exponential_sleep_generator
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
@@ -100,7 +98,7 @@ def __init__(
         else:
             self.request = query._to_pb(target)
         self.target = target
-        self._predicate = _read_rows_predicate_with_exceptions(*retryable_exceptions)
+        self._predicate = _rst_stream_aware_predicate(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
 

google/cloud/bigtable/data/_helpers.py

@@ -132,31 +132,31 @@ def _retry_exception_factory(
     return source_exc, cause_exc
 
 
-def _read_rows_predicate_with_exceptions(*exception_types: type[Exception]) -> Callable[[Exception], bool]:
-    """A custom retry predicate for ReadRows.
-    
+def _rst_stream_aware_predicate(
+    *exception_types: type[Exception],
+) -> Callable[[Exception], bool]:
+    """A custom retry predicate.
+
     This predicate treats Internal error messages with RST_STREAM errors as
     ServiceUnavailable errors and will retry them if the Unavailable exception is retryable.
 
     Args:
         exception_types: Exception types to be retried during operation
-    
+
     Returns:
         Callable[[Exception], bool]: A retry predicate that takes in an exception and
             returns whether or not that exception is retryable
     """
-    is_exception_type = retries.if_exception_type(*exception_types)
-    
-    def predicate(exception: Exception) -> bool:
-        return (isinstance(exception, core_exceptions.InternalServerError) and any(m in exception.message.lower() for m in _RETRYABLE_INTERNAL_ERROR_MESSAGES)) or is_exception_type(exception)
-
-    # Treating RST_STREAM internal errors as unavailable errors is only done if ServiceUnavailable is one of the
-    # given exception types. If InternalServerError is also a retryable exception, we don't necessarily need the
-    # custom predicate either.
-    if core_exceptions.ServiceUnavailable in exception_types and core_exceptions.InternalServerError not in exception_types:
-        return predicate
-
-    return is_exception_type
+    # predicate to check for retryable error types
+    if_exception_type = retries.if_exception_type(*exception_types)
+    # special case: treat InternalServerError with rst_stream error message as ServiceUnavailable
+    rst_check = (
+        lambda e: core_exceptions.ServiceUnavailable in exception_types
+        and isinstance(e, core_exceptions.InternalServerError)
+        and any(m in e.message.lower() for m in _RETRYABLE_INTERNAL_ERROR_MESSAGES)
+    )
+
+    return lambda e: if_exception_type(e) or rst_check(e)
 
 
 def _get_timeouts(

google/cloud/bigtable/data/_sync_autogen/_read_rows.py

@@ -29,7 +29,7 @@
 from google.cloud.bigtable.data.exceptions import _ResetRow
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._helpers import _read_rows_predicate_with_exceptions
+from google.cloud.bigtable.data._helpers import _rst_stream_aware_predicate
 from google.api_core.retry import exponential_sleep_generator
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -88,7 +88,7 @@ def __init__(
         else:
             self.request = query._to_pb(target)
         self.target = target
-        self._predicate = _read_rows_predicate_with_exceptions(*retryable_exceptions)
+        self._predicate = _rst_stream_aware_predicate(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
 

google/cloud/bigtable/data/exceptions.py

@@ -142,7 +142,7 @@ def __repr__(self):
 
 
 # TODO: When working on mutations batcher, rework exception handling to guarantee that
-# MutationsExceptionGroup only stores FailedMutationEntryErrors. 
+# MutationsExceptionGroup only stores FailedMutationEntryErrors.
 class MutationsExceptionGroup(_BigtableExceptionGroup):
     """
     Represents one or more exceptions that occur during a bulk mutation operation

tests/unit/data/_async/test_client.py

@@ -1454,7 +1454,7 @@ async def test_customizable_retryable_errors(
         # Read Rows has its own custom predicate builder that also takes in
         # a list of exceptions
         if is_read_rows_fn:
-            predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._read_rows_predicate_with_exceptions"
+            predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._rst_stream_aware_predicate"
         else:
             predicate_builder = "google.api_core.retry.if_exception_type"
 

tests/unit/data/_sync_autogen/test_client.py

@@ -1160,7 +1160,7 @@ def test_customizable_retryable_errors(
         retry_fn = f"CrossSync._Sync_Impl.{retry_fn}"
         subpackage = "_sync_autogen"
         if is_read_rows_fn:
-            predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._read_rows_predicate_with_exceptions"
+            predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._rst_stream_aware_predicate"
         else:
             predicate_builder = "google.api_core.retry.if_exception_type"
         with mock.patch(

tests/unit/data/test__helpers.py

@@ -222,19 +222,45 @@ def test_get_timeouts_invalid(self, input_times):
             _helpers._align_timeouts(input_times[0], input_times[1])
 
 
-class TestReadRowsPredicateWithException:
+class TestRstStreamAwarePredicate:
     @pytest.mark.parametrize(
         "retryable_exceptions,exception,expected_is_retryable",
         [
-            ([core_exceptions.Aborted, core_exceptions.InternalServerError], core_exceptions.InternalServerError("Sorry"), True),
-            ([core_exceptions.Aborted, core_exceptions.InternalServerError], core_exceptions.DataLoss("Sorry"), False),
-            ([core_exceptions.ServiceUnavailable, core_exceptions.Aborted], core_exceptions.InternalServerError("Sorry"), False),
-            ([core_exceptions.ServiceUnavailable, core_exceptions.Aborted], core_exceptions.InternalServerError(_helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0]), True),
-            ([core_exceptions.InternalServerError, core_exceptions.Aborted], core_exceptions.InternalServerError(_helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0]), True),
-        ]
+            (
+                [core_exceptions.Aborted, core_exceptions.InternalServerError],
+                core_exceptions.InternalServerError("Sorry"),
+                True,
+            ),
+            (
+                [core_exceptions.Aborted, core_exceptions.InternalServerError],
+                core_exceptions.DataLoss("Sorry"),
+                False,
+            ),
+            (
+                [core_exceptions.ServiceUnavailable, core_exceptions.Aborted],
+                core_exceptions.InternalServerError("Sorry"),
+                False,
+            ),
+            (
+                [core_exceptions.ServiceUnavailable, core_exceptions.Aborted],
+                core_exceptions.InternalServerError(
+                    _helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0]
+                ),
+                True,
+            ),
+            (
+                [core_exceptions.InternalServerError, core_exceptions.Aborted],
+                core_exceptions.InternalServerError(
+                    _helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0]
+                ),
+                True,
+            ),
+        ],
     )
-    def test_ctor_retryable_exceptions_predicate(self, retryable_exceptions, exception, expected_is_retryable):
-        predicate = _helpers._read_rows_predicate_with_exceptions(*retryable_exceptions)
+    def test_rst_stream_aware_predicate(
+        self, retryable_exceptions, exception, expected_is_retryable
+    ):
+        predicate = _helpers._rst_stream_aware_predicate(*retryable_exceptions)
         assert predicate(exception) is expected_is_retryable