Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

feat: Added rst_stream exception handling for ReadRows. #1298

Merged
gkevinzheng merged 4 commits into from
Feb 25, 2026
Merged

feat: Added rst_stream exception handling for ReadRows. #1298

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d32616b
feat: Added rst_stream exception handling for ReadRows.
gkevinzheng Feb 24, 2026
e915411
addressed Gemini comments
gkevinzheng Feb 24, 2026
b266a29
chore: remove 3.9 tests from kokoro (#1293)
daniel-sanche Feb 12, 2026
6c716b2
Renamed and cleaned up predicate
gkevinzheng Feb 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion google/cloud/bigtable/data/_async/_read_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
39 changes: 38 additions & 1 deletion google/cloud/bigtable/data/_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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 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(
Expand Down Expand Up @@ -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]:

@daniel-sanche daniel-sanche Feb 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we call this something more specific, since it's not really tied to read_rows? Maybe _rst_stream_aware_predicate?

"""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:
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

@daniel-sanche daniel-sanche Feb 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I find the logic a bit hard to follow, with the multiple predicate functions being passed around. Can we do something like this, to put all the cases in one place?

def _read_rows_predicate_with_exceptions(*exception_types: type[Exception]) -> Callable[[Exception], bool]:
    def predicate(exception: Exception)  -> bool:
        if isinstance(exception, exception_types):
           # handle standard retryable errors
           return True
        elif core_exceptions.ServiceUnavailable in exception_types and isinstance(exception, InternalServerError) and any(m in exception.message.lower() for m in _RETRYABLE_INTERNAL_ERROR_MESSAGES)):
            # special case: treat InternalServerError with specific message as ServiceUnavailable
            return True
        else:
            return False
    return predicate

or even this?

def _read_rows_predicate_with_exceptions(*exception_types: type[Exception]) -> Callable[[Exception], bool]:
    # predicate to check for retryable error types
    if_exception_type = retries.if_exception_type(*exception_types)
    # special case: treat InternalServerError with specific message as ServiceUnavailable
    rst_check = lambda e: core_exceptions.ServiceUnavailable in exception_types and isinstance(e, 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,
Expand Down
4 changes: 2 additions & 2 deletions google/cloud/bigtable/data/_sync_autogen/_read_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down
29 changes: 23 additions & 6 deletions tests/unit/data/_async/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,),
),
],
Expand Down Expand Up @@ -1426,6 +1433,7 @@ async def test_customizable_retryable_errors(
expected_retryables,
fn_name,
fn_args,
is_read_rows_fn,
is_stream,
extra_retryables,
):
Expand All @@ -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
Expand All @@ -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(
Expand Down
27 changes: 16 additions & 11 deletions tests/unit/data/_sync_autogen/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,),
),
],
Expand Down Expand Up @@ -1147,6 +1148,7 @@ def test_customizable_retryable_errors(
expected_retryables,
fn_name,
fn_args,
is_read_rows_fn,
is_stream,
extra_retryables,
):
Expand All @@ -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)
Expand Down
16 changes: 16 additions & 0 deletions tests/unit/data/test__helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading