feat: Add support for Proto and Enum types

Author: trollyxia

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

@@ -657,6 +657,7 @@ async def execute_query(
             DeadlineExceeded,
             ServiceUnavailable,
         ),
+        column_info: dict[str, Any] | None = None,
     ) -> "ExecuteQueryIteratorAsync":
         """
         Executes an SQL query on an instance.
@@ -705,6 +706,13 @@ 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: 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:
             ExecuteQueryIteratorAsync: an asynchronous iterator that yields rows returned by the query
         Raises:
@@ -771,6 +779,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

@@ -87,6 +87,7 @@ def __init__(
         operation_timeout: float,
         req_metadata: Sequence[Tuple[str, str]] = (),
         retryable_excs: Sequence[type[Exception]] = (),
+        column_info: Dict[str, Any] | None = None,
     ) -> None:
         """
         Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
@@ -107,6 +108,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 +138,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 +206,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,12 @@
 # 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 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 +34,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, Any] | 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, Any] | None = None,
+) -> Any:
     """
     used for parsing a map represented as a protobuf to a python dict.
 
@@ -64,10 +80,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 +99,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, Any] | None = None,
+) -> Struct:
     """
     used for parsing a struct represented as a protobuf to a
     google.cloud.bigtable.data.execute_query.Struct
@@ -88,29 +115,96 @@ 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, Any] | 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, Any] | None = None,
+) -> Message | bytes:
+    """
+    Parses a serialized protobuf message into a Message object.
+    """
+    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, Any] | None = None,
+) -> int | Any:
+    """
+    Parses an integer value into a Protobuf enum.
+    """
+    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
+
+
+_TYPE_PARSERS: Dict[
+    Type[SqlType.Type], Callable[[PBValue, Any, str | None, dict[str, Any] | None], Any]
+] = {
     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, Any] | None = None,
+) -> Any:
     """
     used for converting the value represented as a protobufs to a python object.
     """
@@ -126,7 +220,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)

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

@@ -11,8 +11,10 @@
 # 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,
     List,
     TypeVar,
     Generic,
@@ -54,7 +56,10 @@ class _Reader(ABC, Generic[T]):
 
     @abstractmethod
     def consume(
-        self, batches_to_consume: List[bytes], metadata: Metadata
+        self,
+        batches_to_consume: List[bytes],
+        metadata: Metadata,
+        column_info: dict[str, Any] | None = None,
     ) -> Optional[Iterable[T]]:
         """This method receives a list of batches of bytes to be parsed as ProtoRows messages.
         It then uses the metadata to group the values in the parsed messages into rows. Returns
@@ -64,6 +69,8 @@ def consume(
                 :meth:`google.cloud.bigtable.byte_cursor._ByteCursor.consume`
                 method.
             metadata: metadata used to transform values to rows
+            column_info: (Optional) dict with mappings between column names and additional column information
+                for protobuf deserialization.
 
         Returns:
             Iterable[T] or None: Iterable if gathered values can form one or more instances of T,
@@ -89,7 +96,10 @@ def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
         return proto_rows.values
 
     def _construct_query_result_row(
-        self, values: Sequence[PBValue], metadata: Metadata
+        self,
+        values: Sequence[PBValue],
+        metadata: Metadata,
+        column_info: dict[str, Any] | None = None,
     ) -> QueryResultRow:
         result = QueryResultRow()
         columns = metadata.columns
@@ -99,20 +109,29 @@ def _construct_query_result_row(
         ), "This function should be called only when count of values matches count of columns."
 
         for column, value in zip(columns, values):
-            parsed_value = _parse_pb_value_to_python_value(value, column.column_type)
+            parsed_value = _parse_pb_value_to_python_value(
+                value, column.column_type, column.column_name, column_info
+            )
             result.add_field(column.column_name, parsed_value)
         return result
 
     def consume(
-        self, batches_to_consume: List[bytes], metadata: Metadata
+        self,
+        batches_to_consume: List[bytes],
+        metadata: Metadata,
+        column_info: dict[str, Any] | None = None,
     ) -> Optional[Iterable[QueryResultRow]]:
         num_columns = len(metadata.columns)
         rows = []
         for batch_bytes in batches_to_consume:
             values = self._parse_proto_rows(batch_bytes)
             for row_data in batched(values, n=num_columns):
                 if len(row_data) == num_columns:
-                    rows.append(self._construct_query_result_row(row_data, metadata))
+                    rows.append(
+                        self._construct_query_result_row(
+                            row_data, metadata, column_info
+                        )
+                    )
                 else:
                     raise ValueError(
                         "Unexpected error, recieved bad number of values. "