diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py index 9244d9cf193e..3e6c2b533b1e 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py @@ -18,12 +18,12 @@ import time from typing import TYPE_CHECKING, Sequence -from google.api_core import retry as retries from grpc import StatusCode from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( _attempt_timeout_generator, + _rst_stream_aware_predicate, ) from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry from google.cloud.bigtable.data.exceptions import ( @@ -108,7 +108,7 @@ def __init__( else: self.request = query._to_pb(target) self.target = target - self._predicate = retries.if_exception_type(*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 self._operation_metric = metric diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py index c21b50064a6d..ceecf0eecaa1 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -20,9 +20,17 @@ import enum import time from collections import namedtuple -from typing import TYPE_CHECKING, List, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Callable, + List, + Sequence, + Tuple, + Union, +) from google.api_core import exceptions as core_exceptions +from google.api_core import retry as retries from google.api_core.retry import RetryFailureReason, exponential_sleep_generator from google.cloud.bigtable.data.exceptions import RetryExceptionGroup @@ -50,6 +58,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 treated 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( @@ -126,6 +143,35 @@ def _retry_exception_factory( return source_exc, cause_exc +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 + """ + # 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 + def rst_check(e): + return ( + 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( operation: float | TABLE_DEFAULT, attempt: float | None | TABLE_DEFAULT, diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 3c2f99283789..ce29226cac53 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -21,11 +21,13 @@ import time from typing import TYPE_CHECKING, Sequence -from google.api_core import retry as retries from grpc import StatusCode from google.cloud.bigtable.data._cross_sync import CrossSync -from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import ( + _attempt_timeout_generator, + _rst_stream_aware_predicate, +) from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry from google.cloud.bigtable.data.exceptions import ( InvalidChunk, @@ -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 = _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 self._operation_metric = metric diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index 815f26a367e3..7f7e45209313 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1510,49 +1510,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,), ), ], @@ -1588,6 +1595,7 @@ async def test_customizable_retryable_errors( expected_retryables, fn_name, fn_args, + is_read_rows_fn, is_stream, extra_retryables, ): @@ -1600,8 +1608,18 @@ 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._rst_stream_aware_predicate" + 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: @@ -1609,9 +1627,7 @@ async def test_customizable_retryable_errors( 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 diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 08d628f4aef7..899041e3055c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -1263,19 +1263,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,), ), ], @@ -1310,6 +1311,7 @@ def test_customizable_retryable_errors( expected_retryables, fn_name, fn_args, + is_read_rows_fn, is_stream, extra_retryables, ): @@ -1319,6 +1321,11 @@ 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._rst_stream_aware_predicate" + 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: @@ -1326,9 +1333,7 @@ def test_customizable_retryable_errors( 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) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py index 289b5ba7408b..ef65b5f30625 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py @@ -222,6 +222,48 @@ def test_get_timeouts_invalid(self, input_times): _helpers._align_timeouts(input_times[0], input_times[1]) +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, + ), + ], + ) + 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 + + class TestGetRetryableErrors: @pytest.mark.parametrize( "input_codes,input_table,expected",