feat: Add support for Proto and Enum types

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

@@ -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
@@ -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.
@@ -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
+                          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:
@@ -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)
@@ -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__")

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

@@ -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
@@ -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.
@@ -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:
@@ -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)
@@ -592,6 +652,7 @@ def execute_query(
             attempt_timeout,
             operation_timeout,
             retryable_excs=retryable_excs,
+            column_info=column_info,
         )
 
     def __enter__(self):

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

@@ -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 (
@@ -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.
@@ -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
@@ -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,
@@ -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