Proto col support [DRAFT] [DO NOT MERGE]

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

@@ -36,8 +36,11 @@
 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 +68,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 java.util.stream.Collectors;
@@ -385,6 +389,10 @@ private Object writeReplace() {
           case JSON:
             builder.set(fieldName).to(Value.json((String) value));
             break;
+          case PROTO:
+            builder.set(fieldName).to(Value.protoMessage((byte[]) value));
+          case PROTO_ENUM:
+            builder.set(fieldName).to(Value.protoEnum((long) value));
           case PG_JSONB:
             builder.set(fieldName).to(Value.pgJsonb((String) value));
             break;
@@ -490,6 +498,7 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
           checkType(fieldType, proto, KindCase.BOOL_VALUE);
           return proto.getBoolValue();
         case INT64:
+        case PROTO_ENUM:
           checkType(fieldType, proto, KindCase.STRING_VALUE);
           return Long.parseLong(proto.getStringValue());
         case FLOAT64:
@@ -504,6 +513,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:
@@ -655,6 +665,19 @@ public boolean isNull(int columnIndex) {
       return rowData.get(columnIndex) == null;
     }
 
+    @Override
+    public <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T m)
+        throws InvalidProtocolBufferException {
+      return (T) m.toBuilder().mergeFrom(((ByteArray) rowData.get(columnIndex)).toByteArray())
+          .build();
+    }
+
+    @Override
+    protected <T extends ProtocolMessageEnum> T getProtoEnumInternal(
+        int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+      return (T) method.apply((int) getLongInternal(columnIndex));
+    }
+
     @Override
     protected boolean getBooleanInternal(int columnIndex) {
       return (Boolean) rowData.get(columnIndex);
@@ -1368,6 +1391,18 @@ protected String getStringInternal(int columnIndex) {
     return currRow().getStringInternal(columnIndex);
   }
 
+  @Override
+  protected <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T m)
+      throws InvalidProtocolBufferException {
+    return currRow().getProtoMessageInternal(columnIndex, m);
+  }
+
+  @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);

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

@@ -21,9 +21,13 @@
 import com.google.cloud.ByteArray;
 import com.google.cloud.Date;
 import com.google.cloud.Timestamp;
