feat: Add support for Proto and Enum types

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

@@ -23,6 +23,7 @@
     Set,
     Sequence,
     TYPE_CHECKING,
+    Union,
 )
 
 import abc
@@ -58,6 +59,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 +660,7 @@ async def execute_query(
             DeadlineExceeded,
             ServiceUnavailable,
         ),
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> "ExecuteQueryIteratorAsync":
         """
         Executes an SQL query on an instance.
@@ -705,6 +709,54 @@ 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, or if deserialization fails:
+                    - 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.
+
+                    Examples:
+                        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,
+                        }
+
         Returns:
             ExecuteQueryIteratorAsync: an asynchronous iterator that yields rows returned by the query
         Raises:
@@ -771,6 +823,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

@@ -485,6 +485,7 @@ def execute_query(
             DeadlineExceeded,
             ServiceUnavailable,
         ),
+        column_info: dict[str, Any] | None = None,
     ) -> "ExecuteQueryIterator":
         """Executes an SQL query on an instance.
         Returns an iterator to asynchronously stream back columns from selected rows.
@@ -532,6 +533,13 @@ 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: Dictionary with mappings between column names and additional column information.
+                An object where column names as keys and custom objects as corresponding
+                values for deserialization. It's specifically useful for data types like
+                protobuf where deserialization logic is on user-specific code. When provided,
+                the custom object enables deserialization of backend-received column data.
+                If not provided, data remains serialized as bytes for Proto Messages and
+                integer for Proto Enums.
         Returns:
             ExecuteQueryIterator: an asynchronous iterator that yields rows returned by the query
         Raises:
@@ -592,6 +600,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

@@ -21,8 +21,11 @@
     Sequence,
     Tuple,
     TYPE_CHECKING,
