Skip to content

fix(bigtable): Added rst_stream exception handling for ReadRows. #18197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
daniel-sanche wants to merge 2 commits into from
Draft

fix(bigtable): Added rst_stream exception handling for ReadRows. #18197

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0e40e98
feat: Added rst_stream exception handling for ReadRows. (#1298)
gkevinzheng Feb 25, 2026
dd45c75
resolve merge conflicts
daniel-sanche Aug 21, 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: 2 additions & 2 deletions packages/google-cloud-bigtable/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 @@ -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 (
Expand Down Expand Up @@ -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
Expand Down
48 changes: 47 additions & 1 deletion packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
)
Comment on lines +165 to +170

@gemini-code-assist gemini-code-assist Bot Aug 21, 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.

medium

If e.message is None (which can happen if an InternalServerError is instantiated without a message or with None), calling e.message.lower() will raise an AttributeError. To ensure robust defensive programming and prevent potential crashes in the retry loop, we should verify that e.message is a string before performing string operations on it.

Suggested change
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)
)
def rst_check(e):
return (
core_exceptions.ServiceUnavailable in exception_types
and isinstance(e, core_exceptions.InternalServerError)
and isinstance(e.message, str)
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
8 changes: 5 additions & 3 deletions packages/google-cloud-bigtable/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 @@ -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,
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 = _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
Expand Down
26 changes: 21 additions & 5 deletions packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,),
),
],
Expand Down Expand Up @@ -1588,6 +1595,7 @@ async def test_customizable_retryable_errors(
expected_retryables,
fn_name,
fn_args,
is_read_rows_fn,
is_stream,
extra_retryables,
):
Expand All @@ -1600,18 +1608,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._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:
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 Down
27 changes: 16 additions & 11 deletions packages/google-cloud-bigtable/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 @@ -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,),
),
],
Expand Down Expand Up @@ -1310,6 +1311,7 @@ def test_customizable_retryable_errors(
expected_retryables,
fn_name,
fn_args,
is_read_rows_fn,
is_stream,
extra_retryables,
):
Expand All @@ -1319,16 +1321,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._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:
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
42 changes: 42 additions & 0 deletions packages/google-cloud-bigtable/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,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",
Expand Down
Loading