diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py index 0c1565737..ab0358285 100644 --- a/google/cloud/bigtable/row_data.py +++ b/google/cloud/bigtable/row_data.py @@ -474,7 +474,11 @@ def _read_next(self): def _read_next_response(self): """Helper for :meth:`__iter__`.""" - return self.retry(self._read_next, on_error=self._on_error)() + resp_protoplus = self.retry(self._read_next, on_error=self._on_error)() + # unwrap the underlying protobuf, there is a significant amount of + # overhead that protoplus imposes for very little gain. The protos + # are not user visible, so we just use the raw protos for merging. + return data_messages_v2_pb2.ReadRowsResponse.pb(resp_protoplus) def __iter__(self): """Consume the ``ReadRowsResponse`` s from the stream. @@ -543,11 +547,12 @@ def _process_chunk(self, chunk): def _update_cell(self, chunk): if self._cell is None: qualifier = None - if "qualifier" in chunk: - qualifier = chunk.qualifier + if chunk.HasField("qualifier"): + qualifier = chunk.qualifier.value + family = None - if "family_name" in chunk: - family = chunk.family_name + if chunk.HasField("family_name"): + family = chunk.family_name.value self._cell = PartialCellData( chunk.row_key, @@ -577,8 +582,8 @@ def _validate_chunk_reset_row(self, chunk): # No reset with other keys _raise_if(chunk.row_key) - _raise_if("family_name" in chunk) - _raise_if("qualifier" in chunk) + _raise_if(chunk.HasField("family_name")) + _raise_if(chunk.HasField("qualifier")) _raise_if(chunk.timestamp_micros) _raise_if(chunk.labels) _raise_if(chunk.value_size) diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 9b329dc9f..94a90aa24 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -637,15 +637,15 @@ def test_partial_rows_data__copy_from_previous_filled(): def test_partial_rows_data_valid_last_scanned_row_key_on_start(): client = _Client() - response = _ReadRowsResponseV2(chunks=(), last_scanned_row_key="2.AFTER") + response = _ReadRowsResponseV2([], last_scanned_row_key=b"2.AFTER") iterator = _MockCancellableIterator(response) client._data_stub = mock.MagicMock() client._data_stub.read_rows.side_effect = [iterator] request = object() yrd = _make_partial_rows_data(client._data_stub.read_rows, request) - yrd.last_scanned_row_key = "1.BEFORE" + yrd.last_scanned_row_key = b"1.BEFORE" _partial_rows_data_consume_all(yrd) - assert yrd.last_scanned_row_key == "2.AFTER" + assert yrd.last_scanned_row_key == b"2.AFTER" def test_partial_rows_data_invalid_empty_chunk(): @@ -666,6 +666,7 @@ def test_partial_rows_data_invalid_empty_chunk(): def test_partial_rows_data_state_cell_in_progress(): from google.cloud.bigtable_v2.services.bigtable import BigtableClient + from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 LABELS = ["L1", "L2"] @@ -682,6 +683,9 @@ def test_partial_rows_data_state_cell_in_progress(): value=VALUE, labels=LABELS, ) + # _update_cell expects to be called after the protoplus wrapper has been + # shucked + chunk = messages_v2_pb2.ReadRowsResponse.CellChunk.pb(chunk) yrd._update_cell(chunk) more_cell_data = _ReadRowsResponseCellChunkPB(value=VALUE) @@ -1455,10 +1459,12 @@ def __init__(self, **kw): self.__dict__.update(kw) -class _ReadRowsResponseV2(object): - def __init__(self, chunks, last_scanned_row_key=""): - self.chunks = chunks - self.last_scanned_row_key = last_scanned_row_key +def _ReadRowsResponseV2(chunks, last_scanned_row_key=b""): + from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 + + return messages_v2_pb2.ReadRowsResponse( + chunks=chunks, last_scanned_row_key=last_scanned_row_key + ) def _generate_cell_chunks(chunk_text_pbs): diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 883f713d8..a89e02e8c 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -2206,10 +2206,12 @@ def next(self): __next__ = next -class _ReadRowsResponseV2(object): - def __init__(self, chunks, last_scanned_row_key=""): - self.chunks = chunks - self.last_scanned_row_key = last_scanned_row_key +def _ReadRowsResponseV2(chunks, last_scanned_row_key=b""): + from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 + + return messages_v2_pb2.ReadRowsResponse( + chunks=chunks, last_scanned_row_key=last_scanned_row_key + ) def _TablePB(*args, **kw):