Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: Add support for Proto and Enum types #1202

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
daniel-sanche merged 7 commits into from
Sep 12, 2025
Merged

feat: Add support for Proto and Enum types #1202

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c85c31e
feat: Add support for Proto and Enum types
trollyxia Aug 28, 2025
9132e80
Address review comments
trollyxia Sep 5, 2025
f65d45a
Minor change, address comments
trollyxia Sep 8, 2025
e1fac0c
Update docstring
trollyxia Sep 8, 2025
13236e6
Address review comments
trollyxia Sep 10, 2025
d785b8b
Modify docstring to indicate message.DecodeError when proto parsing f…
trollyxia Sep 11, 2025
7dea490
Disable proto/enum parsing for unnamed proto/enum fields inside a struct
trollyxia Sep 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 61 additions & 0 deletions google/cloud/bigtable/data/_async/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@
from google.api_core.exceptions import DeadlineExceeded
from google.api_core.exceptions import ServiceUnavailable
from google.api_core.exceptions import Aborted
from google.protobuf.message import Message
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper

import google.auth.credentials
import google.auth._default
Expand Down Expand Up @@ -657,6 +659,7 @@ async def execute_query(
DeadlineExceeded,
ServiceUnavailable,
),
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
) -> "ExecuteQueryIteratorAsync":
"""
Executes an SQL query on an instance.
Expand Down Expand Up @@ -705,6 +708,62 @@ async def execute_query(
If None, defaults to prepare_operation_timeout.
prepare_retryable_errors: a list of errors that will be retried if encountered during prepareQuery.
Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable)
column_info: (Optional) A dictionary mapping column names to Protobuf message classes or EnumTypeWrapper objects.
This dictionary provides the necessary type information for deserializing PROTO and
ENUM column values from the query results. When an entry is provided
for a PROTO or ENUM column, the client library will attempt to deserialize the raw data.

- For PROTO columns: The value in the dictionary should be the
Protobuf Message class (e.g., ``my_pb2.MyMessage``).
- For ENUM columns: The value should be the Protobuf EnumTypeWrapper
object (e.g., ``my_pb2.MyEnum``).

Example::

import my_pb2

column_info = {
"my_proto_column": my_pb2.MyMessage,
"my_enum_column": my_pb2.MyEnum
}

If ``column_info`` is not provided, or if a specific column name is not found
in the dictionary:

- PROTO columns will be returned as raw bytes.
- ENUM columns will be returned as integers.

Note for Nested PROTO or ENUM Fields:

To specify types for PROTO or ENUM fields within STRUCTs or MAPs, use a dot-separated
path from the top-level column name.

- For STRUCTs: ``struct_column_name.field_name``
- For MAPs: ``map_column_name.key`` or ``map_column_name.value`` to specify types
Copy link
Contributor

@jackdingilian jackdingilian Sep 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the with_history maps that have nested struct values I assume this doesn't work? Or does it handle arbitrary nesting?

Copy link
Contributor

@jackdingilian jackdingilian Sep 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah cool it looks like this works. Can you add an example of that to the doc as well?

Copy link
Member Author

@trollyxia trollyxia Sep 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done!

for the map keys or values, respectively.

Example::

import my_pb2

column_info = {
# Top-level column
"my_proto_column": my_pb2.MyMessage,
"my_enum_column": my_pb2.MyEnum,

# Nested field in a STRUCT column named 'my_struct'
"my_struct.nested_proto_field": my_pb2.OtherMessage,
"my_struct.nested_enum_field": my_pb2.AnotherEnum,

# Nested field in a MAP column named 'my_map'
"my_map.key": my_pb2.MapKeyEnum, # If map keys were enums
"my_map.value": my_pb2.MapValueMessage,

# PROTO field inside a STRUCT, where the STRUCT is the value in a MAP column
"struct_map.value.nested_proto_field": my_pb2.DeeplyNestedProto,
"struct_map.value.nested_enum_field": my_pb2.DeeplyNestedEnum
}

Returns:
ExecuteQueryIteratorAsync: an asynchronous iterator that yields rows returned by the query
Raises:
Expand All @@ -714,6 +773,7 @@ async def execute_query(
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
a parameter is passed without an explicit type, and the type cannot be infered
google.protobuf.message.DecodeError: raised if the deserialization of a PROTO/ENUM value fails.
"""
instance_name = self._gapic_client.instance_path(self.project, instance_id)
converted_param_types = _to_param_types(parameters, parameter_types)
Expand Down Expand Up @@ -771,6 +831,7 @@ async def execute_query(
attempt_timeout,
operation_timeout,
retryable_excs=retryable_excs,
column_info=column_info,
)

