feat: Support for Proto Messages & Enums (#2155)

* feat: Importing Proto Changes

Commit will be reverted, once PROTO changes are available publicly.

* feat: Proto Message Implementation

* feat: Adding support for enum

* feat: Code refactoring

Adding default implementation for newly added methods
ByteArray compatability changes for Proto Messages

* docs: Adding Java docs for all the newly added methods.

* test: Sample Proto & Generated classes for unit test

* feat: Adding bytes/proto & int64/enum compatability

Adding Additional check for ChecksumResultSet

* test: Adding unit tests

* test: Adding unit tests for ValueBinder.java

* feat: refactoring to add support for getValue & other minor changes

* feat: Minor refactoring

1. Adding docs and formatting the code.
2. Adding additional methods for enum and message which accepts descriptors.

* feat: Adding bytes/message & int64/enum compatability in Value

* refactor: Minor refactoring

* feat: Adding Proto Array Implementation

* test: Implementing unit tests for array of protos and enums

* refactor: adding clirr ignores

* feat: Adding support for enum as Primary Key

* feat: Code Review Changes, minor refactoring and adding docs

* feat: Addressing review comments

-Modified Docs/Comments
-Minor Refactoring

* refactor: Using Column instead of column to avoid test failures

* feat: Minor refactoring

-code review comments
-adding function docs

Author: gauravpurohit06

google-cloud-spanner/clirr-ignored-differences.xml

@@ -192,6 +192,46 @@
     <className>com/google/cloud/spanner/StructReader</className>
     <method>java.util.List getPgJsonbList(java.lang.String)</method>
   </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>com.google.protobuf.ProtocolMessageEnum getProtoEnum(int, java.util.function.Function)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>com.google.protobuf.ProtocolMessageEnum getProtoEnum(java.lang.String, java.util.function.Function)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>com.google.protobuf.AbstractMessage getProtoMessage(int, com.google.protobuf.AbstractMessage)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>com.google.protobuf.AbstractMessage getProtoMessage(java.lang.String, com.google.protobuf.AbstractMessage)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>java.util.List getProtoEnumList(int, java.util.function.Function)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>java.util.List getProtoEnumList(java.lang.String, java.util.function.Function)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>java.util.List getProtoMessageList(int, com.google.protobuf.AbstractMessage)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/StructReader</className>
+    <method>java.util.List getProtoMessageList(java.lang.String, com.google.protobuf.AbstractMessage)</method>
+  </difference>
   <difference>
     <differenceType>7012</differenceType>
     <className>com/google/cloud/spanner/BatchClient</className>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

@@ -32,12 +32,16 @@
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.Uninterruptibles;
+import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.ListValue;
+import com.google.protobuf.ProtocolMessageEnum;
 import com.google.protobuf.Value.KindCase;
 import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.ResultSetMetadata;
@@ -65,6 +69,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -391,6 +396,14 @@ private Object writeReplace() {
           case JSON:
             builder.set(fieldName).to(Value.json((String) value));
             break;
+          case PROTO:
+            builder
+                .set(fieldName)
+                .to(Value.protoMessage((ByteArray) value, fieldType.getProtoTypeFqn()));
+            break;
+          case ENUM:
+            builder.set(fieldName).to(Value.protoEnum((Long) value, fieldType.getProtoTypeFqn()));
+            break;
           case PG_JSONB:
             builder.set(fieldName).to(Value.pgJsonb((String) value));
             break;
@@ -410,6 +423,7 @@ private Object writeReplace() {
                 builder.set(fieldName).toBoolArray((Iterable<Boolean>) value);
                 break;
               case INT64:
+              case ENUM:
                 builder.set(fieldName).toInt64Array((Iterable<Long>) value);
                 break;
               case FLOAT64:
@@ -431,6 +445,7 @@ private Object writeReplace() {
                 builder.set(fieldName).toPgJsonbArray((Iterable<String>) value);
                 break;
               case BYTES:
+              case PROTO:
                 builder.set(fieldName).toBytesArray((Iterable<ByteArray>) value);
                 break;
               case TIMESTAMP:
@@ -496,6 +511,7 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
           checkType(fieldType, proto, KindCase.BOOL_VALUE);
           return proto.getBoolValue();
         case INT64:
+        case ENUM:
           checkType(fieldType, proto, KindCase.STRING_VALUE);
           return Long.parseLong(proto.getStringValue());
         case FLOAT64:
@@ -510,6 +526,7 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
           checkType(fieldType, proto, KindCase.STRING_VALUE);
           return proto.getStringValue();
         case BYTES:
+        case PROTO:
           checkType(fieldType, proto, KindCase.STRING_VALUE);
           return ByteArray.fromBase64(proto.getStringValue());
         case TIMESTAMP:
@@ -547,7 +564,8 @@ private static Struct decodeStructValue(Type structType, ListValue structValue)
     static Object decodeArrayValue(Type elementType, ListValue listValue) {
       switch (elementType.getCode()) {
         case INT64:
-          // For int64/float64 types, use custom containers.  These avoid wrapper object
+        case ENUM:
+          // For int64/float64/enum types, use custom containers.  These avoid wrapper object
           // creation for non-null arrays.
           return new Int64Array(listValue);
         case FLOAT64:
@@ -562,6 +580,7 @@ static Object decodeArrayValue(Type elementType, ListValue listValue) {
         case TIMESTAMP:
         case DATE:
         case STRUCT:
+        case PROTO:
           return Lists.transform(
               listValue.getValuesList(), input -> decodeValue(elementType, input));
         default:
@@ -597,6 +616,30 @@ public boolean isNull(int columnIndex) {
       return rowData.get(columnIndex) == null;
     }
 
+    @Override
+    protected <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T message) {
+      Preconditions.checkNotNull(
+          message,
+          "Proto message may not be null. Use MyProtoClass.getDefaultInstance() as a parameter value.");
+      try {
+        return (T)
+            message
+                .toBuilder()
+                .mergeFrom(((ByteArray) rowData.get(columnIndex)).toByteArray())
+                .build();
+      } catch (InvalidProtocolBufferException e) {
+        throw SpannerExceptionFactory.asSpannerException(e);
+      }
+    }
+
+    @Override
+    protected <T extends ProtocolMessageEnum> T getProtoEnumInternal(
+        int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+      Preconditions.checkNotNull(
+          method, "Method may not be null. Use 'MyProtoEnum::forNumber' as a parameter value.");
+      return (T) method.apply((int) getLongInternal(columnIndex));
+    }
+
     @Override
     protected boolean getBooleanInternal(int columnIndex) {
       return (Boolean) rowData.get(columnIndex);
@@ -658,6 +701,8 @@ protected Value getValueInternal(int columnIndex) {
           return Value.bool(isNull ? null : getBooleanInternal(columnIndex));
         case INT64:
           return Value.int64(isNull ? null : getLongInternal(columnIndex));
+        case ENUM:
+          return Value.protoEnum(getLongInternal(columnIndex), columnType.getProtoTypeFqn());
         case NUMERIC:
           return Value.numeric(isNull ? null : getBigDecimalInternal(columnIndex));
         case PG_NUMERIC:
@@ -672,6 +717,8 @@ protected Value getValueInternal(int columnIndex) {
           return Value.pgJsonb(isNull ? null : getPgJsonbInternal(columnIndex));
         case BYTES:
           return Value.bytes(isNull ? null : getBytesInternal(columnIndex));
+        case PROTO:
+          return Value.protoMessage(getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
         case TIMESTAMP:
           return Value.timestamp(isNull ? null : getTimestampInternal(columnIndex));
         case DATE:
@@ -699,6 +746,12 @@ protected Value getValueInternal(int columnIndex) {
               return Value.pgJsonbArray(isNull ? null : getPgJsonbListInternal(columnIndex));
             case BYTES:
               return Value.bytesArray(isNull ? null : getBytesListInternal(columnIndex));
+            case PROTO:
+              return Value.protoMessageArray(
+                  isNull ? null : getBytesListInternal(columnIndex), elementType.getProtoTypeFqn());
+            case ENUM:
+              return Value.protoEnumArray(
+                  isNull ? null : getLongListInternal(columnIndex), elementType.getProtoTypeFqn());
             case TIMESTAMP:
               return Value.timestampArray(isNull ? null : getTimestampListInternal(columnIndex));
             case DATE:
@@ -778,6 +831,52 @@ protected List<String> getJsonListInternal(int columnIndex) {
       return Collections.unmodifiableList((List<String>) rowData.get(columnIndex));
     }
 
+    @Override
+    @SuppressWarnings("unchecked") // We know ARRAY<PROTO> produces a List<ByteArray>.
+    protected <T extends AbstractMessage> List<T> getProtoMessageListInternal(
+        int columnIndex, T message) {
+      Preconditions.checkNotNull(
+          message,
+          "Proto message may not be null. Use MyProtoClass.getDefaultInstance() as a parameter value.");
+
+      List<ByteArray> bytesArray = (List<ByteArray>) rowData.get(columnIndex);
+
+      try {
+        List<T> protoMessagesList = new ArrayList<>(bytesArray.size());
+        for (ByteArray protoMessageBytes : bytesArray) {
+          if (protoMessageBytes == null) {
+            protoMessagesList.add(null);
+          } else {
+            protoMessagesList.add(
+                (T) message.toBuilder().mergeFrom(protoMessageBytes.toByteArray()).build());
+          }
+        }
+        return protoMessagesList;
+      } catch (InvalidProtocolBufferException e) {
+        throw SpannerExceptionFactory.asSpannerException(e);
+      }
+    }
+
+    @Override
+    @SuppressWarnings("unchecked") // We know ARRAY<ENUM> produces a List<Long>.
+    protected <T extends ProtocolMessageEnum> List<T> getProtoEnumListInternal(
+        int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+      Preconditions.checkNotNull(
+          method, "Method may not be null. Use 'MyProtoEnum::forNumber' as a parameter value.");
+
+      List<Long> enumIntArray = (List<Long>) rowData.get(columnIndex);
+      List<T> protoEnumList = new ArrayList<>(enumIntArray.size());
+      for (Long enumIntValue : enumIntArray) {
+        if (enumIntValue == null) {
+          protoEnumList.add(null);
+        } else {
+          protoEnumList.add((T) method.apply(enumIntValue.intValue()));
+        }
+      }
+
+      return protoEnumList;
+    }
+
     @Override
     @SuppressWarnings("unchecked") // We know ARRAY<JSONB> produces a List<String>.
     protected List<String> getPgJsonbListInternal(int columnIndex) {
@@ -1310,6 +1409,17 @@ protected String getStringInternal(int columnIndex) {
     return currRow().getStringInternal(columnIndex);
   }
 
+  @Override
+  protected <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T message) {
+    return currRow().getProtoMessageInternal(columnIndex, message);
+  }
+
+  @Override
+  protected <T extends ProtocolMessageEnum> T getProtoEnumInternal(
+      int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+    return currRow().getProtoEnumInternal(columnIndex, method);
+  }
+
   @Override
   protected String getJsonInternal(int columnIndex) {
     return currRow().getJsonInternal(columnIndex);
@@ -1395,6 +1505,18 @@ protected List<ByteArray> getBytesListInternal(int columnIndex) {
     return currRow().getBytesListInternal(columnIndex);
   }
 
+  @Override
+  protected <T extends AbstractMessage> List<T> getProtoMessageListInternal(
+      int columnIndex, T message) {
+    return currRow().getProtoMessageListInternal(columnIndex, message);
+  }
+
+  @Override
+  protected <T extends ProtocolMessageEnum> List<T> getProtoEnumListInternal(
+      int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+    return currRow().getProtoEnumListInternal(columnIndex, method);
+  }
+
   @Override
   protected List<Timestamp> getTimestampListInternal(int columnIndex) {
     return currRow().getTimestampListInternal(columnIndex);