Address review comments

trollyxia · trollyxia · commit 9132e803e592 · 2025-09-05T11:55:25.000-04:00
diff --git a/google/cloud/bigtable/data/_async/client.py b/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,7 +660,7 @@ async def execute_query(
             DeadlineExceeded,
             ServiceUnavailable,
         ),
-        column_info: dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> "ExecuteQueryIteratorAsync":
         """
         Executes an SQL query on an instance.
@@ -706,13 +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: 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.
+            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:
diff --git a/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py b/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,7 +90,7 @@ def __init__(
         operation_timeout: float,
         req_metadata: Sequence[Tuple[str, str]] = (),
         retryable_excs: Sequence[type[Exception]] = (),
-        column_info: Dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> None:
         """
         Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
diff --git a/google/cloud/bigtable/data/execute_query/_query_result_parsing_utils.py b/google/cloud/bigtable/data/execute_query/_query_result_parsing_utils.py
@@ -13,7 +13,8 @@
 # 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
@@ -43,7 +44,7 @@ def _parse_array_type(
     value: PBValue,
     metadata_type: SqlType.Array,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> Any:
     """
     used for parsing an array represented as a protobuf to a python list.
@@ -62,7 +63,7 @@ def _parse_map_type(
     value: PBValue,
     metadata_type: SqlType.Map,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> Any:
     """
     used for parsing a map represented as a protobuf to a python dict.
@@ -103,7 +104,7 @@ def _parse_struct_type(
     value: PBValue,
     metadata_type: SqlType.Struct,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> Struct:
     """
     used for parsing a struct represented as a protobuf to a
@@ -137,7 +138,7 @@ def _parse_timestamp_type(
     value: PBValue,
     metadata_type: SqlType.Timestamp,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> DatetimeWithNanoseconds:
     """
     used for parsing a timestamp represented as a protobuf to DatetimeWithNanoseconds
@@ -149,10 +150,26 @@ def _parse_proto_type(
     value: PBValue,
     metadata_type: SqlType.Proto,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> Message | bytes:
     """
-    Parses a serialized protobuf message into a Message object.
+    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
@@ -171,10 +188,27 @@ def _parse_enum_type(
     value: PBValue,
     metadata_type: SqlType.Enum,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
-) -> int | Any:
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
+) -> int | str:
     """
-    Parses an integer value into a Protobuf enum.
+    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
@@ -187,9 +221,12 @@ def _parse_enum_type(
     return value.int_value
 
 
-_TYPE_PARSERS: Dict[
-    Type[SqlType.Type], Callable[[PBValue, Any, str | None, dict[str, Any] | None], Any]
-] = {
+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,
@@ -203,7 +240,7 @@ def _parse_pb_value_to_python_value(
     value: PBValue,
     metadata_type: SqlType.Type,
     column_name: str | None,
-    column_info: dict[str, Any] | None = None,
+    column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
 ) -> Any:
     """
     used for converting the value represented as a protobufs to a python object.
diff --git a/google/cloud/bigtable/data/execute_query/_reader.py b/google/cloud/bigtable/data/execute_query/_reader.py
@@ -14,15 +14,17 @@
 from __future__ import annotations
 
 from typing import (
-    Any,
     List,
     TypeVar,
     Generic,
     Iterable,
     Optional,
     Sequence,
+    Union,
 )
 from abc import ABC, abstractmethod
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 
 from google.cloud.bigtable_v2 import ProtoRows, Value as PBValue
 
@@ -59,7 +61,7 @@ def consume(
         self,
         batches_to_consume: List[bytes],
         metadata: Metadata,
-        column_info: dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | 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
@@ -99,7 +101,7 @@ def _construct_query_result_row(
         self,
         values: Sequence[PBValue],
         metadata: Metadata,
-        column_info: dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> QueryResultRow:
         result = QueryResultRow()
         columns = metadata.columns
@@ -119,7 +121,7 @@ def consume(
         self,
         batches_to_consume: List[bytes],
         metadata: Metadata,
-        column_info: dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> Optional[Iterable[QueryResultRow]]:
         num_columns = len(metadata.columns)
         rows = []
diff --git a/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py b/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py
@@ -16,8 +16,10 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
-from typing import Any, Dict, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Optional, 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 (
     _attempt_timeout_generator,
@@ -63,7 +65,7 @@ def __init__(
         operation_timeout: float,
         req_metadata: Sequence[Tuple[str, str]] = (),
         retryable_excs: Sequence[type[Exception]] = (),
-        column_info: Dict[str, Any] | None = None,
+        column_info: dict[str, Union[Message, EnumTypeWrapper]] | None = None,
     ) -> None:
         """Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
 
