Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 34ceb86

Browse files
trollyxiatrollyxia
authored
feat: Add support for Proto and Enum types (#1202)
1 parent 6fa3008 commit 34ceb86

File tree

16 files changed

+1462
-61
lines changed

16 files changed

+1462
-61
lines changed

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@
5858
from google.api_core.exceptions import DeadlineExceeded
5959
from google.api_core.exceptions import ServiceUnavailable
6060
from google.api_core.exceptions import Aborted
61+
from google.protobuf.message import Message
62+
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
6163

6264
import google.auth.credentials
6365
import google.auth._default
@@ -684,6 +686,7 @@ async def execute_query(
684686
DeadlineExceeded,
685687
ServiceUnavailable,
686688
),
689+
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
687690
) -> "ExecuteQueryIteratorAsync":
688691
"""
689692
Executes an SQL query on an instance.
@@ -732,6 +735,62 @@ async def execute_query(
732735
If None, defaults to prepare_operation_timeout.
733736
prepare_retryable_errors: a list of errors that will be retried if encountered during prepareQuery.
734737
Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable)
738+
column_info: (Optional) A dictionary mapping column names to Protobuf message classes or EnumTypeWrapper objects.
739+
This dictionary provides the necessary type information for deserializing PROTO and
740+
ENUM column values from the query results. When an entry is provided
741+
for a PROTO or ENUM column, the client library will attempt to deserialize the raw data.
742+
743+
- For PROTO columns: The value in the dictionary should be the
744+
Protobuf Message class (e.g., ``my_pb2.MyMessage``).
745+
- For ENUM columns: The value should be the Protobuf EnumTypeWrapper
746+
object (e.g., ``my_pb2.MyEnum``).
747+
748+
Example::
749+
750+
import my_pb2
751+
752+
column_info = {
753+
"my_proto_column": my_pb2.MyMessage,
754+
"my_enum_column": my_pb2.MyEnum
755+
}
756+
757+
If ``column_info`` is not provided, or if a specific column name is not found
758+
in the dictionary:
759+
760+
- PROTO columns will be returned as raw bytes.
761+
- ENUM columns will be returned as integers.
762+
763+
Note for Nested PROTO or ENUM Fields:
764+
765+
To specify types for PROTO or ENUM fields within STRUCTs or MAPs, use a dot-separated
766+
path from the top-level column name.
767+
768+
- For STRUCTs: ``struct_column_name.field_name``
769+
- For MAPs: ``map_column_name.key`` or ``map_column_name.value`` to specify types
770+
for the map keys or values, respectively.
771+
772+
Example::
773+
774+
import my_pb2
775+
776+
column_info = {
777+
# Top-level column
778+
"my_proto_column": my_pb2.MyMessage,
779+
"my_enum_column": my_pb2.MyEnum,
780+
781+
# Nested field in a STRUCT column named 'my_struct'
782+
"my_struct.nested_proto_field": my_pb2.OtherMessage,
783+
"my_struct.nested_enum_field": my_pb2.AnotherEnum,
784+
785+
# Nested field in a MAP column named 'my_map'
786+
"my_map.key": my_pb2.MapKeyEnum, # If map keys were enums
787+
"my_map.value": my_pb2.MapValueMessage,
788+
789+
# PROTO field inside a STRUCT, where the STRUCT is the value in a MAP column
790+
"struct_map.value.nested_proto_field": my_pb2.DeeplyNestedProto,
791+
"struct_map.value.nested_enum_field": my_pb2.DeeplyNestedEnum
792+
}
793+
735794
Returns:
736795
ExecuteQueryIteratorAsync: an asynchronous iterator that yields rows returned by the query
737796
Raises:
@@ -741,6 +800,7 @@ async def execute_query(
741800
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
742801
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
743802
a parameter is passed without an explicit type, and the type cannot be infered
803+
google.protobuf.message.DecodeError: raised if the deserialization of a PROTO/ENUM value fails.
744804
"""
745805
instance_name = self._gapic_client.instance_path(self.project, instance_id)
746806
converted_param_types = _to_param_types(parameters, parameter_types)
@@ -798,6 +858,7 @@ async def execute_query(
798858
attempt_timeout,
799859
operation_timeout,
800860
retryable_excs=retryable_excs,
861+
column_info=column_info,
801862
)
802863