@CrossSync.convert(sync_name="__enter__")
Expand Down
61 changes: 61 additions & 0 deletions google/cloud/bigtable/data/_sync_autogen/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@
from google.api_core.exceptions import DeadlineExceeded
from google.api_core.exceptions import ServiceUnavailable
from google.api_core.exceptions import Aborted
from google.protobuf.message import Message
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
import google.auth.credentials
import google.auth._default
from google.api_core import client_options as client_options_lib
Expand Down Expand Up @@ -485,6 +487,7 @@ def execute_query(
DeadlineExceeded,
ServiceUnavailable,
),
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
) -> "ExecuteQueryIterator":
"""Executes an SQL query on an instance.
Returns an iterator to asynchronously stream back columns from selected rows.
Expand Down Expand Up @@ -532,6 +535,62 @@ def execute_query(
If None, defaults to prepare_operation_timeout.
prepare_retryable_errors: a list of errors that will be retried if encountered during prepareQuery.
Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable)
column_info: (Optional) A dictionary mapping column names to Protobuf message classes or EnumTypeWrapper objects.
This dictionary provides the necessary type information for deserializing PROTO and
ENUM column values from the query results. When an entry is provided
for a PROTO or ENUM column, the client library will attempt to deserialize the raw data.

- For PROTO columns: The value in the dictionary should be the
Protobuf Message class (e.g., ``my_pb2.MyMessage``).
- For ENUM columns: The value should be the Protobuf EnumTypeWrapper
object (e.g., ``my_pb2.MyEnum``).

Example::

import my_pb2

column_info = {
"my_proto_column": my_pb2.MyMessage,
"my_enum_column": my_pb2.MyEnum
}

If ``column_info`` is not provided, or if a specific column name is not found
in the dictionary:

- PROTO columns will be returned as raw bytes.
- ENUM columns will be returned as integers.

Note for Nested PROTO or ENUM Fields:

To specify types for PROTO or ENUM fields within STRUCTs or MAPs, use a dot-separated
path from the top-level column name.

- For STRUCTs: ``struct_column_name.field_name``
- For MAPs: ``map_column_name.key`` or ``map_column_name.value`` to specify types
for the map keys or values, respectively.

Example::

import my_pb2

column_info = {
# Top-level column
"my_proto_column": my_pb2.MyMessage,
"my_enum_column": my_pb2.MyEnum,

# Nested field in a STRUCT column named 'my_struct'
"my_struct.nested_proto_field": my_pb2.OtherMessage,
"my_struct.nested_enum_field": my_pb2.AnotherEnum,

# Nested field in a MAP column named 'my_map'
"my_map.key": my_pb2.MapKeyEnum, # If map keys were enums
"my_map.value": my_pb2.MapValueMessage,

# PROTO field inside a STRUCT, where the STRUCT is the value in a MAP column
"struct_map.value.nested_proto_field": my_pb2.DeeplyNestedProto,
"struct_map.value.nested_enum_field": my_pb2.DeeplyNestedEnum
}

Returns:
ExecuteQueryIterator: an asynchronous iterator that yields rows returned by the query
Raises:
Expand All @@ -541,6 +600,7 @@ def execute_query(
google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
a parameter is passed without an explicit type, and the type cannot be infered
google.protobuf.message.DecodeError: raised if the deserialization of a PROTO/ENUM value fails.
"""
instance_name = self._gapic_client.instance_path(self.project, instance_id)
converted_param_types = _to_param_types(parameters, parameter_types)
Expand Down Expand Up @@ -592,6 +652,7 @@ def execute_query(
attempt_timeout,
operation_timeout,
retryable_excs=retryable_excs,
column_info=column_info,
)

def __enter__(self):
Expand Down
10 changes: 9 additions & 1 deletion google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
TYPE_CHECKING,
)
from google.api_core import retry as retries
from google.protobuf.message import Message
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper

from google.cloud.bigtable.data.execute_query._byte_cursor import _ByteCursor
from google.cloud.bigtable.data._helpers import (
Expand Down Expand Up @@ -87,6 +89,7 @@ def __init__(
operation_timeout: float,
req_metadata: Sequence[Tuple[str, str]] = (),
retryable_excs: Sequence[type[Exception]] = (),
column_info: dict[str, Message | EnumTypeWrapper] | None = None,
) -> None:
"""
Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
Expand All @@ -107,6 +110,8 @@ def __init__(
Failed requests will be retried within the budget
req_metadata: metadata used while sending the gRPC request
retryable_excs: a list of errors that will be retried if encountered.
column_info: dict with mappings between column names and additional column information
for protobuf deserialization.
Raises:
{NO_LOOP}
:class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
Expand Down Expand Up @@ -135,6 +140,7 @@ def __init__(
exception_factory=_retry_exception_factory,
)
self._req_metadata = req_metadata
self._column_info = column_info
try:
self._register_instance_task = CrossSync.create_task(
self._client._register_instance,
Expand Down Expand Up @@ -202,7 +208,9 @@ async def _next_impl(self) -> CrossSync.Iterator[QueryResultRow]:
raise ValueError(
"Error parsing response before finalizing metadata"
)
results = self._reader.consume(batches_to_parse, self.metadata)
results = self._reader.consume(
batches_to_parse, self.metadata, self._column_info
)
if results is None:
continue

Expand Down
Loading
Loading