feat(spanner): add support for Proto Columns

google/cloud/spanner_v1/_helpers.py

@@ -18,9 +18,12 @@
 import decimal
 import math
 import time
+import base64
 
 from google.protobuf.struct_pb2 import ListValue
 from google.protobuf.struct_pb2 import Value
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 
 from google.api_core import datetime_helpers
 from google.cloud._helpers import _date_from_iso8601_date
@@ -204,6 +207,12 @@ def _make_value_pb(value):
             return Value(null_value="NULL_VALUE")
         else:
             return Value(string_value=value)
+    if isinstance(value, Message):
+        value = value.SerializeToString()
+        if value is None:
+            return Value(null_value="NULL_VALUE")
+        else:
+            return Value(string_value=base64.b64encode(value))
 
     raise ValueError("Unknown type: %s" % (value,))
 
@@ -232,7 +241,7 @@ def _make_list_value_pbs(values):
     return [_make_list_value_pb(row) for row in values]
 
 
-def _parse_value_pb(value_pb, field_type):
+def _parse_value_pb(value_pb, field_type, field_name, column_info=None):
     """Convert a Value protobuf to cell data.
 
     :type value_pb: :class:`~google.protobuf.struct_pb2.Value`
@@ -241,6 +250,12 @@ def _parse_value_pb(value_pb, field_type):
     :type field_type: :class:`~google.cloud.spanner_v1.types.Type`
     :param field_type: type code for the value
 
+    :type field_name: str
+    :param field_name: column name
+
+    :type column_info: dict
+    :param column_info: (Optional) dict of column name and column information
+
     :rtype: varies on field_type
     :returns: value extracted from value_pb
     :raises ValueError: if unknown type is passed
@@ -273,18 +288,38 @@ def _parse_value_pb(value_pb, field_type):
         return DatetimeWithNanoseconds.from_rfc3339(value_pb.string_value)
     elif type_code == TypeCode.ARRAY:
         return [
-            _parse_value_pb(item_pb, field_type.array_element_type)
+            _parse_value_pb(
+                item_pb, field_type.array_element_type, field_name, column_info
+            )
             for item_pb in value_pb.list_value.values
         ]
     elif type_code == TypeCode.STRUCT:
         return [
-            _parse_value_pb(item_pb, field_type.struct_type.fields[i].type_)
+            _parse_value_pb(
+                item_pb, field_type.struct_type.fields[i].type_, field_name, column_info
+            )
             for (i, item_pb) in enumerate(value_pb.list_value.values)
         ]
     elif type_code == TypeCode.NUMERIC:
         return decimal.Decimal(value_pb.string_value)
     elif type_code == TypeCode.JSON:
         return JsonObject.from_str(value_pb.string_value)
+    elif type_code == TypeCode.PROTO:
+        bytes_value = base64.b64decode(value_pb.string_value)
+        if column_info is not None and column_info.get(field_name) is not None:
+            default_proto_message = column_info.get(field_name)
+            if isinstance(default_proto_message, Message):
+                proto_message = type(default_proto_message)()
+                proto_message.ParseFromString(bytes_value)
+                return proto_message
+        return bytes_value
+    elif type_code == TypeCode.ENUM:
+        int_value = int(value_pb.string_value)
+        if column_info is not None and column_info.get(field_name) is not None:
+            proto_enum = column_info.get(field_name)
+            if isinstance(proto_enum, EnumTypeWrapper):
+                return proto_enum.Name(int_value)
+        return int_value
     else:
         raise ValueError("Unknown type: %s" % (field_type,))
 
@@ -305,7 +340,7 @@ def _parse_list_value_pbs(rows, row_type):
     for row in rows:
         row_data = []
         for value_pb, field in zip(row.values, row_type.fields):
-            row_data.append(_parse_value_pb(value_pb, field.type_))
+            row_data.append(_parse_value_pb(value_pb, field.type_, field.name))
         result.append(row_data)
     return result
 

google/cloud/spanner_v1/data_types.py

@@ -15,6 +15,10 @@
 """Custom data types for spanner."""
 
 import json
+import types
+
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 
 
 class JsonObject(dict):
@@ -71,3 +75,109 @@ def serialize(self):
             return json.dumps(self._array_value, sort_keys=True, separators=(",", ":"))
 
         return json.dumps(self, sort_keys=True, separators=(",", ":"))
+
+
+def _proto_message(bytes_val, proto_message_object):
+    """Helper for :func:`get_proto_message`.
+    parses serialized protocol buffer bytes data into proto message.
+
+    Args:
+        bytes_val (bytes): bytes object.
+        proto_message_object (Message): Message object for parsing
+
+    Returns:
+        Message: parses serialized protocol buffer data into this message.
+
+    Raises:
+        ValueError: if the input proto_message_object is not of type Message
+    """
+    if isinstance(bytes_val, types.NoneType):
+        return None
+
+    if not isinstance(bytes_val, bytes):
+        raise ValueError("Expected input bytes_val to be a string")
+
+    proto_message = proto_message_object.__deepcopy__()
+    proto_message.ParseFromString(bytes_val)
+    return proto_message
+
+
+def _proto_enum(int_val, proto_enum_object):
+    """Helper for :func:`get_proto_enum`.
+    parses int value into string containing the name of an enum value.
+
+    Args:
+        int_val (int): integer value.
+        proto_enum_object (EnumTypeWrapper): Enum object.
+
+    Returns:
+        str: string containing the name of an enum value.
+
+    Raises:
+        ValueError: if the input proto_enum_object is not of type EnumTypeWrapper
+    """
+    if isinstance(int_val, types.NoneType):
+        return None
+
+    if not isinstance(int_val, int):
+        raise ValueError("Expected input int_val to be a integer")
+
+    return proto_enum_object.Name(int_val)
+
+
+def get_proto_message(bytes_string, proto_message_object):
+    """parses serialized protocol buffer bytes' data or its list into proto message or list of proto message.
+
+    Args:
+        bytes_string (bytes or list[bytes]): bytes object.
+        proto_message_object (Message): Message object for parsing
+
+    Returns:
+        Message or list[Message]: parses serialized protocol buffer data into this message.
+
+    Raises:
+        ValueError: if the input proto_message_object is not of type Message
+    """
+    if isinstance(bytes_string, types.NoneType):
+        return None
+
+    if not isinstance(proto_message_object, Message):
+        raise ValueError("Input proto_message_object should be of type Message")
+
+    if not isinstance(bytes_string, (bytes, list)):
+        raise ValueError(
+            "Expected input bytes_string to be a string or list of strings"
+        )
+
+    if isinstance(bytes_string, list):
+        return [_proto_message(item, proto_message_object) for item in bytes_string]
+
+    return _proto_message(bytes_string, proto_message_object)
+
+
+def get_proto_enum(int_value, proto_enum_object):
+    """parses int or list of int values into enum or list of enum values.
+
+    Args:
+        int_value (int or list[int]): list of integer value.
+        proto_enum_object (EnumTypeWrapper): Enum object.
+
+    Returns:
+        str or list[str]: list of strings containing the name of enum value.
+
+    Raises:
+        ValueError: if the input int_list is not of type list
+    """
+    if isinstance(int_value, types.NoneType):
+        return None
+
+    if not isinstance(proto_enum_object, EnumTypeWrapper):
+        raise ValueError("Input proto_enum_object should be of type EnumTypeWrapper")
+
+    if not isinstance(int_value, (int, list)):
+        raise ValueError("Expected input int_value to be a integer or list of integers")
+
+    if isinstance(int_value, list):
+        return [_proto_enum(item, proto_enum_object) for item in int_value]
+
+    return _proto_enum(int_value, proto_enum_object)

google/cloud/spanner_v1/database.py

@@ -137,6 +137,9 @@ class Database(object):
     :type enable_drop_protection: boolean
     :param enable_drop_protection: (Optional) Represents whether the database
         has drop protection enabled or not.
+    :type proto_descriptors: bytes
+    :param proto_descriptors: (Optional) Proto descriptors used by CREATE/ALTER PROTO BUNDLE
+                              statements in 'ddl_statements' above.
     """
 
     _spanner_api = None
