feat: support conversion of BigQuery schema to Protobuf

Author: Regan Koopmans

packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/__init__.py

@@ -28,7 +28,7 @@
     # this code path once we drop support for Python 3.7
     import importlib_metadata as metadata
 
-from google.cloud.bigquery_storage_v1 import client, types
+from google.cloud.bigquery_storage_v1 import client, schema, types
 
 
 class BigQueryReadClient(client.BigQueryReadClient):
@@ -140,4 +140,6 @@ def _get_version(dependency_name):
     # google.cloud.bigquery_storage_v1.client
     "BigQueryReadClient",
     "BigQueryWriteClient",
+    # google.cloud.bigquery_storage_v1.schema
+    "schema",
 )

packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/schema.py

@@ -0,0 +1,259 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Utilities for converting BigQuery schemas to Protocol Buffer descriptors.
+
+This module provides functionality to dynamically generate Protocol Buffer
+descriptors from BigQuery table schemas, eliminating the need to manually
+create and compile .proto files when using the BigQuery Storage Write API.
+"""
+
+from typing import Dict, List, Optional
+
+from google.cloud.bigquery_storage_v1 import types
+from google.protobuf import descriptor_pb2
+
+
+# Mapping from BigQuery types to Protocol Buffer field types
+_BQ_TO_PROTO_TYPE_MAP: Dict[types.TableFieldSchema.Type, int] = {
+    types.TableFieldSchema.Type.STRING: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    types.TableFieldSchema.Type.INT64: descriptor_pb2.FieldDescriptorProto.TYPE_INT64,
+    types.TableFieldSchema.Type.BOOL: descriptor_pb2.FieldDescriptorProto.TYPE_BOOL,
+    types.TableFieldSchema.Type.BYTES: descriptor_pb2.FieldDescriptorProto.TYPE_BYTES,
+    types.TableFieldSchema.Type.DOUBLE: descriptor_pb2.FieldDescriptorProto.TYPE_DOUBLE,
+    # DATE is represented as days since epoch
+    types.TableFieldSchema.Type.DATE: descriptor_pb2.FieldDescriptorProto.TYPE_INT32,
+    # DATETIME is represented as a formatted string
+    types.TableFieldSchema.Type.DATETIME: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # TIME is represented as a formatted string
+    types.TableFieldSchema.Type.TIME: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # TIMESTAMP is represented as microseconds since epoch
+    types.TableFieldSchema.Type.TIMESTAMP: descriptor_pb2.FieldDescriptorProto.TYPE_INT64,
+    # NUMERIC and BIGNUMERIC are represented as strings
+    types.TableFieldSchema.Type.NUMERIC: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    types.TableFieldSchema.Type.BIGNUMERIC: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # GEOGRAPHY is represented as WKT string
+    types.TableFieldSchema.Type.GEOGRAPHY: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # JSON is represented as a string
+    types.TableFieldSchema.Type.JSON: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # INTERVAL is represented as a string
+    types.TableFieldSchema.Type.INTERVAL: descriptor_pb2.FieldDescriptorProto.TYPE_STRING,
+    # STRUCT and RANGE use TYPE_MESSAGE and are handled specially
+    types.TableFieldSchema.Type.STRUCT: descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE,
+    types.TableFieldSchema.Type.RANGE: descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE,
+}
+
+
+def _get_field_label(mode: types.TableFieldSchema.Mode) -> int:
+    """Convert BigQuery field mode to Protocol Buffer field label.
+
+    Args:
+        mode: The BigQuery field mode (NULLABLE, REQUIRED, or REPEATED).
+
+    Returns:
+        The corresponding Protocol Buffer field label constant.
+    """
+    if mode == types.TableFieldSchema.Mode.REQUIRED:
+        return descriptor_pb2.FieldDescriptorProto.LABEL_REQUIRED
+    elif mode == types.TableFieldSchema.Mode.REPEATED:
+        return descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED
+    else:  # NULLABLE or MODE_UNSPECIFIED
+        return descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
+
+
+def _create_range_descriptor(
+    field_name: str,
+    element_type: types.TableFieldSchema.Type,
+) -> descriptor_pb2.DescriptorProto:
+    """Create a nested message descriptor for a RANGE field.
+
+    Args:
+        field_name: The name of the RANGE field.
+        element_type: The element type of the RANGE (e.g., DATE, DATETIME, TIMESTAMP).
+
+    Returns:
+        A DescriptorProto representing the RANGE structure with start and end fields.
+    """
+    range_descriptor = descriptor_pb2.DescriptorProto()
+    range_descriptor.name = f"{field_name}_Range"
+
+    # Get the proto type for the element
+    element_proto_type = _BQ_TO_PROTO_TYPE_MAP.get(element_type)
+    if element_proto_type is None:
+        raise ValueError(
+            f"Unsupported element type '{element_type}' for RANGE field '{field_name}'. "
+            f"Supported types are DATE, DATETIME, and TIMESTAMP."
+        )
+
+    # Add 'start' field
+    start_field = range_descriptor.field.add()
+    start_field.name = "start"
+    start_field.number = 1
+    start_field.type = element_proto_type
+    start_field.label = descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
+
+    # Add 'end' field
+    end_field = range_descriptor.field.add()
+    end_field.name = "end"
+    end_field.number = 2
+    end_field.type = element_proto_type
+    end_field.label = descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
+
+    return range_descriptor
+
+
+def _convert_fields_to_proto(
+    fields: List[types.TableFieldSchema],
+) -> tuple[List[descriptor_pb2.FieldDescriptorProto], List[descriptor_pb2.DescriptorProto]]:
+    """Convert BigQuery fields to Protocol Buffer field descriptors.
+
+    Args:
+        fields: List of BigQuery table field schemas.
+
+    Returns:
+        A tuple of (field_descriptors, nested_descriptors):
+        - field_descriptors: List of FieldDescriptorProto objects
+        - nested_descriptors: List of nested DescriptorProto objects for STRUCT/RANGE types
+    """
+    field_descriptors = []
+    nested_descriptors = []
+
+    for field_number, bq_field in enumerate(fields, start=1):
+        field_descriptor = descriptor_pb2.FieldDescriptorProto()
+        field_descriptor.name = bq_field.name
+        field_descriptor.number = field_number
+
+        # Handle STRUCT fields (nested messages)
+        if bq_field.type_ == types.TableFieldSchema.Type.STRUCT:
+            field_descriptor.type = descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE
+
+            # Create nested message descriptor
+            nested_descriptor = descriptor_pb2.DescriptorProto()
+            nested_descriptor.name = f"{bq_field.name}_Struct"
+
+            # Recursively convert nested fields
+            nested_fields, deeply_nested = _convert_fields_to_proto(
+                bq_field.fields
+            )
+
+            for nested_field in nested_fields:
+                nested_descriptor.field.append(nested_field)
+
+            for deeply_nested_type in deeply_nested:
+                nested_descriptor.nested_type.append(deeply_nested_type)
+
+            nested_descriptors.append(nested_descriptor)
+            field_descriptor.type_name = nested_descriptor.name
+
+        # Handle RANGE fields
+        elif bq_field.type_ == types.TableFieldSchema.Type.RANGE:
+            field_descriptor.type = descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE
+
+            # Get the element type from range_element_type
+            if not bq_field.range_element_type or not bq_field.range_element_type.type_:
+                raise ValueError(
+                    f"RANGE field '{bq_field.name}' is missing range_element_type. "
+                    f"RANGE fields must specify an element type (DATE, DATETIME, or TIMESTAMP)."
+                )
+            element_type = bq_field.range_element_type.type_
+
+            range_descriptor = _create_range_descriptor(bq_field.name, element_type)
+            nested_descriptors.append(range_descriptor)
+            field_descriptor.type_name = range_descriptor.name
+
+        # Handle primitive types
+        else:
+            proto_type = _BQ_TO_PROTO_TYPE_MAP.get(bq_field.type_)
+            if proto_type is None:
+                raise ValueError(
+                    f"Unsupported BigQuery type: {bq_field.type_} for field {bq_field.name}"
+                )
+            field_descriptor.type = proto_type
+
+        # Set field label based on mode
+        mode = bq_field.mode or types.TableFieldSchema.Mode.NULLABLE
+        field_descriptor.label = _get_field_label(mode)
+
+        field_descriptors.append(field_descriptor)
+
+    return field_descriptors, nested_descriptors
+
+
+def table_schema_to_proto_descriptor(
+    table_schema: types.TableSchema,
+    message_name: str = "TableRow",
+) -> descriptor_pb2.DescriptorProto:
+    """Convert a BigQuery TableSchema to a Protocol Buffer DescriptorProto.
+
+    This function generates a Protocol Buffer descriptor that can be used with
+    the BigQuery Storage Write API without needing to create and compile .proto
+    files. The generated descriptor uses proto2 wire format, which is required
+    by the Write API.
+
+    Args:
+        table_schema: The BigQuery table schema to convert.
+        message_name: Optional name for the root message type. Defaults to "TableRow".
+
+    Returns:
+        A DescriptorProto that can be used with ProtoSchema in the Write API.
+
+    Raises:
+        ValueError: If the schema contains unsupported field types.
+
+    Example:
+        >>> from google.cloud.bigquery_storage_v1 import schema, types
+        >>>
+        >>> # Define a BigQuery schema
+        >>> table_schema = types.TableSchema(fields=[
+        ...     types.TableFieldSchema(
+        ...         name="id",
+        ...         type_=types.TableFieldSchema.Type.INT64,
+        ...         mode=types.TableFieldSchema.Mode.REQUIRED
+        ...     ),
+        ...     types.TableFieldSchema(
+        ...         name="name",
+        ...         type_=types.TableFieldSchema.Type.STRING
+        ...     ),
+        ... ])
+        >>>
+        >>> # Convert to proto descriptor
+        >>> proto_descriptor = schema.table_schema_to_proto_descriptor(table_schema)
+        >>>
+        >>> # Use with Write API
+        >>> proto_schema = types.ProtoSchema()
+        >>> proto_schema.proto_descriptor = proto_descriptor
+
+    Note:
+        For detailed information about BigQuery to Protocol Buffer type mappings,
+        see: https://cloud.google.com/bigquery/docs/write-api#data_type_conversions
+    """
+    # Create the root descriptor
+    descriptor = descriptor_pb2.DescriptorProto()
+    descriptor.name = message_name
+
+    # Convert fields
+    field_descriptors, nested_descriptors = _convert_fields_to_proto(table_schema.fields)
+
+    # Add field descriptors to the message
+    for field_descriptor in field_descriptors:
+        descriptor.field.append(field_descriptor)
+
+    # Add nested type descriptors
+    for nested_descriptor in nested_descriptors:
+        descriptor.nested_type.append(nested_descriptor)
+
+    return descriptor
+
+
+__all__ = ("table_schema_to_proto_descriptor",)