+import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.math.BigDecimal;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Function;
 
 /**
  * Base class for assisting {@link StructReader} implementations.
@@ -58,6 +62,12 @@ protected String getPgJsonbInternal(int columnIndex) {
 
   protected abstract Date getDateInternal(int columnIndex);
 
+  protected abstract <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T m)
+      throws InvalidProtocolBufferException;
+
+  protected abstract <T extends ProtocolMessageEnum> T getProtoEnumInternal(
+      int columnIndex, Function<Integer, ProtocolMessageEnum> method);
+
   protected Value getValueInternal(int columnIndex) {
     throw new UnsupportedOperationException("method should be overwritten");
   }
@@ -212,14 +222,16 @@ public String getPgJsonb(String columnName) {
 
   @Override
   public ByteArray getBytes(int columnIndex) {
-    checkNonNullOfType(columnIndex, Type.bytes(), columnIndex);
+    checkNonNullOfTypes(columnIndex, Arrays.asList(Type.bytes(), Type.proto()), columnIndex,
+        "BYTES, PROTO");
     return getBytesInternal(columnIndex);
   }
 
   @Override
   public ByteArray getBytes(String columnName) {
     int columnIndex = getColumnIndex(columnName);
-    checkNonNullOfType(columnIndex, Type.bytes(), columnName);
+    checkNonNullOfTypes(columnIndex, Arrays.asList(Type.bytes(), Type.proto()), columnName,
+        "BYTES, PROTO");
     return getBytesInternal(columnIndex);
   }
 
@@ -249,6 +261,36 @@ public Date getDate(String columnName) {
     return getDateInternal(columnIndex);
   }
 
+  @Override
+  public <T extends ProtocolMessageEnum> T getProtoEnum(
+      int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+    checkNonNullOfType(columnIndex, Type.protoEnum(), columnIndex);
+    return getProtoEnumInternal(columnIndex, method);
+  }
+
+  @Override
+  public <T extends ProtocolMessageEnum> T getProtoEnum(
+      String columnName, Function<Integer, ProtocolMessageEnum> method) {
+    int columnIndex = getColumnIndex(columnName);
+    checkNonNullOfType(columnIndex, Type.proto(), columnName);
+    return getProtoEnumInternal(columnIndex, method);
+  }
+
+  @Override
+  public <T extends AbstractMessage> T getProtoMessage(int columnIndex, T m)
+      throws InvalidProtocolBufferException {
+    checkNonNullOfType(columnIndex, Type.proto(), columnIndex);
+    return getProtoMessageInternal(columnIndex, m);
+  }
+
+  @Override
+  public <T extends AbstractMessage> T getProtoMessage(String columnName, T m)
+      throws InvalidProtocolBufferException {
+    int columnIndex = getColumnIndex(columnName);
+    checkNonNullOfType(columnIndex, Type.proto(), columnName);
+    return getProtoMessageInternal(columnIndex, m);
+  }
+
   @Override
   public Value getValue(int columnIndex) {
     checkNonNull(columnIndex, columnIndex);

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

@@ -22,8 +22,12 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
+import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.math.BigDecimal;
 import java.util.List;
+import java.util.function.Function;
 
 /** Forwarding implements of StructReader */
 public class ForwardingStructReader implements StructReader {
@@ -382,6 +386,34 @@ public List<Struct> getStructList(String columnName) {
     return delegate.get().getStructList(columnName);
   }
 
+  @Override
+  public <T extends AbstractMessage> T getProtoMessage(int columnIndex, T m)
+      throws InvalidProtocolBufferException {
+    checkValidState();
+    return delegate.get().getProtoMessage(columnIndex, m);
+  }
+
+  @Override
+  public <T extends AbstractMessage> T getProtoMessage(String columnName, T m)
+      throws InvalidProtocolBufferException {
+    checkValidState();
+    return delegate.get().getProtoMessage(columnName, m);
+  }
+
+  @Override
+  public <T extends ProtocolMessageEnum> T getProtoEnum(
+      int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+    checkValidState();
+    return delegate.get().getProtoEnum(columnIndex, method);
+  }
+
+  @Override
+  public <T extends ProtocolMessageEnum> T getProtoEnum(
+      String columnName, Function<Integer, ProtocolMessageEnum> method) {
+    checkValidState();
+    return delegate.get().getProtoEnum(columnName, method);
+  }
+
   @Override
   public Value getValue(int columnIndex) {
     checkValidState();

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2022 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
+ *
+ *       http://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.
+ */
+
+package com.google.cloud.spanner;
+
+import com.google.protobuf.ProtocolMessageEnum;
+
+public class ProtoEnumWrapper {
+  ProtocolMessageEnum protocolMessageEnum;
+  long enumValue;
+
+  ProtoEnumWrapper(ProtocolMessageEnum protoEnum) {
+    this.protocolMessageEnum = protoEnum;
+    this.enumValue = protoEnum.getNumber();
+  }
+
+  ProtoEnumWrapper(long val) {
+    this.enumValue = val;
+  }
+}

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2022 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
+ *
+ *       http://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.
+ */
+
+package com.google.cloud.spanner;
+
+import com.google.protobuf.AbstractMessage;
+import java.util.Arrays;
+
+public class ProtoMessageWrapper {
+  AbstractMessage message;
+  byte[] serializedMessage;
+
+  ProtoMessageWrapper(AbstractMessage m) {
+    this.message = m;
+    this.serializedMessage = m.toByteArray();
+  }
+
+  ProtoMessageWrapper(byte[] serializedMessage) {
+    this.serializedMessage = serializedMessage;
+  }
+
+  @Override
+  public String toString() {
+    if (message != null) {
+      return message.toString();
+    }
+    return Arrays.toString(serializedMessage);
+  }
+}

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

@@ -29,9 +29,13 @@
 import com.google.common.base.Supplier;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolMessageEnum;
 import com.google.spanner.v1.ResultSetStats;
 import java.math.BigDecimal;
 import java.util.List;
+import java.util.function.Function;
 
 /** Utility methods for working with {@link com.google.cloud.spanner.ResultSet}. */
 public final class ResultSets {
@@ -293,6 +297,31 @@ public Date getDate(String columnName) {
       return getCurrentRowAsStruct().getDate(columnName);
     }
 
+
+    @Override
+    public <T extends AbstractMessage> T getProtoMessage(int columnIndex, T m)
+        throws InvalidProtocolBufferException {
+      return getCurrentRowAsStruct().getProtoMessage(columnIndex, m);
+    }
+
+    @Override
+    public <T extends AbstractMessage> T getProtoMessage(String columnName, T m)
+        throws InvalidProtocolBufferException {
+      return getCurrentRowAsStruct().getProtoMessage(columnName, m);
+    }
+
+    @Override
+    public <T extends ProtocolMessageEnum> T getProtoEnum(
+        int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+      return getCurrentRowAsStruct().getProtoEnum(columnIndex, method);
+    }
+
+    @Override
+    public <T extends ProtocolMessageEnum> T getProtoEnum(
+        String columnName, Function<Integer, ProtocolMessageEnum> method) {
+      return getCurrentRowAsStruct().getProtoEnum(columnName, method);
+    }
+
     @Override
     public Value getValue(int columnIndex) {
       return getCurrentRowAsStruct().getValue(columnIndex);

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

@@ -28,11 +28,15 @@
 import com.google.common.primitives.Booleans;
 import com.google.common.primitives.Doubles;
 import com.google.common.primitives.Longs;
+import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.io.Serializable;
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Function;
 import javax.annotation.concurrent.Immutable;
 
 /**
@@ -217,6 +221,18 @@ protected Date getDateInternal(int columnIndex) {
       return values.get(columnIndex).getDate();
     }
 
+    @Override
+    protected <T extends AbstractMessage> T getProtoMessageInternal(int columnIndex, T m)
+        throws InvalidProtocolBufferException {
+      return values.get(columnIndex).getProtoMessage(m);
+    }
+
+    @Override
+    protected <T extends ProtocolMessageEnum> T getProtoEnumInternal(
+        int columnIndex, Function<Integer, ProtocolMessageEnum> method) {
+      return values.get(columnIndex).getProtoEnum(method);
+    }
+
     @Override
     protected Value getValueInternal(int columnIndex) {
       return values.get(columnIndex);
@@ -354,6 +370,7 @@ private Object getAsObject(int columnIndex) {
       case BOOL:
         return getBooleanInternal(columnIndex);
       case INT64:
+      case PROTO_ENUM:
         return getLongInternal(columnIndex);
       case FLOAT64:
         return getDoubleInternal(columnIndex);
@@ -368,6 +385,7 @@ private Object getAsObject(int columnIndex) {
       case PG_JSONB:
         return getPgJsonbInternal(columnIndex);
       case BYTES:
+      case PROTO:
         return getBytesInternal(columnIndex);
       case TIMESTAMP:
         return getTimestampInternal(columnIndex);