@@ -152,6 +155,7 @@ def __init__(
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
         enable_drop_protection=False,
+        proto_descriptors=None,
     ):
         self.database_id = database_id
         self._instance = instance
@@ -173,6 +177,7 @@ def __init__(
         self._enable_drop_protection = enable_drop_protection
         self._reconciling = False
         self._directed_read_options = self._instance._client.directed_read_options
+        self._proto_descriptors = proto_descriptors
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -382,6 +387,14 @@ def enable_drop_protection(self):
     def enable_drop_protection(self, value):
         self._enable_drop_protection = value
 
+    @property
+    def proto_descriptors(self):
+        """Proto Descriptors for this database.
+        :rtype: bytes
+        :returns: bytes representing the proto descriptors for this database
+        """
+        return self._proto_descriptors
+
     @property
     def logger(self):
         """Logger used by the database.
@@ -465,6 +478,7 @@ def create(self):
             extra_statements=list(self._ddl_statements),
             encryption_config=self._encryption_config,
             database_dialect=self._database_dialect,
+            proto_descriptors=self._proto_descriptors,
         )
         future = api.create_database(request=request, metadata=metadata)
         return future
@@ -501,6 +515,7 @@ def reload(self):
         metadata = _metadata_with_prefix(self.name)
         response = api.get_database_ddl(database=self.name, metadata=metadata)
         self._ddl_statements = tuple(response.statements)
+        self._proto_descriptors = response.proto_descriptors
         response = api.get_database(name=self.name, metadata=metadata)
         self._state = DatabasePB.State(response.state)
         self._create_time = response.create_time
@@ -514,7 +529,7 @@ def reload(self):
         self._enable_drop_protection = response.enable_drop_protection
         self._reconciling = response.reconciling
 
-    def update_ddl(self, ddl_statements, operation_id=""):
+    def update_ddl(self, ddl_statements, operation_id="", proto_descriptors=None):
         """Update DDL for this database.
 
         Apply any configured schema from :attr:`ddl_statements`.
@@ -526,6 +541,8 @@ def update_ddl(self, ddl_statements, operation_id=""):
         :param ddl_statements: a list of DDL statements to use on this database
         :type operation_id: str
         :param operation_id: (optional) a string ID for the long-running operation
+        :type proto_descriptors: bytes
+        :param proto_descriptors: (optional) Proto descriptors used by CREATE/ALTER PROTO BUNDLE statements
 
         :rtype: :class:`google.api_core.operation.Operation`
         :returns: an operation instance
@@ -539,6 +556,7 @@ def update_ddl(self, ddl_statements, operation_id=""):
             database=self.name,
             statements=ddl_statements,
             operation_id=operation_id,
+            proto_descriptors=proto_descriptors,
         )
 
         future = api.update_database_ddl(request=request, metadata=metadata)

