feat: Added rst_stream exception handling for ReadRows.

Author: gkevinzheng

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

@@ -29,8 +29,10 @@
 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.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
@@ -98,7 +100,7 @@ def __init__(
         else:
             self.request = query._to_pb(target)
         self.target = target
-        self._predicate = retries.if_exception_type(*retryable_exceptions)
+        self._predicate = _read_rows_predicate_with_exceptions(*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

@@ -16,13 +16,14 @@
 """
 from __future__ import annotations
 
-from typing import Sequence, List, Tuple, TYPE_CHECKING, Union
+from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union
 import time
 import enum
 from collections import namedtuple
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 
 from google.api_core import exceptions as core_exceptions
+from google.api_core import retry as retries
 from google.api_core.retry import RetryFailureReason
 from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
 
@@ -47,6 +48,15 @@
 # used by every data client as a default project name for testing on Bigtable emulator.
 _DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
 
+# Internal error messages that can be retried during ReadRows. Internal error messages with this error
+# text should be streated as Unavailable error messages with the same error text, and will therefore be
+# treated as Unavailable errors rather than Internal errors.
+_RETRYABLE_INTERNAL_ERROR_MESSAGES = (
+    "rst_stream",
+    "rst stream",
+    "received unexpected eos on data frame from server",
+)
+
 # used to identify an active bigtable resource that needs to be warmed through PingAndWarm
 # each instance/app_profile_id pair needs to be individually tracked
 _WarmedInstanceKey = namedtuple(
@@ -122,6 +132,33 @@ 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.
+    
+    This predicate treats Internal error messages with RST_STREAM errors as
+    ServiceUnavailable errors and will retry them if the Unavailable exception is retryable.
+
+    Args:
+        retryable_exceptions: tuple of 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 exception.message.lower() 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
+
+
 def _get_timeouts(
     operation: float | TABLE_DEFAULT,
     attempt: float | None | TABLE_DEFAULT,

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.api_core import retry as retries
+from google.cloud.bigtable.data._helpers import _read_rows_predicate_with_exceptions
 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 = retries.if_exception_type(*retryable_exceptions)
+        self._predicate = _read_rows_predicate_with_exceptions(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
 

tests/unit/data/_async/test_client.py

@@ -1349,49 +1349,56 @@ def test_table_ctor_sync(self):
     @CrossSync.pytest
     # iterate over all retryable rpcs
     @pytest.mark.parametrize(
-        "fn_name,fn_args,is_stream,extra_retryables",
+        "fn_name,fn_args,is_read_rows_fn,is_stream,extra_retryables",
         [
             (
                 "read_rows_stream",
                 (ReadRowsQuery(),),
                 True,
+                True,
                 (),
             ),
             (
                 "read_rows",
                 (ReadRowsQuery(),),
                 True,
+                True,
                 (),
             ),
             (
                 "read_row",
                 (b"row_key",),
                 True,
+                True,
                 (),
             ),
             (
                 "read_rows_sharded",
                 ([ReadRowsQuery()],),
                 True,
+                True,
                 (),
             ),
             (
                 "row_exists",
                 (b"row_key",),
                 True,
+                True,
                 (),
             ),
-            ("sample_row_keys", (), False, ()),
+            ("sample_row_keys", (), False, False, ()),
             (
                 "mutate_row",
                 (b"row_key", [DeleteAllFromRow()]),
                 False,
+                False,
                 (),
             ),
             (
                 "bulk_mutate_rows",
                 ([mutations.RowMutationEntry(b"key", [DeleteAllFromRow()])],),
                 False,
+                False,
                 (_MutateRowsIncomplete,),
             ),
         ],
@@ -1426,6 +1433,7 @@ async def test_customizable_retryable_errors(
         expected_retryables,
         fn_name,
         fn_args,
+        is_read_rows_fn,
         is_stream,
         extra_retryables,
     ):
@@ -1438,18 +1446,26 @@ async def test_customizable_retryable_errors(
             retry_fn += "_stream"
         if CrossSync.is_async:
             retry_fn = f"CrossSync.{retry_fn}"
+            subpackage = "_async"
         else:
             retry_fn = f"CrossSync._Sync_Impl.{retry_fn}"
+            subpackage = "_sync_autogen"
+
+        # 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"
+        else:
+            predicate_builder = "google.api_core.retry.if_exception_type"
+
         with mock.patch(
             f"google.cloud.bigtable.data._cross_sync.{retry_fn}"
         ) as retry_fn_mock:
             async with self._make_client() as client:
                 table = client.get_table("instance-id", "table-id")
                 expected_predicate = expected_retryables.__contains__
                 retry_fn_mock.side_effect = RuntimeError("stop early")
-                with mock.patch(
-                    "google.api_core.retry.if_exception_type"
-                ) as predicate_builder_mock:
+                with mock.patch(predicate_builder) as predicate_builder_mock:
                     predicate_builder_mock.return_value = expected_predicate
                     with pytest.raises(Exception):
                         # we expect an exception from attempting to call the mock
@@ -1460,7 +1476,8 @@ async def test_customizable_retryable_errors(
                         *expected_retryables, *extra_retryables
                     )
                     retry_call_args = retry_fn_mock.call_args_list[0].args
-                    # output of if_exception_type should be sent in to retry constructor
+
+                    # output of the predicate builder should be sent in to retry constructor
                     assert retry_call_args[1] is expected_predicate
 
     @pytest.mark.parametrize(

tests/unit/data/_sync_autogen/test_client.py

@@ -1101,19 +1101,20 @@ def test_ctor_invalid_timeout_values(self):
         client.close()
 
     @pytest.mark.parametrize(
-        "fn_name,fn_args,is_stream,extra_retryables",
+        "fn_name,fn_args,is_read_rows_fn,is_stream,extra_retryables",
         [
-            ("read_rows_stream", (ReadRowsQuery(),), True, ()),
-            ("read_rows", (ReadRowsQuery(),), True, ()),
-            ("read_row", (b"row_key",), True, ()),
-            ("read_rows_sharded", ([ReadRowsQuery()],), True, ()),
-            ("row_exists", (b"row_key",), True, ()),
-            ("sample_row_keys", (), False, ()),
-            ("mutate_row", (b"row_key", [DeleteAllFromRow()]), False, ()),
+            ("read_rows_stream", (ReadRowsQuery(),), True, True, ()),
+            ("read_rows", (ReadRowsQuery(),), True, True, ()),
+            ("read_row", (b"row_key",), True, True, ()),
+            ("read_rows_sharded", ([ReadRowsQuery()],), True, True, ()),
+            ("row_exists", (b"row_key",), True, True, ()),
+            ("sample_row_keys", (), False, False, ()),
+            ("mutate_row", (b"row_key", [DeleteAllFromRow()]), False, False, ()),
             (
                 "bulk_mutate_rows",
                 ([mutations.RowMutationEntry(b"key", [DeleteAllFromRow()])],),
                 False,
+                False,
                 (_MutateRowsIncomplete,),
             ),
         ],
@@ -1147,6 +1148,7 @@ def test_customizable_retryable_errors(
         expected_retryables,
         fn_name,
         fn_args,
+        is_read_rows_fn,
         is_stream,
         extra_retryables,
     ):
@@ -1156,16 +1158,19 @@ def test_customizable_retryable_errors(
         if is_stream:
             retry_fn += "_stream"
         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"
+        else:
+            predicate_builder = "google.api_core.retry.if_exception_type"
         with mock.patch(
             f"google.cloud.bigtable.data._cross_sync.{retry_fn}"
         ) as retry_fn_mock:
             with self._make_client() as client:
                 table = client.get_table("instance-id", "table-id")
                 expected_predicate = expected_retryables.__contains__
                 retry_fn_mock.side_effect = RuntimeError("stop early")
-                with mock.patch(
-                    "google.api_core.retry.if_exception_type"
-                ) as predicate_builder_mock:
+                with mock.patch(predicate_builder) as predicate_builder_mock:
                     predicate_builder_mock.return_value = expected_predicate
                     with pytest.raises(Exception):
                         test_fn = table.__getattribute__(fn_name)

tests/unit/data/test__helpers.py

@@ -222,6 +222,22 @@ def test_get_timeouts_invalid(self, input_times):
             _helpers._align_timeouts(input_times[0], input_times[1])
 
 
+class TestReadRowsPredicateWithException:
+    @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),
+        ]
+    )
+    def test_ctor_retryable_exceptions_predicate(self, retryable_exceptions, exception, expected_is_retryable):
+        predicate = _helpers._read_rows_predicate_with_exceptions(*retryable_exceptions)
+        assert predicate(exception) is expected_is_retryable
+
+
 class TestGetRetryableErrors:
     @pytest.mark.parametrize(
         "input_codes,input_table,expected",