Visitar URL original
feat: Implement updated execute query protocol · googleapis/python-bigtable@867cd6b · GitHub
Skip to content
This repository was archived by the owner on Aug 30, 2026. It is now read-only.

Commit 867cd6b

Browse files
jackdingilianjackdingilian
committed
feat: Implement updated execute query protocol
1 parent 5527282 commit 867cd6b

15 files changed

Lines changed: 555 additions & 424 deletions

File tree

‎google/cloud/bigtable/data/_async/client.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from functools import partial
3535
from grpc import Channel
3636

37+
from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
3738
from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
3839
from google.cloud.bigtable.data.execute_query.metadata import (
3940
SqlType,
@@ -218,6 +219,7 @@ def __init__(
218219
self._executor = (
219220
concurrent.futures.ThreadPoolExecutor() if not CrossSync.is_async else None
220221
)
222+
self._crc32c = _CRC32C()
221223
if self._emulator_host is None:
222224
# attempt to start background channel refresh tasks
223225
try:
@@ -667,6 +669,7 @@ async def execute_query(
667669
app_profile_id,
668670
request_body,
669671
prepare_metadata,
672+
self._crc32c,
670673
attempt_timeout,
671674
operation_timeout,
672675
retryable_excs=retryable_excs,

‎google/cloud/bigtable/data/_sync_autogen/client.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import concurrent.futures
2626
from functools import partial
2727
from grpc import Channel
28+
from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
2829
from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
2930
from google.cloud.bigtable.data.execute_query.metadata import (
3031
SqlType,
@@ -164,6 +165,7 @@ def __init__(
164165
if not CrossSync._Sync_Impl.is_async
165166
else None
166167
)
168+
self._crc32c = _CRC32C()
167169
if self._emulator_host is None:
168170
try:
169171
self._start_background_channel_refresh()
@@ -522,6 +524,7 @@ def execute_query(
522524
app_profile_id,
523525
request_body,
524526
prepare_metadata,
527+
self._crc32c,
525528
attempt_timeout,
526529
operation_timeout,
527530
retryable_excs=retryable_excs,

‎google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
EarlyMetadataCallError,
3434
InvalidExecuteQueryResponse,
3535
)
36+
from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
3637
from google.cloud.bigtable.data.execute_query.values import QueryResultRow
3738
from google.cloud.bigtable.data.execute_query.metadata import Metadata
3839
from google.cloud.bigtable.data.execute_query._reader import (
@@ -83,6 +84,7 @@ def __init__(
8384
app_profile_id: Optional[str],
8485
request_body: Dict[str, Any],
8586
prepare_metadata: Metadata,
87+
crc32c_implementation: _CRC32C,
8688
attempt_timeout: float | None,
8789
operation_timeout: float,
8890
req_metadata: Sequence[Tuple[str, str]] = (),
@@ -113,7 +115,7 @@ def __init__(
113115
self._instance_id = instance_id
114116
self._prepare_metadata = prepare_metadata
115117
self._final_metadata = None
116-
self._byte_cursor = _ByteCursor()
118+
self._byte_cursor = _ByteCursor(crc32c_implementation)
117119
self._reader: _Reader[QueryResultRow] = _QueryResultRowReader()
118120
self.has_received_token = False
119121
self._result_generator = self._next_impl()

‎google/cloud/bigtable/data/execute_query/_byte_cursor.py‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,9 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Optional
15+
from typing import List, Optional
1616

17+
from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
1718
from google.cloud.bigtable_v2 import ExecuteQueryResponse
1819

1920

@@ -30,9 +31,16 @@ class _ByteCursor:
3031
See :class:`google.cloud.bigtable.execute_query_reader._Reader` for more context.
3132
"""
3233

33-
def __init__(self):
34-
self._buffer = bytearray()
34+
def __init__(self, crc32c_impl: _CRC32C):
35+
self._batch_buffer = bytearray()
36+
self._batches: List[bytes] = []
3537
self._resume_token = None
38+
self._last_response_results_field = None
39+
self._crc32c = crc32c_impl
40+
41+
def reset(self):
42+
self._batch_buffer = bytearray()
43+
self._batches = []
3644

3745
def prepare_for_new_request(self):
3846
"""
@@ -50,13 +58,15 @@ def prepare_for_new_request(self):
5058
Returns:
5159
bytes: Last received resume_token.
5260
"""
53-
self._buffer = bytearray()
61+
# The first response of any retried stream will always contain reset, so
62+
# this isn't actually necessary, but we do it for safety
63+
self.reset()
5464
return self._resume_token
5565

5666
def empty(self) -> bool:
57-
return len(self._buffer) == 0
67+
return not self._batch_buffer and not self._batches
5868

59-
def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
69+
def consume(self, response: ExecuteQueryResponse) -> Optional[List[bytes]]:
6070
"""
6171
Reads results bytes from an ``ExecuteQuery`` response and adds them to a buffer.
6272
@@ -72,7 +82,8 @@ def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
7282
Response obtained from the stream.
7383
7484
Returns:
75-
bytes or None: bytes if buffers were flushed or None otherwise.
85+
bytes or None: List of bytes if buffers were flushed or None otherwise.
86+
Each element in the list represents the bytes of a `ProtoRows` message.
7687
7788
Raises:
7889
ValueError: If provided ``ExecuteQueryResponse`` is not valid
@@ -83,15 +94,32 @@ def consume(self, response: ExecuteQueryResponse) -> Optional[bytes]:
8394

8495
if response_pb.HasField("results"):
8596
results = response_pb.results
97+
if results.reset:
98+
self.reset()
8699
if results.HasField("proto_rows_batch"):
87-
self._buffer.extend(results.proto_rows_batch.batch_data)
100+
self._batch_buffer.extend(results.proto_rows_batch.batch_data)
101+
# Note that 0 is a valid checksum so we must check for field presence
102+
if results.HasField("batch_checksum"):
103+
if self._crc32c.enabled():
104+
expected_checksum = results.batch_checksum
105+
checksum = self._crc32c.checksum(self._batch_buffer)
106+
if expected_checksum != checksum:
107+
raise ValueError(
108+
f"Unexpected checksum mismatch. Expected: {expected_checksum}, got: {checksum}"
109+
)
110+
# We have a complete batch so we move it to batches and reset the
111+
# batch_buffer
112+
self._batches.append(memoryview(self._batch_buffer))
113+
self._batch_buffer = bytearray()
88114

89115
if results.resume_token:
90116
self._resume_token = results.resume_token
91117

92-
if self._buffer:
93-
return_value = memoryview(self._buffer)
94-
self._buffer = bytearray()
118+
if self._batches:
119+
if self._batch_buffer:
120+
raise ValueError("Unexpected resume_token without checksum")
121+
return_value = self._batches
122+
self._batches = []
95123
return return_value
96124
else:
97125
raise ValueError(f"Unexpected ExecuteQueryResponse: {response}")

‎google/cloud/bigtable/data/execute_query/_checksum.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
import warnings
17+
18+
19+
class _CRC32C(object):
20+
"""
21+
Wrapper around ``google_crc32c`` library that logs a warning when it is not available
22+
on the current architechure.
23+
"""
24+
25+
def __init__(self):
26+
try:
27+
import google_crc32c # type: ignore
28+
29+
self._google_crc32c = google_crc32c
30+
self.import_failed = False
31+
except ImportError:
32+
self.import_failed = True
33+
warnings.warn(
34+
"Unable to import 'google-crc32 library on this python version and architecture."
35+
" Disabling ExecuteQuery checksum validation. Please upgrade to a supported version"
36+
" and architecture if possible",
37+
category=RuntimeWarning,
38+
)
39+
40+
def enabled(self) -> bool:
41+
"""
42+
True whenever ``google_crc32c`` can be imported. False, otherwise.
43+
"""
44+
return not self.import_failed
45+
46+
def checksum(self, val: bytearray) -> int:
47+
"""
48+
Returns the crc32c checksum of the data. ``enabled`` should always be checked before calling.
49+
"""
50+
# google_crc32c wants bytes. Use a memoryview to avoid a copy
51+
memory_view = memoryview(val)
52+
return self._google_crc32c.value(bytes(memory_view))

‎google/cloud/bigtable/data/execute_query/_reader.py‎

Lines changed: 22 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
from typing import (
16+
List,
1617
TypeVar,
1718
Generic,
1819
Iterable,
@@ -54,14 +55,13 @@ class _Reader(ABC, Generic[T]):
5455

5556
@abstractmethod
5657
def consume(
57-
self, bytes_to_consume: bytes, metadata: Metadata
58+
self, batches_to_consume: List[bytes], metadata: Metadata
5859
) -> Optional[Iterable[T]]:
59-
"""This method receives a parsable chunk of bytes and returns either a None if there is
60-
not enough chunks to return to the user yet (e.g. we haven't received all columns in a
61-
row yet), or a list of appropriate values gathered from one or more parsable chunks.
62-
60+
"""This method receives a list of batches of bytes to be parsed as ProtoRows messages.
61+
It then uses the metadata to group the values in the parsed messages into rows. Returns
62+
None if batches_to_consume is empty
6363
Args:
64-
bytes_to_consume (bytes): chunk of parsable bytes received from
64+
bytes_to_consume (bytes): chunk of parsable byte batches received from
6565
:meth:`google.cloud.bigtable.byte_cursor._ByteCursor.consume`
6666
method.
6767
metadata: metadata used to transform values to rows
@@ -85,11 +85,9 @@ class _QueryResultRowReader(_Reader[QueryResultRow]):
8585
:class:`google.cloud.bigtable.byte_cursor._ByteCursor` passed in the constructor.
8686
"""
8787

88-
def __init__(self):
89-
"""
90-
Constructs new instance of ``_QueryResultRowReader``.
91-
"""
92-
self._values = []
88+
def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
89+
proto_rows = ProtoRows.pb().FromString(bytes_to_parse)
90+
return proto_rows.values
9391

9492
def _construct_query_result_row(
9593
self, values: Sequence[PBValue], metadata: ProtoMetadata
@@ -106,36 +104,23 @@ def _construct_query_result_row(
106104
result.add_field(column.column_name, parsed_value)
107105
return result
108106

109-
def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
110-
proto_rows = ProtoRows.pb().FromString(bytes_to_parse)
111-
return proto_rows.values
112-
113107
def consume(
114-
self, bytes_to_consume: bytes, metadata: Metadata
108+
self, batches_to_consume: List[bytes], metadata: Metadata
115109
) -> Optional[Iterable[QueryResultRow]]:
116-
if bytes_to_consume is None:
117-
raise ValueError("bytes_to_consume shouldn't be None")
118-
119-
self._values.extend(self._parse_proto_rows(bytes_to_consume))
120-
121-
# The logic, not defined by mypy types, ensures that the value of
122-
# "metadata" is never null at the time it is retrieved here
123110
proto_metadata = cast(ProtoMetadata, metadata)
124111
num_columns = len(proto_metadata.columns)
125-
126-
if len(self._values) < num_columns:
127-
return None
128-
129112
rows = []
130-
for batch in batched(self._values, n=num_columns):
131-
if len(batch) == num_columns:
132-
rows.append(self._construct_query_result_row(batch, proto_metadata))
133-
else:
134-
raise ValueError(
135-
"Server error, recieved bad number of values. "
136-
f"Expected {num_columns} got {len(batch)}."
137-
)
138-
139-
self._values = []
113+
for batch_bytes in batches_to_consume:
114+
values = self._parse_proto_rows(batch_bytes)
115+
for row_data in batched(values, n=num_columns):
116+
if len(row_data) == num_columns:
117+
rows.append(
118+
self._construct_query_result_row(row_data, proto_metadata)
119+
)
120+
else:
121+
raise ValueError(
122+
"Unexpected error, recieved bad number of values. "
123+
f"Expected {num_columns} got {len(row_data)}."
124+
)
140125

141126
return rows

‎google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
EarlyMetadataCallError,
2828
InvalidExecuteQueryResponse,
2929
)
30+
from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
3031
from google.cloud.bigtable.data.execute_query.values import QueryResultRow
3132
from google.cloud.bigtable.data.execute_query.metadata import Metadata
3233
from google.cloud.bigtable.data.execute_query._reader import (
@@ -59,6 +60,7 @@ def __init__(
5960
app_profile_id: Optional[str],
6061
request_body: Dict[str, Any],
6162
prepare_metadata: Metadata,
63+
crc32c_implementation: _CRC32C,
6264
attempt_timeout: float | None,
6365
operation_timeout: float,
6466
req_metadata: Sequence[Tuple[str, str]] = (),
@@ -87,7 +89,7 @@ def __init__(
8789
self._instance_id = instance_id
8890
self._prepare_metadata = prepare_metadata
8991
self._final_metadata = None
90-
self._byte_cursor = _ByteCursor()
92+
self._byte_cursor = _ByteCursor(crc32c_implementation)
9193
self._reader: _Reader[QueryResultRow] = _QueryResultRowReader()
9294
self.has_received_token = False
9395
self._result_generator = self._next_impl()

‎setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
"proto-plus >= 1.22.3, <2.0.0dev",
4545
"proto-plus >= 1.25.0, <2.0.0dev; python_version>='3.13'",
4646
"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",
47+
"google-crc32c>=1.5.0, <2.0.0dev",
4748
]
4849
extras = {"libcst": "libcst >= 0.2.5"}
4950

0 commit comments

Comments
 (0)