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 3e6c2b533b1e..f2fccf4f2eed 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 @@ -176,6 +176,7 @@ async def chunk_stream( Yields: ReadRowsResponsePB.CellChunk: the next chunk in the stream """ + current_key = None async for resp in await stream: # extract proto from proto-plus wrapper resp = resp._pb @@ -190,12 +191,11 @@ async def chunk_stream( raise InvalidChunk("last scanned out of order") self._last_yielded_row_key = resp.last_scanned_row_key - current_key = None # process each chunk in the response for c in resp.chunks: if current_key is None: current_key = c.row_key - if current_key is None: + if not current_key: raise InvalidChunk("first chunk is missing a row key") elif ( self._last_yielded_row_key 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 ce29226cac53..18e7e3319058 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 @@ -17,36 +17,22 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations - import time from typing import TYPE_CHECKING, Sequence - 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._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, - _ResetRow, - _RowSetComplete, -) +from google.cloud.bigtable.data.exceptions import InvalidChunk, _ResetRow, _RowSetComplete from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.row import Cell, Row from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB from google.cloud.bigtable_v2.types import RowRange as RowRangePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB - if TYPE_CHECKING: from google.cloud.bigtable.data._metrics import ActiveOperationMetric - from google.cloud.bigtable.data._sync_autogen.client import ( - _DataApiTarget as TargetType, - ) - + from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget as TargetType class _ReadRowsOperation: """ @@ -68,35 +54,13 @@ class _ReadRowsOperation: metric: the metric object representing the active operation retryable_exceptions: A list of exceptions that should trigger a retry """ + __slots__ = ('attempt_timeout_gen', 'operation_timeout', 'request', 'target', '_predicate', '_last_yielded_row_key', '_remaining_count', '_operation_metric') - __slots__ = ( - "attempt_timeout_gen", - "operation_timeout", - "request", - "target", - "_predicate", - "_last_yielded_row_key", - "_remaining_count", - "_operation_metric", - ) - - def __init__( - self, - query: ReadRowsQuery, - target: TargetType, - operation_timeout: float, - attempt_timeout: float, - metric: ActiveOperationMetric, - retryable_exceptions: Sequence[type[Exception]] = (), - ): - self.attempt_timeout_gen = _attempt_timeout_generator( - attempt_timeout, operation_timeout - ) + def __init__(self, query: ReadRowsQuery, target: TargetType, operation_timeout: float, attempt_timeout: float, metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]]=()): + self.attempt_timeout_gen = _attempt_timeout_generator(attempt_timeout, operation_timeout) self.operation_timeout = operation_timeout if isinstance(query, dict): - self.request = ReadRowsRequestPB( - **query, **target._request_path, app_profile_id=target.app_profile_id - ) + self.request = ReadRowsRequestPB(**query, **target._request_path, app_profile_id=target.app_profile_id) else: self.request = query._to_pb(target) self.target = target @@ -108,75 +72,53 @@ def __init__( def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: """Start the read_rows operation, retrying on retryable errors. - Yields: - Row: The next row in the stream""" - return tracked_retry( - retry_fn=CrossSync._Sync_Impl.retry_target_stream, - operation=self._operation_metric, - target=self._read_rows_attempt, - predicate=self._predicate, - timeout=self.operation_timeout, - ) +Yields: + Row: The next row in the stream""" + return tracked_retry(retry_fn=CrossSync._Sync_Impl.retry_target_stream, operation=self._operation_metric, target=self._read_rows_attempt, predicate=self._predicate, timeout=self.operation_timeout) def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: """Attempt a single read_rows rpc call. - This function is intended to be wrapped by retry logic, - which will call this function until it succeeds or - a non-retryable error is raised. +This function is intended to be wrapped by retry logic, +which will call this function until it succeeds or +a non-retryable error is raised. - Yields: - Row: The next row in the stream""" +Yields: + Row: The next row in the stream""" self._operation_metric.start_attempt() if self._last_yielded_row_key is not None: try: - self.request.rows = self._revise_request_rowset( - row_set=self.request.rows, - last_seen_row_key=self._last_yielded_row_key, - ) + self.request.rows = self._revise_request_rowset(row_set=self.request.rows, last_seen_row_key=self._last_yielded_row_key) except _RowSetComplete: return self.merge_rows(None) if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: return self.merge_rows(None) - gapic_stream = self.target.client._gapic_client.read_rows( - self.request, timeout=next(self.attempt_timeout_gen), retry=None - ) + gapic_stream = self.target.client._gapic_client.read_rows(self.request, timeout=next(self.attempt_timeout_gen), retry=None) chunked_stream = self.chunk_stream(gapic_stream) return self.merge_rows(chunked_stream) - def chunk_stream( - self, - stream: CrossSync._Sync_Impl.Awaitable[ - CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB] - ], - ) -> CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk]: + def chunk_stream(self, stream: CrossSync._Sync_Impl.Awaitable[CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB]]) -> CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk]: """process chunks out of raw read_rows stream - Args: - stream: the raw read_rows stream from the gapic client - Yields: - ReadRowsResponsePB.CellChunk: the next chunk in the stream""" +Args: + stream: the raw read_rows stream from the gapic client +Yields: + ReadRowsResponsePB.CellChunk: the next chunk in the stream""" + current_key = None for resp in stream: resp = resp._pb if resp.last_scanned_row_key: - if ( - self._last_yielded_row_key is not None - and resp.last_scanned_row_key <= self._last_yielded_row_key - ): - raise InvalidChunk("last scanned out of order") + if self._last_yielded_row_key is not None and resp.last_scanned_row_key <= self._last_yielded_row_key: + raise InvalidChunk('last scanned out of order') self._last_yielded_row_key = resp.last_scanned_row_key - current_key = None for c in resp.chunks: if current_key is None: current_key = c.row_key - if current_key is None: - raise InvalidChunk("first chunk is missing a row key") - elif ( - self._last_yielded_row_key - and current_key <= self._last_yielded_row_key - ): - raise InvalidChunk("row keys should be strictly increasing") + if not current_key: + raise InvalidChunk('first chunk is missing a row key') + elif self._last_yielded_row_key and current_key <= self._last_yielded_row_key: + raise InvalidChunk('row keys should be strictly increasing') yield c if c.reset_row: current_key = None @@ -185,18 +127,16 @@ def chunk_stream( if self._remaining_count is not None: self._remaining_count -= 1 if self._remaining_count < 0: - raise InvalidChunk("emit count exceeds row limit") + raise InvalidChunk('emit count exceeds row limit') current_key = None - def merge_rows( - self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None - ) -> CrossSync._Sync_Impl.Iterable[Row]: + def merge_rows(self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows - Args: - chunks: the chunk stream to merge - Yields: - Row: the next row in the stream""" +Args: + chunks: the chunk stream to merge +Yields: + Row: the next row in the stream""" try: if chunks is None: self._operation_metric.end_with_success() @@ -210,7 +150,7 @@ def merge_rows( return row_key = c.row_key if not row_key: - raise InvalidChunk("first row chunk is missing key") + raise InvalidChunk('first row chunk is missing key') cells = [] family: str | None = None qualifier: bytes | None = None @@ -220,23 +160,23 @@ def merge_rows( raise _ResetRow(c) k = c.row_key f = c.family_name.value - q = c.qualifier.value if c.HasField("qualifier") else None + q = c.qualifier.value if c.HasField('qualifier') else None if k and k != row_key: - raise InvalidChunk("unexpected new row key") + raise InvalidChunk('unexpected new row key') if f: family = f if q is not None: qualifier = q else: - raise InvalidChunk("new family without qualifier") + raise InvalidChunk('new family without qualifier') elif family is None: - raise InvalidChunk("missing family") + raise InvalidChunk('missing family') elif q is not None: if family is None: - raise InvalidChunk("new qualifier without family") + raise InvalidChunk('new qualifier without family') qualifier = q elif qualifier is None: - raise InvalidChunk("missing qualifier") + raise InvalidChunk('missing qualifier') ts = c.timestamp_micros labels = c.labels if c.labels else [] value = c.value @@ -247,52 +187,35 @@ def merge_rows( t = c.timestamp_micros cl = c.labels k = c.row_key - if ( - c.HasField("family_name") - and c.family_name.value != family - ): - raise InvalidChunk("family changed mid cell") - if ( - c.HasField("qualifier") - and c.qualifier.value != qualifier - ): - raise InvalidChunk("qualifier changed mid cell") + if c.HasField('family_name') and c.family_name.value != family: + raise InvalidChunk('family changed mid cell') + if c.HasField('qualifier') and c.qualifier.value != qualifier: + raise InvalidChunk('qualifier changed mid cell') if t and t != ts: - raise InvalidChunk("timestamp changed mid cell") + raise InvalidChunk('timestamp changed mid cell') if cl and cl != labels: - raise InvalidChunk("labels changed mid cell") + raise InvalidChunk('labels changed mid cell') if k and k != row_key: - raise InvalidChunk("row key changed mid cell") + raise InvalidChunk('row key changed mid cell') if c.reset_row: raise _ResetRow(c) buffer.append(c.value) - value = b"".join(buffer) - cells.append( - Cell(value, row_key, family, qualifier, ts, list(labels)) - ) + value = b''.join(buffer) + cells.append(Cell(value, row_key, family, qualifier, ts, list(labels))) if c.commit_row: block_time = time.monotonic_ns() yield Row(row_key, cells) if self._operation_metric.active_attempt is not None: - self._operation_metric.active_attempt.application_blocking_time_ns += ( - time.monotonic_ns() - block_time - ) + self._operation_metric.active_attempt.application_blocking_time_ns += time.monotonic_ns() - block_time break c = it.__next__() except _ResetRow as e: c = e.chunk - if ( - c.row_key - or c.HasField("family_name") - or c.HasField("qualifier") - or c.timestamp_micros - or c.labels - or c.value - ): - raise InvalidChunk("reset row with data") + if c.row_key or c.HasField('family_name') or c.HasField('qualifier') or c.timestamp_micros or c.labels or c.value: + raise InvalidChunk('reset row with data') continue except CrossSync._Sync_Impl.StopIteration: - raise InvalidChunk("premature end of stream") + raise InvalidChunk('premature end of stream') except GeneratorExit as close_exception: self._operation_metric.end_with_status(StatusCode.CANCELLED) raise close_exception @@ -303,19 +226,17 @@ def merge_rows( def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB: """Revise the rows in the request to avoid ones we've already processed. - Args: - row_set: the row set from the request - last_seen_row_key: the last row key encountered - Returns: - RowSetPB: the new rowset after adusting for the last seen key - Raises: - _RowSetComplete: if there are no rows left to process after the revision""" +Args: + row_set: the row set from the request + last_seen_row_key: the last row key encountered +Returns: + RowSetPB: the new rowset after adusting for the last seen key +Raises: + _RowSetComplete: if there are no rows left to process after the revision""" if row_set is None or (not row_set.row_ranges and (not row_set.row_keys)): last_seen = last_seen_row_key return RowSetPB(row_ranges=[RowRangePB(start_key_open=last_seen)]) - adjusted_keys: list[bytes] = [ - k for k in row_set.row_keys if k > last_seen_row_key - ] + adjusted_keys: list[bytes] = [k for k in row_set.row_keys if k > last_seen_row_key] adjusted_ranges: list[RowRangePB] = [] for row_range in row_set.row_ranges: end_key = row_range.end_key_closed or row_range.end_key_open or None @@ -327,4 +248,4 @@ def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSe adjusted_ranges.append(new_range) if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0: raise _RowSetComplete() - return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges) + return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges) \ No newline at end of file diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py index 24b92d2de5ad..9dc51aaafc87 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py @@ -391,3 +391,53 @@ async def mock_stream(): with pytest.raises(InvalidChunk) as exc: await stream.__anext__() assert "row keys should be strictly increasing" in str(exc.value) + + @CrossSync.pytest + @CrossSync.convert(replace_symbols={"__anext__": "__next__"}) + async def test_chunk_stream_cross_response_row_key(self): + """ + A row spanning multiple ReadRowsResponse messages must not clobber + _last_yielded_row_key with b"" when the continuation chunk omits + row_key (protobuf default). Regression test for the bug where + current_key was reset to None at the start of each response. + """ + from google.cloud.bigtable_v2.types import ReadRowsResponse + + row_key = b"row1" + + async def mock_awaitable_stream(): + async def mock_stream(): + # First response: row starts but does not commit + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk( + row_key=row_key, + family_name="f", + qualifier=b"q", + value=b"part1", + ) + ] + ) + # Second response: continuation chunk commits; row_key is omitted (b"") + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk( + value=b"part2", + commit_row=True, + ) + ] + ) + + return mock_stream() + + instance = mock.Mock() + instance._last_yielded_row_key = None + instance._remaining_count = None + stream = self._get_target_class().chunk_stream( + instance, mock_awaitable_stream() + ) + # consume both chunks + async for _ in stream: + pass + # _last_yielded_row_key must be the actual row key, not b"" + assert instance._last_yielded_row_key == row_key diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py index b6cac9cafbc0..502a48475fae 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py @@ -16,13 +16,10 @@ # This file is automatically generated by CrossSync. Do not edit manually. from unittest import mock - import pytest - from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._metrics import ActiveOperationMetric - class TestReadRowsOperation: """ Tests helper functions in the ReadRowsOperation class @@ -39,7 +36,6 @@ def _make_one(self, *args, **kwargs): def test_ctor(self): from google.cloud.bigtable.data import ReadRowsQuery - row_limit = 91 query = ReadRowsQuery(limit=row_limit) client = mock.Mock() @@ -47,58 +43,35 @@ def test_ctor(self): client.read_rows.return_value = None table = mock.Mock() table._client = client - table._request_path = {"table_name": "test_table"} - table.app_profile_id = "test_profile" + table._request_path = {'table_name': 'test_table'} + table.app_profile_id = 'test_profile' expected_operation_timeout = 42 expected_request_timeout = 44 time_gen_mock = mock.Mock() expected_metric = mock.Mock() - subpath = "_async" if CrossSync._Sync_Impl.is_async else "_sync_autogen" - with mock.patch( - f"google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator", - time_gen_mock, - ): - instance = self._make_one( - query, - table, - operation_timeout=expected_operation_timeout, - attempt_timeout=expected_request_timeout, - metric=expected_metric, - ) + subpath = '_async' if CrossSync._Sync_Impl.is_async else '_sync_autogen' + with mock.patch(f'google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator', time_gen_mock): + instance = self._make_one(query, table, operation_timeout=expected_operation_timeout, attempt_timeout=expected_request_timeout, metric=expected_metric) assert time_gen_mock.call_count == 1 - time_gen_mock.assert_called_once_with( - expected_request_timeout, expected_operation_timeout - ) + time_gen_mock.assert_called_once_with(expected_request_timeout, expected_operation_timeout) assert instance._last_yielded_row_key is None assert instance._remaining_count == row_limit assert instance.operation_timeout == expected_operation_timeout assert client.read_rows.call_count == 0 - assert instance.request.table_name == "test_table" + assert instance.request.table_name == 'test_table' assert instance.request.app_profile_id == table.app_profile_id assert instance.request.rows_limit == row_limit assert instance._operation_metric == expected_metric - @pytest.mark.parametrize( - "in_keys,last_key,expected", - [ - (["b", "c", "d"], "a", ["b", "c", "d"]), - (["a", "b", "c"], "b", ["c"]), - (["a", "b", "c"], "c", []), - (["a", "b", "c"], "d", []), - (["d", "c", "b", "a"], "b", ["d", "c"]), - ], - ) - @pytest.mark.parametrize("with_range", [True, False]) - def test_revise_request_rowset_keys_with_range( - self, in_keys, last_key, expected, with_range - ): + @pytest.mark.parametrize('in_keys,last_key,expected', [(['b', 'c', 'd'], 'a', ['b', 'c', 'd']), (['a', 'b', 'c'], 'b', ['c']), (['a', 'b', 'c'], 'c', []), (['a', 'b', 'c'], 'd', []), (['d', 'c', 'b', 'a'], 'b', ['d', 'c'])]) + @pytest.mark.parametrize('with_range', [True, False]) + def test_revise_request_rowset_keys_with_range(self, in_keys, last_key, expected, with_range): from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable_v2.types import RowRange as RowRangePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB - - in_keys = [key.encode("utf-8") for key in in_keys] - expected = [key.encode("utf-8") for key in expected] - last_key = last_key.encode("utf-8") + in_keys = [key.encode('utf-8') for key in in_keys] + expected = [key.encode('utf-8') for key in expected] + last_key = last_key.encode('utf-8') if with_range: sample_range = [RowRangePB(start_key_open=last_key)] else: @@ -112,68 +85,16 @@ def test_revise_request_rowset_keys_with_range( assert revised.row_keys == expected assert revised.row_ranges == sample_range - @pytest.mark.parametrize( - "in_ranges,last_key,expected", - [ - ( - [{"start_key_open": "b", "end_key_closed": "d"}], - "a", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_closed": "b", "end_key_closed": "d"}], - "a", - [{"start_key_closed": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_open": "a", "end_key_closed": "d"}], - "b", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_closed": "a", "end_key_open": "d"}], - "b", - [{"start_key_open": "b", "end_key_open": "d"}], - ), - ( - [{"start_key_closed": "b", "end_key_closed": "d"}], - "b", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ([{"start_key_closed": "b", "end_key_closed": "d"}], "d", []), - ([{"start_key_closed": "b", "end_key_open": "d"}], "d", []), - ([{"start_key_closed": "b", "end_key_closed": "d"}], "e", []), - ([{"start_key_closed": "b"}], "z", [{"start_key_open": "z"}]), - ([{"start_key_closed": "b"}], "a", [{"start_key_closed": "b"}]), - ( - [{"end_key_closed": "z"}], - "a", - [{"start_key_open": "a", "end_key_closed": "z"}], - ), - ( - [{"end_key_open": "z"}], - "a", - [{"start_key_open": "a", "end_key_open": "z"}], - ), - ], - ) - @pytest.mark.parametrize("with_key", [True, False]) - def test_revise_request_rowset_ranges( - self, in_ranges, last_key, expected, with_key - ): + @pytest.mark.parametrize('in_ranges,last_key,expected', [([{'start_key_open': 'b', 'end_key_closed': 'd'}], 'a', [{'start_key_open': 'b', 'end_key_closed': 'd'}]), ([{'start_key_closed': 'b', 'end_key_closed': 'd'}], 'a', [{'start_key_closed': 'b', 'end_key_closed': 'd'}]), ([{'start_key_open': 'a', 'end_key_closed': 'd'}], 'b', [{'start_key_open': 'b', 'end_key_closed': 'd'}]), ([{'start_key_closed': 'a', 'end_key_open': 'd'}], 'b', [{'start_key_open': 'b', 'end_key_open': 'd'}]), ([{'start_key_closed': 'b', 'end_key_closed': 'd'}], 'b', [{'start_key_open': 'b', 'end_key_closed': 'd'}]), ([{'start_key_closed': 'b', 'end_key_closed': 'd'}], 'd', []), ([{'start_key_closed': 'b', 'end_key_open': 'd'}], 'd', []), ([{'start_key_closed': 'b', 'end_key_closed': 'd'}], 'e', []), ([{'start_key_closed': 'b'}], 'z', [{'start_key_open': 'z'}]), ([{'start_key_closed': 'b'}], 'a', [{'start_key_closed': 'b'}]), ([{'end_key_closed': 'z'}], 'a', [{'start_key_open': 'a', 'end_key_closed': 'z'}]), ([{'end_key_open': 'z'}], 'a', [{'start_key_open': 'a', 'end_key_open': 'z'}])]) + @pytest.mark.parametrize('with_key', [True, False]) + def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected, with_key): from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable_v2.types import RowRange as RowRangePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB - - next_key = (last_key + "a").encode("utf-8") - last_key = last_key.encode("utf-8") - in_ranges = [ - RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) - for r in in_ranges - ] - expected = [ - RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) for r in expected - ] + next_key = (last_key + 'a').encode('utf-8') + last_key = last_key.encode('utf-8') + in_ranges = [RowRangePB(**{k: v.encode('utf-8') for k, v in r.items()}) for r in in_ranges] + expected = [RowRangePB(**{k: v.encode('utf-8') for k, v in r.items()}) for r in expected] if with_key: row_keys = [next_key] else: @@ -187,17 +108,14 @@ def test_revise_request_rowset_ranges( assert revised.row_keys == row_keys assert revised.row_ranges == expected - @pytest.mark.parametrize("last_key", ["a", "b", "c"]) + @pytest.mark.parametrize('last_key', ['a', 'b', 'c']) def test_revise_request_full_table(self, last_key): from google.cloud.bigtable_v2.types import RowRange as RowRangePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB - - last_key = last_key.encode("utf-8") + last_key = last_key.encode('utf-8') row_set = RowSetPB() for selected_set in [row_set, None]: - revised = self._get_target_class()._revise_request_rowset( - selected_set, last_key - ) + revised = self._get_target_class()._revise_request_rowset(selected_set, last_key) assert revised.row_keys == [] assert len(revised.row_ranges) == 1 assert revised.row_ranges[0] == RowRangePB(start_key_open=last_key) @@ -207,114 +125,72 @@ def test_revise_to_empty_rowset(self): from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable_v2.types import RowRange as RowRangePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB - - row_keys = [b"a", b"b", b"c"] - row_range = RowRangePB(end_key_open=b"c") + row_keys = [b'a', b'b', b'c'] + row_range = RowRangePB(end_key_open=b'c') row_set = RowSetPB(row_keys=row_keys, row_ranges=[row_range]) with pytest.raises(_RowSetComplete): - self._get_target_class()._revise_request_rowset(row_set, b"d") + self._get_target_class()._revise_request_rowset(row_set, b'd') - @pytest.mark.parametrize( - "start_limit,emit_num,expected_limit", - [ - (10, 0, 10), - (10, 1, 9), - (10, 10, 0), - (None, 10, None), - (None, 0, None), - (4, 2, 2), - ], - ) + @pytest.mark.parametrize('start_limit,emit_num,expected_limit', [(10, 0, 10), (10, 1, 9), (10, 10, 0), (None, 10, None), (None, 0, None), (4, 2, 2)]) def test_revise_limit(self, start_limit, emit_num, expected_limit): """revise_limit should revise the request's limit field - - if limit is 0 (unlimited), it should never be revised - - if start_limit-emit_num == 0, the request should end early - - if the number emitted exceeds the new limit, an exception should - should be raised (tested in test_revise_limit_over_limit)""" +- if limit is 0 (unlimited), it should never be revised +- if start_limit-emit_num == 0, the request should end early +- if the number emitted exceeds the new limit, an exception should + should be raised (tested in test_revise_limit_over_limit)""" from google.cloud.bigtable.data import ReadRowsQuery from google.cloud.bigtable_v2.types import ReadRowsResponse def awaitable_stream(): + def mock_stream(): for i in range(emit_num): - yield ReadRowsResponse( - chunks=[ - ReadRowsResponse.CellChunk( - row_key=str(i).encode(), - family_name="b", - qualifier=b"c", - value=b"d", - commit_row=True, - ) - ] - ) - + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(row_key=str(i).encode(), family_name='b', qualifier=b'c', value=b'd', commit_row=True)]) return mock_stream() - query = ReadRowsQuery(limit=start_limit) table = mock.Mock() - table._request_path = {"table_name": "table_name"} - table.app_profile_id = "app_profile_id" - instance = self._make_one( - query, table, 10, 10, ActiveOperationMetric("READ_ROWS") - ) + table._request_path = {'table_name': 'table_name'} + table.app_profile_id = 'app_profile_id' + instance = self._make_one(query, table, 10, 10, ActiveOperationMetric('READ_ROWS')) assert instance._remaining_count == start_limit for val in instance.chunk_stream(awaitable_stream()): pass assert instance._remaining_count == expected_limit - @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) + @pytest.mark.parametrize('start_limit,emit_num', [(5, 10), (3, 9), (1, 10)]) def test_revise_limit_over_limit(self, start_limit, emit_num): """Should raise runtime error if we get in state where emit_num > start_num - (unless start_num == 0, which represents unlimited)""" +(unless start_num == 0, which represents unlimited)""" from google.cloud.bigtable.data import ReadRowsQuery from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable_v2.types import ReadRowsResponse def awaitable_stream(): + def mock_stream(): for i in range(emit_num): - yield ReadRowsResponse( - chunks=[ - ReadRowsResponse.CellChunk( - row_key=str(i).encode(), - family_name="b", - qualifier=b"c", - value=b"d", - commit_row=True, - ) - ] - ) - + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(row_key=str(i).encode(), family_name='b', qualifier=b'c', value=b'd', commit_row=True)]) return mock_stream() - query = ReadRowsQuery(limit=start_limit) table = mock.Mock() - table._request_path = {"table_name": "table_name"} - table.app_profile_id = "app_profile_id" - instance = self._make_one( - query, table, 10, 10, ActiveOperationMetric("READ_ROWS") - ) + table._request_path = {'table_name': 'table_name'} + table.app_profile_id = 'app_profile_id' + instance = self._make_one(query, table, 10, 10, ActiveOperationMetric('READ_ROWS')) assert instance._remaining_count == start_limit with pytest.raises(InvalidChunk) as e: for val in instance.chunk_stream(awaitable_stream()): pass - assert "emit count exceeds row limit" in str(e.value) + assert 'emit count exceeds row limit' in str(e.value) def test_close(self): """should be able to close a stream safely with close. - Closed generators should raise StopAsyncIteration on next yield""" +Closed generators should raise StopAsyncIteration on next yield""" def mock_stream(): while True: yield 1 - - with mock.patch.object( - self._get_target_class(), "_read_rows_attempt" - ) as mock_attempt: - instance = self._make_one( - mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric("READ_ROWS") - ) + with mock.patch.object(self._get_target_class(), '_read_rows_attempt') as mock_attempt: + instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric('READ_ROWS')) wrapped_gen = mock_stream() mock_attempt.return_value = wrapped_gen gen = instance.start_operation() @@ -330,32 +206,42 @@ def test_retryable_ignore_repeated_rows(self): """Duplicate rows should cause an invalid chunk error""" from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable_v2.types import ReadRowsResponse - - row_key = b"duplicate" + row_key = b'duplicate' def mock_awaitable_stream(): + def mock_stream(): while True: - yield ReadRowsResponse( - chunks=[ - ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True) - ] - ) - yield ReadRowsResponse( - chunks=[ - ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True) - ] - ) - + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True)]) + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True)]) return mock_stream() - instance = mock.Mock() instance._last_yielded_row_key = None instance._remaining_count = None - stream = self._get_target_class().chunk_stream( - instance, mock_awaitable_stream() - ) + stream = self._get_target_class().chunk_stream(instance, mock_awaitable_stream()) stream.__next__() with pytest.raises(InvalidChunk) as exc: stream.__next__() - assert "row keys should be strictly increasing" in str(exc.value) + assert 'row keys should be strictly increasing' in str(exc.value) + + def test_chunk_stream_cross_response_row_key(self): + """A row spanning multiple ReadRowsResponse messages must not clobber +_last_yielded_row_key with b"" when the continuation chunk omits +row_key (protobuf default). Regression test for the bug where +current_key was reset to None at the start of each response.""" + from google.cloud.bigtable_v2.types import ReadRowsResponse + row_key = b'row1' + + def mock_awaitable_stream(): + + def mock_stream(): + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(row_key=row_key, family_name='f', qualifier=b'q', value=b'part1')]) + yield ReadRowsResponse(chunks=[ReadRowsResponse.CellChunk(value=b'part2', commit_row=True)]) + return mock_stream() + instance = mock.Mock() + instance._last_yielded_row_key = None + instance._remaining_count = None + stream = self._get_target_class().chunk_stream(instance, mock_awaitable_stream()) + for _ in stream: + pass + assert instance._last_yielded_row_key == row_key \ No newline at end of file