+    Union,
 )
 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 +90,7 @@ def __init__(
         operation_timeout: float,
         req_metadata: Sequence[Tuple[str, str]] = (),
         retryable_excs: Sequence[type[Exception]] = (),
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> None:
         """
         Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
@@ -107,6 +111,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 +141,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 +209,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
 

google/cloud/bigtable/data/execute_query/_query_result_parsing_utils.py

@@ -11,8 +11,13 @@
 # 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.
+from __future__ import annotations
 
-from typing import Any, Callable, Dict, Type
+from typing import Any, Callable, Dict, Type, Union, Optional
+from typing_extensions import TypeAlias
+
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 from google.cloud.bigtable.data.execute_query.values import Struct
 from google.cloud.bigtable.data.execute_query.metadata import SqlType
 from google.cloud.bigtable_v2 import Value as PBValue
@@ -30,24 +35,36 @@
     SqlType.Struct: "array_value",
     SqlType.Array: "array_value",
     SqlType.Map: "array_value",
+    SqlType.Proto: "bytes_value",
+    SqlType.Enum: "int_value",
 }
 
 
-def _parse_array_type(value: PBValue, metadata_type: SqlType.Array) -> Any:
+def _parse_array_type(
+    value: PBValue,
+    metadata_type: SqlType.Array,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> Any:
     """
     used for parsing an array represented as a protobuf to a python list.
     """
     return list(
         map(
             lambda val: _parse_pb_value_to_python_value(
-                val, metadata_type.element_type
+                val, metadata_type.element_type, column_name, column_info
             ),
             value.array_value.values,
         )
     )
 
 
-def _parse_map_type(value: PBValue, metadata_type: SqlType.Map) -> Any:
+def _parse_map_type(
+    value: PBValue,
+    metadata_type: SqlType.Map,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> Any:
     """
     used for parsing a map represented as a protobuf to a python dict.
 
@@ -64,10 +81,16 @@ def _parse_map_type(value: PBValue, metadata_type: SqlType.Map) -> Any:
             map(
                 lambda map_entry: (
                     _parse_pb_value_to_python_value(
-                        map_entry.array_value.values[0], metadata_type.key_type
+                        map_entry.array_value.values[0],
+                        metadata_type.key_type,
+                        f"{column_name}.key" if column_name is not None else None,
+                        column_info,
                     ),
                     _parse_pb_value_to_python_value(
-                        map_entry.array_value.values[1], metadata_type.value_type
+                        map_entry.array_value.values[1],
+                        metadata_type.value_type,
+                        f"{column_name}.value" if column_name is not None else None,
+                        column_info,
                     ),
                 ),
                 value.array_value.values,
@@ -77,7 +100,12 @@ def _parse_map_type(value: PBValue, metadata_type: SqlType.Map) -> Any:
         raise ValueError("Invalid map entry - less or more than two values.")
 
 
-def _parse_struct_type(value: PBValue, metadata_type: SqlType.Struct) -> Struct:
+def _parse_struct_type(
+    value: PBValue,
+    metadata_type: SqlType.Struct,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> Struct:
     """
     used for parsing a struct represented as a protobuf to a
     google.cloud.bigtable.data.execute_query.Struct
@@ -88,29 +116,132 @@ def _parse_struct_type(value: PBValue, metadata_type: SqlType.Struct) -> Struct:
     struct = Struct()
     for value, field in zip(value.array_value.values, metadata_type.fields):
         field_name, field_type = field
-        struct.add_field(field_name, _parse_pb_value_to_python_value(value, field_type))
+        nested_column_name: str | None
+        if column_name is None:
+            nested_column_name = None
+        else:
+            # qualify the column name for nested lookups
+            nested_column_name = (
+                f"{column_name}.{field_name}" if field_name else column_name
+            )
+        struct.add_field(
+            field_name,
+            _parse_pb_value_to_python_value(
+                value, field_type, nested_column_name, column_info
+            ),
+        )
 
     return struct
 
 
 def _parse_timestamp_type(
-    value: PBValue, metadata_type: SqlType.Timestamp
+    value: PBValue,
+    metadata_type: SqlType.Timestamp,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> DatetimeWithNanoseconds:
     """
     used for parsing a timestamp represented as a protobuf to DatetimeWithNanoseconds
     """
     return DatetimeWithNanoseconds.from_timestamp_pb(value.timestamp_value)
 
 
-_TYPE_PARSERS: Dict[Type[SqlType.Type], Callable[[PBValue, Any], Any]] = {
+def _parse_proto_type(
+    value: PBValue,
+    metadata_type: SqlType.Proto,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> Message | bytes:
+    """
+    Parses a serialized protobuf message into a Message object using type information
+    provided in column_info.
+
+    Args:
+        value: The value to parse, expected to have a bytes_value attribute.
+        metadata_type: The expected SQL type (Proto).
+        column_name: The name of the column.
+        column_info: (Optional) A dictionary mapping column names to their
+            corresponding Protobuf Message classes. This information is used
+            to deserialize the raw bytes.
+
+    Returns:
+        A deserialized Protobuf Message object if parsing is successful.
+        If parsing fails for any reason, or if the required type information
+        is not found in column_info, the function returns the original
+        serialized data as bytes (value.bytes_value). This fallback
+        ensures that the raw data is still accessible.
+    """
+    if (
+        column_name is not None
+        and column_info is not None
+        and column_info.get(column_name) is not None
+    ):
+        default_proto_message = column_info.get(column_name)
+        if isinstance(default_proto_message, Message):
+            proto_message = type(default_proto_message)()
+            proto_message.ParseFromString(value.bytes_value)
+            return proto_message
+    return value.bytes_value
+
+
+def _parse_enum_type(
+    value: PBValue,
+    metadata_type: SqlType.Enum,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> int | str:
+    """
+    Parses an integer value into a Protobuf enum name string using type information
+    provided in column_info.
+
+    Args:
+        value: The value to parse, expected to have an int_value attribute.
+        metadata_type: The expected SQL type (Enum).
+        column_name: The name of the column.
+        column_info: (Optional) A dictionary mapping column names to their
+            corresponding Protobuf EnumTypeWrapper objects. This information
+            is used to convert the integer to an enum name.
+
+    Returns:
+        A string representing the name of the enum value if conversion is successful.
+        If conversion fails for any reason, such as the required EnumTypeWrapper
+        not being found in column_info, or if an error occurs during the name lookup
+        (e.g., the integer is not a valid enum value), the function returns the
+        original integer value (value.int_value). This fallback ensures the
+        raw integer representation is still accessible.
+    """
+    if (
+        column_name is not None
+        and column_info is not None
+        and column_info.get(column_name) is not None
+    ):
+        proto_enum = column_info.get(column_name)
+        if isinstance(proto_enum, EnumTypeWrapper):
+            return proto_enum.Name(value.int_value)
+    return value.int_value
+
+
+ParserCallable: TypeAlias = Callable[
+    [PBValue, Any, Optional[str], Optional[Dict[str, Union[Message, EnumTypeWrapper]]]],
+    Any,
+]
+
+_TYPE_PARSERS: Dict[Type[SqlType.Type], ParserCallable] = {
     SqlType.Timestamp: _parse_timestamp_type,
     SqlType.Struct: _parse_struct_type,
     SqlType.Array: _parse_array_type,
     SqlType.Map: _parse_map_type,
+    SqlType.Proto: _parse_proto_type,
+    SqlType.Enum: _parse_enum_type,
 }
 
 
-def _parse_pb_value_to_python_value(value: PBValue, metadata_type: SqlType.Type) -> Any:
+def _parse_pb_value_to_python_value(
+    value: PBValue,
+    metadata_type: SqlType.Type,
+    column_name: str | None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> Any:
     """
     used for converting the value represented as a protobufs to a python object.
     """
@@ -126,7 +257,7 @@ def _parse_pb_value_to_python_value(value: PBValue, metadata_type: SqlType.Type)
 
     if kind in _TYPE_PARSERS:
         parser = _TYPE_PARSERS[kind]
-        return parser(value, metadata_type)
+        return parser(value, metadata_type, column_name, column_info)
     elif kind in _REQUIRED_PROTO_FIELDS:
         field_name = _REQUIRED_PROTO_FIELDS[kind]
         return getattr(value, field_name)