803864
@CrossSync.convert(sync_name="__enter__")

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@
4949
from google.api_core.exceptions import DeadlineExceeded
5050
from google.api_core.exceptions import ServiceUnavailable
5151
from google.api_core.exceptions import Aborted
52+
from google.protobuf.message import Message
53+
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
5254
import google.auth.credentials
5355
import google.auth._default
5456
from google.api_core import client_options as client_options_lib
@@ -508,6 +510,7 @@ def execute_query(
508510
DeadlineExceeded,
509511
ServiceUnavailable,
510512
),
513+
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
511514
) -> "ExecuteQueryIterator":
512515
"""Executes an SQL query on an instance.
513516
Returns an iterator to asynchronously stream back columns from selected rows.
@@ -555,6 +558,62 @@ def execute_query(
555558
If None, defaults to prepare_operation_timeout.
556559
prepare_retryable_errors: a list of errors that will be retried if encountered during prepareQuery.
557560
Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable)
561+
column_info: (Optional) A dictionary mapping column names to Protobuf message classes or EnumTypeWrapper objects.
562+
This dictionary provides the necessary type information for deserializing PROTO and
563+
ENUM column values from the query results. When an entry is provided
564+
for a PROTO or ENUM column, the client library will attempt to deserialize the raw data.
565+
566+
- For PROTO columns: The value in the dictionary should be the
567+
Protobuf Message class (e.g., ``my_pb2.MyMessage``).
568+
- For ENUM columns: The value should be the Protobuf EnumTypeWrapper
569+
object (e.g., ``my_pb2.MyEnum``).
570+
571+
Example::
572+
573+
import my_pb2
574+
575+
column_info = {
576+
"my_proto_column": my_pb2.MyMessage,
577+
"my_enum_column": my_pb2.MyEnum
578+
}
579+
580+
If ``column_info`` is not provided, or if a specific column name is not found
581+
in the dictionary:
582+
583+
- PROTO columns will be returned as raw bytes.
584+
- ENUM columns will be returned as integers.
585+
586+
Note for Nested PROTO or ENUM Fields:
587+
588+
To specify types for PROTO or ENUM fields within STRUCTs or MAPs, use a dot-separated
589+
path from the top-level column name.
590+
591+
- For STRUCTs: ``struct_column_name.field_name``
592+
- For MAPs: ``map_column_name.key`` or ``map_column_name.value`` to specify types
593+
for the map keys or values, respectively.
594+
595+
Example::
596+
597+
import my_pb2
598+
599+
column_info = {
600+
# Top-level column
601+
"my_proto_column": my_pb2.MyMessage,
602+
"my_enum_column": my_pb2.MyEnum,
603+
604+
# Nested field in a STRUCT column named 'my_struct'
605+
"my_struct.nested_proto_field": my_pb2.OtherMessage,
606+
"my_struct.nested_enum_field": my_pb2.AnotherEnum,
607+
608+
# Nested field in a MAP column named 'my_map'
609+
"my_map.key": my_pb2.MapKeyEnum, # If map keys were enums
610+
"my_map.value": my_pb2.MapValueMessage,
611+
612+
# PROTO field inside a STRUCT, where the STRUCT is the value in a MAP column
613+
"struct_map.value.nested_proto_field": my_pb2.DeeplyNestedProto,
614+
"struct_map.value.nested_enum_field": my_pb2.DeeplyNestedEnum
615+
}
616+
558617
Returns:
559618
ExecuteQueryIterator: an asynchronous iterator that yields rows returned by the query
560619
Raises:
@@ -564,6 +623,7 @@ def execute_query(
564623
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
565624
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
566625
a parameter is passed without an explicit type, and the type cannot be infered
626+
google.protobuf.message.DecodeError: raised if the deserialization of a PROTO/ENUM value fails.
567627
"""
568628
instance_name = self._gapic_client.instance_path(self.project, instance_id)
569629
converted_param_types = _to_param_types(parameters, parameter_types)
@@ -615,6 +675,7 @@ def execute_query(
615675
attempt_timeout,
616676
operation_timeout,
617677
retryable_excs=retryable_excs,
678+
column_info=column_info,
618679
)
619680

620681
def __enter__(self):

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
TYPE_CHECKING,
2424
)
2525
from google.api_core import retry as retries
26+
from google.protobuf.message import Message
27+
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
2628

2729
from google.cloud.bigtable.data.execute_query._byte_cursor import _ByteCursor
2830
from google.cloud.bigtable.data._helpers import (
@@ -87,6 +89,7 @@ def __init__(
8789
operation_timeout: float,
8890
req_metadata: Sequence[Tuple[str, str]] = (),
8991
retryable_excs: Sequence[type[Exception]] = (),
92+
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
9093
) -> None:
9194
"""
9295
Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
@@ -107,6 +110,8 @@ def __init__(
107110
Failed requests will be retried within the budget
108111
req_metadata: metadata used while sending the gRPC request
109112
retryable_excs: a list of errors that will be retried if encountered.
113+
column_info: dict with mappings between column names and additional column information
114+
for protobuf deserialization.
110115
Raises:
111116
{NO_LOOP}
112117
:class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
@@ -135,6 +140,7 @@ def __init__(
135140
exception_factory=_retry_exception_factory,
136141
)
137142
self._req_metadata = req_metadata
143+
self._column_info = column_info
138144
try:
139145
self._register_instance_task = CrossSync.create_task(
140146
self._client._register_instance,
@@ -202,7 +208,9 @@ async def _next_impl(self) -> CrossSync.Iterator[QueryResultRow]:
202208
raise ValueError(
203209
"Error parsing response before finalizing metadata"
204210
)
205-
results = self._reader.consume(batches_to_parse, self.metadata)
211+
results = self._reader.consume(
212+
batches_to_parse, self.metadata, self._column_info
213+
)
206214
if results is None:
207215
continue
208216

0 commit comments

Comments
 (0)