google/cloud/spanner_v1/instance.py

@@ -435,6 +435,7 @@ def database(
         enable_drop_protection=False,
         # should be only set for tests if tests want to use interceptors
         enable_interceptors_in_tests=False,
+        proto_descriptors=None,
     ):
         """Factory to create a database within this instance.
 
@@ -478,9 +479,14 @@ def database(
         :param enable_interceptors_in_tests: (Optional) should only be set to True
             for tests if the tests want to use interceptors.
 
+        :type proto_descriptors: bytes
+        :param proto_descriptors: (Optional) Proto descriptors used by CREATE/ALTER PROTO BUNDLE
+                                  statements in 'ddl_statements' above.
+
         :rtype: :class:`~google.cloud.spanner_v1.database.Database`
         :returns: a database owned by this instance.
         """
+
         if not enable_interceptors_in_tests:
             return Database(
                 database_id,
@@ -492,6 +498,7 @@ def database(
                 database_dialect=database_dialect,
                 database_role=database_role,
                 enable_drop_protection=enable_drop_protection,
+                proto_descriptors=proto_descriptors,
             )
         else:
             return TestDatabase(

google/cloud/spanner_v1/param_types.py

@@ -18,6 +18,8 @@
 from google.cloud.spanner_v1 import TypeAnnotationCode
 from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import StructType
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 
 
 # Scalar parameter types
@@ -73,3 +75,35 @@ def Struct(fields):
     :returns: the appropriate struct-type protobuf
     """
     return Type(code=TypeCode.STRUCT, struct_type=StructType(fields=fields))
+
+
+def ProtoMessage(proto_message_object):
+    """Construct a proto message type description protobuf.
+
+    :type proto_message_object: :class:`google.protobuf.message.Message`
+    :param proto_message_object: the proto message instance
+
+    :rtype: :class:`type_pb2.Type`
+    :returns: the appropriate proto-message-type protobuf
+    """
+    if not isinstance(proto_message_object, Message):
+        raise ValueError("Expected input object of type Proto Message.")
+    return Type(
+        code=TypeCode.PROTO, proto_type_fqn=proto_message_object.DESCRIPTOR.full_name
+    )
+
+
+def ProtoEnum(proto_enum_object):
+    """Construct a proto enum type description protobuf.
+
+    :type proto_enum_object: :class:`google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper`
+    :param proto_enum_object: the proto enum instance
+
+    :rtype: :class:`type_pb2.Type`
+    :returns: the appropriate proto-enum-type protobuf
+    """
+    if not isinstance(proto_enum_object, EnumTypeWrapper):
+        raise ValueError("Expected input object of type Proto Enum")
+    return Type(
+        code=TypeCode.ENUM, proto_type_fqn=proto_enum_object.DESCRIPTOR.full_name
+    )