jackdingilian · Mar 18, 2025 · Mar 6, 2025 · daniel-sanche · Mar 17, 2025 · jackdingilian
diff --git a/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py b/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py
@@ -91,6 +91,9 @@ def __init__(
         """
         Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
 
+        **Please Note** this is not meant to be constructed directly by applications. It should always
+        be created via the client. The constructor is subject to change.
+
         It is **not thread-safe**. It should not be used by multiple {TASK_OR_THREAD}.
 
         Args:
@@ -106,6 +109,7 @@ def __init__(
             retryable_excs: a list of errors that will be retried if encountered.
         Raises:
             {NO_LOOP}
+            :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
         """
         self._table_name = None
         self._app_profile_id = app_profile_id
@@ -217,6 +221,11 @@ async def _next_impl(self) -> CrossSync.Iterator[QueryResultRow]:
 
     @CrossSync.convert(sync_name="__next__", replace_symbols={"__anext__": "__next__"})
     async def __anext__(self) -> QueryResultRow:
+        """
+        Yields QueryResultRows representing the results of the query.
+
+        :raises: :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
+        """
         if self._is_closed:
             raise CrossSync.StopIteration
         return await self._result_generator.__anext__()

diff --git a/google/cloud/bigtable/data/execute_query/_byte_cursor.py b/google/cloud/bigtable/data/execute_query/_byte_cursor.py
@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import List, Optional
 
+from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
 from google.cloud.bigtable_v2 import ExecuteQueryResponse
 
 
@@ -31,9 +32,14 @@ class _ByteCursor:
     """
 
     def __init__(self):
-        self._buffer = bytearray()
+        self._batch_buffer = bytearray()
+        self._batches: List[bytes] = []
         self._resume_token = None
 
+    def reset(self):
+        self._batch_buffer = bytearray()
+        self._batches = []
+
     def prepare_for_new_request(self):
         """
         Prepares this ``_ByteCursor`` for retrying an ``ExecuteQuery`` request.
@@ -50,13 +56,15 @@ def prepare_for_new_request(self):
         Returns:
             bytes: Last received resume_token.
         """
-        self._buffer = bytearray()
+        # The first response of any retried stream will always contain reset, so
+        # this isn't actually necessary, but we do it for safety
+        self.reset()
         return self._resume_token
 
     def empty(self) -> bool:
-        return len(self._buffer) == 0
+        return not self._batch_buffer and not self._batches
 
-    def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
+    def consume(self, response: ExecuteQueryResponse) -> Optional[List[bytes]]:
         """
         Reads results bytes from an ``ExecuteQuery`` response and adds them to a buffer.
 
@@ -72,7 +80,8 @@ def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
                 Response obtained from the stream.
 
         Returns:
-            bytes or None: bytes if buffers were flushed or None otherwise.
+            bytes or None: List of bytes if buffers were flushed or None otherwise.
+            Each element in the list represents the bytes of a `ProtoRows` message.
 
         Raises:
             ValueError: If provided ``ExecuteQueryResponse`` is not valid
@@ -83,15 +92,31 @@ def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
 
         if response_pb.HasField("results"):
             results = response_pb.results
+            if results.reset:
+                self.reset()
             if results.HasField("proto_rows_batch"):
-                self._buffer.extend(results.proto_rows_batch.batch_data)
+                self._batch_buffer.extend(results.proto_rows_batch.batch_data)
+                # Note that 0 is a valid checksum so we must check for field presence
+                if results.HasField("batch_checksum"):
+                    expected_checksum = results.batch_checksum
+                    checksum = _CRC32C.checksum(self._batch_buffer)
+                    if expected_checksum != checksum:
+                        raise ValueError(
+                            f"Unexpected checksum mismatch. Expected: {expected_checksum}, got: {checksum}"
+                        )
+                    # We have a complete batch so we move it to batches and reset the
+                    # batch_buffer
+                    self._batches.append(memoryview(self._batch_buffer))
+                    self._batch_buffer = bytearray()
 
             if results.resume_token:
                 self._resume_token = results.resume_token
 
-                if self._buffer:
-                    return_value = memoryview(self._buffer)
-                    self._buffer = bytearray()
+                if self._batches:
+                    if self._batch_buffer:
+                        raise ValueError("Unexpected resume_token without checksum")
+                    return_value = self._batches
+                    self._batches = []
                     return return_value
         else:
             raise ValueError(f"Unexpected ExecuteQueryResponse: {response}")

diff --git a/google/cloud/bigtable/data/execute_query/_checksum.py b/google/cloud/bigtable/data/execute_query/_checksum.py
@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import warnings
+
+with warnings.catch_warnings(record=True) as import_warning:
+    import google_crc32c  # type: ignore
+
+
+class _CRC32C(object):
+    """
+    Wrapper around ``google_crc32c`` library
+    """
+
+    warn_emitted = False
+
+    @classmethod
+    def checksum(cls, val: bytearray) -> int:
+        """
+        Returns the crc32c checksum of the data.
+        """
+        if import_warning and not cls.warn_emitted:
+            cls.warn_emitted = True
+            warnings.warn(
+                "Using pure python implementation of `google-crc32` for ExecuteQuery response "
+                "validation. This is significantly slower than the c extension. If possible, "
+                "run in an environment that supports the c extension.",
+                RuntimeWarning,
+            )
+        memory_view = memoryview(val)
+        return google_crc32c.value(bytes(memory_view))
diff --git a/google/cloud/bigtable/data/execute_query/_reader.py b/google/cloud/bigtable/data/execute_query/_reader.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import (
+    List,
     TypeVar,
     Generic,
     Iterable,
@@ -54,14 +55,13 @@ class _Reader(ABC, Generic[T]):
 
     @abstractmethod
     def consume(
-        self, bytes_to_consume: bytes, metadata: Metadata
+        self, batches_to_consume: List[bytes], metadata: Metadata
     ) -> Optional[Iterable[T]]:
-        """This method receives a parsable chunk of bytes and returns either a None if there is
-        not enough chunks to return to the user yet (e.g. we haven't received all columns in a
-        row yet), or a list of appropriate values gathered from one or more parsable chunks.
-
+        """This method receives a list of batches of bytes to be parsed as ProtoRows messages.
+        It then uses the metadata to group the values in the parsed messages into rows. Returns
+        None if batches_to_consume is empty
         Args:
-            bytes_to_consume (bytes): chunk of parsable bytes received from
+            bytes_to_consume (bytes): chunk of parsable byte batches received from
                 :meth:`google.cloud.bigtable.byte_cursor._ByteCursor.consume`
                 method.
             metadata: metadata used to transform values to rows
@@ -85,11 +85,9 @@ class _QueryResultRowReader(_Reader[QueryResultRow]):
     :class:`google.cloud.bigtable.byte_cursor._ByteCursor` passed in the constructor.
     """
 
-    def __init__(self):
-        """
-        Constructs new instance of ``_QueryResultRowReader``.
-        """
-        self._values = []
+    def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
+        proto_rows = ProtoRows.pb().FromString(bytes_to_parse)
+        return proto_rows.values
 
     def _construct_query_result_row(
         self, values: Sequence[PBValue], metadata: ProtoMetadata
@@ -106,36 +104,23 @@ def _construct_query_result_row(
             result.add_field(column.column_name, parsed_value)
         return result
 
-    def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
-        proto_rows = ProtoRows.pb().FromString(bytes_to_parse)
-        return proto_rows.values
-
     def consume(
-        self, bytes_to_consume: bytes, metadata: Metadata
+        self, batches_to_consume: List[bytes], metadata: Metadata
     ) -> Optional[Iterable[QueryResultRow]]:
-        if bytes_to_consume is None:
-            raise ValueError("bytes_to_consume shouldn't be None")
-
-        self._values.extend(self._parse_proto_rows(bytes_to_consume))
-
-        # The logic, not defined by mypy types, ensures that the value of
-        # "metadata" is never null at the time it is retrieved here
         proto_metadata = cast(ProtoMetadata, metadata)
         num_columns = len(proto_metadata.columns)
-
-        if len(self._values) < num_columns:
-            return None
-
         rows = []
-        for batch in batched(self._values, n=num_columns):
-            if len(batch) == num_columns:
-                rows.append(self._construct_query_result_row(batch, proto_metadata))
-            else:
-                raise ValueError(
-                    "Server error, recieved bad number of values. "
-                    f"Expected {num_columns} got {len(batch)}."
-                )
-
-        self._values = []
+        for batch_bytes in batches_to_consume:
+            values = self._parse_proto_rows(batch_bytes)
+            for row_data in batched(values, n=num_columns):
+                if len(row_data) == num_columns:
+                    rows.append(
+                        self._construct_query_result_row(row_data, proto_metadata)
+                    )
+                else:
+                    raise ValueError(
+                        "Unexpected error, recieved bad number of values. "
+                        f"Expected {num_columns} got {len(row_data)}."
+                    )
 
         return rows
diff --git a/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py b/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py
@@ -66,6 +66,9 @@ def __init__(
     ) -> None:
         """Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
 
+        **Please Note** this is not meant to be constructed directly by applications. It should always
+        be created via the client. The constructor is subject to change.
+
         It is **not thread-safe**. It should not be used by multiple threads.
 
         Args:
@@ -80,7 +83,9 @@ def __init__(
             req_metadata: metadata used while sending the gRPC request
             retryable_excs: a list of errors that will be retried if encountered.
         Raises:
-            None"""
+            None
+            :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
+        """
         self._table_name = None
         self._app_profile_id = app_profile_id
         self._client = client
@@ -173,6 +178,10 @@ def _next_impl(self) -> CrossSync._Sync_Impl.Iterator[QueryResultRow]:
         self.close()
 
     def __next__(self) -> QueryResultRow:
+        """Yields QueryResultRows representing the results of the query.
+
+        :raises: :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
+        """
         if self._is_closed:
             raise CrossSync._Sync_Impl.StopIteration
         return self._result_generator.__next__()

diff --git a/setup.py b/setup.py
@@ -44,6 +44,7 @@
     "proto-plus >= 1.22.3, <2.0.0dev",
     "proto-plus >= 1.25.0, <2.0.0dev; python_version>='3.13'",
     "protobuf>=3.20.2,<6.0.0dev,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",
+    "google-crc32c>=1.5.0, <2.0.0dev",
 ]
 extras = {"libcst": "libcst >= 0.2.5"}