feat: Implement updated execute query protocol

Author: jackdingilian

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__()

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}")

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))

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

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__()

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