feat: Add support for Proto and Enum types (#2662)

Change-Id: I27c4d06a3e29bb52b0e5391acba9730c05997164

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Rollback plan is reviewed and LGTMed
- [ ] All new data plane features have a completed end to end testing plan

Fixes #<issue_number_goes_here> ☕️

If you write sample code, please follow the [samples format](
https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).

Author: trollyxia

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

@@ -331,6 +331,36 @@
         <className>com/google/cloud/bigtable/data/v2/models/sql/Statement$Builder</className>
         <method>*</method>
     </difference>
+    <difference>
+        <!-- BetaApi was added -->
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/sql/StructReader</className>
+        <method>*getProtoMessage(*)</method>
+    </difference>
+    <difference>
+        <!-- BetaApi was added -->
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/sql/StructReader</className>
+        <method>*getProtoEnum(*)</method>
+    </difference>
+    <difference>
+        <!-- BetaApi was added -->
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/sql/SqlType</className>
+        <method>*protoOf(*)</method>
+    </difference>
+    <difference>
+        <!-- BetaApi was added -->
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/sql/SqlType</className>
+        <method>*enumOf(*)</method>
+    </difference>
+    <difference>
+        <!-- InternalApi was added -->
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/models/sql/SqlType</className>
+        <method>*getUnqualifiedName*</method>
+    </difference>
     <difference>
         <!-- InternalApi was updated -->
         <differenceType>7004</differenceType>

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/common/Type.java

@@ -21,10 +21,16 @@
 import com.google.cloud.bigtable.data.v2.internal.ColumnToIndexMapper;
 import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
 import com.google.common.base.Objects;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Parser;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.time.Instant;
 import java.util.List;
+import java.util.function.Function;
+import javax.annotation.Nonnull;
 
 /**
  * Shared type implementations. Right now this is only used by SqlType but this will become a shared
@@ -385,10 +391,188 @@ public java.lang.String toString() {
     }
   }
 
+  @AutoValue
+  abstract class Proto<T extends AbstractMessage> implements Type, SqlType.Proto<T> {
+
+    public static <T extends AbstractMessage> SqlType.Proto<T> create(T message) {
+      Preconditions.checkNotNull(
+          message,
+          "Proto message may not be null. Use 'MyProtoMessage::getDefaultInstance()' as a parameter value.");
+      return new AutoValue_Type_Proto<>(message);
+    }
+
+    @Nonnull
+    abstract T getMessage();
+
+    @Override
+    public Code getCode() {
+      return Code.PROTO;
+    }
+
+    @Nonnull
+    @Override
+    public Parser<T> getParserForType() {
+      return (Parser<T>) getMessage().getParserForType();
+    }
+
+    @Override
+    public java.lang.String getMessageName() {
+      return getMessage().getDescriptorForType().getFullName();
+    }
+
+    @Override
+    public java.lang.String toString() {
+      return getCode().name() + "{message=" + getMessageName() + "}";
+    }
+  }
+
+  @AutoValue
+  abstract class Enum<T extends ProtocolMessageEnum> implements Type, SqlType.Enum<T> {
+
+    public static <T extends ProtocolMessageEnum> SqlType.Enum<T> create(
+        Function<Integer, T> forNumber) {
+      Preconditions.checkNotNull(
+          forNumber, "Method may not be null. Use 'MyProtoEnum::forNumber' as a parameter value.");
+      return new AutoValue_Type_Enum<>(forNumber);
+    }
+
+    @Nonnull
+    @Override
+    public abstract Function<Integer, T> getForNumber();
+
+    @Override
+    public java.lang.String getEnumName() {
+      T thisEnum = getForNumber().apply(0);
+      if (thisEnum == null) {
+        return "";
+      }
+      return thisEnum.getDescriptorForType().getFullName();
+    }
+
+    @Override
+    public Code getCode() {
+      return Code.ENUM;
+    }
+
+    @Override
+    public java.lang.String toString() {
+      return getCode().name() + "{enum=" + getEnumName() + "}";
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof Type.Enum)) {
+        return false;
+      }
+      Type.Enum<?> that = (Type.Enum<?>) o;
+      // We don't want to compare functions directly, so try to get the enum descriptor and compare
+      // those.
+      T thisEnum = getForNumber().apply(0);
+      Object thatEnum = that.getForNumber().apply(0);
+
+      if (thisEnum == null || thatEnum == null) {
+        // Can't determine equality, fallback to object equality on the function.
+        return getForNumber().equals(that.getForNumber());
+      }
+      return thisEnum
+          .getDescriptorForType()
+          .getFullName()
+          .equals(((ProtocolMessageEnum) thatEnum).getDescriptorForType().getFullName());
+    }
+
+    @Override
+    public int hashCode() {
+      T thisEnum = getForNumber().apply(0);
+      if (thisEnum == null) {
+        return getForNumber().hashCode();
+      }
+      return java.util.Objects.hash(getCode(), thisEnum.getDescriptorForType().getFullName());
+    }
+  }
+
+  /**
+   * This is a special version of proto that is intended to only be used internally, to facilitate
+   * proto schema parsing, which does not have the full information required to parse the protobuf
+   * messages.
+   *
+   * <p>Any attempts to call getParserForType() will throw an exception.
+   */
+  @AutoValue
+  abstract class SchemalessProto implements SqlType.Proto {
+
+    public static SchemalessProto fromProto(com.google.bigtable.v2.Type.Proto proto) {
+      return create(proto.getMessageName());
+    }
+
+    public static SchemalessProto create(java.lang.String messageName) {
+      return new AutoValue_Type_SchemalessProto(messageName);
+    }
+
+    @Override
+    public abstract java.lang.String getMessageName();
+
+    @Override
+    public Parser<AbstractMessage> getParserForType() {
+      throw new UnsupportedOperationException(
+          "Cannot get parser for unresolved proto type. Please use the getProtoMessage overload that takes a message instance.");
+    }
+
+    @Override
+    public Code getCode() {
+      return Code.PROTO;
+    }
+
+    @Override
+    public java.lang.String toString() {
+      return getCode().name() + "{messageName=" + getMessageName() + "}";
+    }
+  }
+
+  /**
+   * This is a special version of enum that is intended to only be used internally, to facilitate
+   * enum schema parsing, which does not have the full information required to parse the protobuf
+   * enum messages.
+   *
+   * <p>Any attempts to call getForNumber() will throw an exception.
+   */
+  @AutoValue
+  abstract class SchemalessEnum implements SqlType.Enum {
+
+    public static SchemalessEnum fromProto(com.google.bigtable.v2.Type.Enum proto) {
+      return create(proto.getEnumName());
+    }
+
+    public static SchemalessEnum create(java.lang.String enumName) {
+      return new AutoValue_Type_SchemalessEnum(enumName);
+    }
+
+    public abstract java.lang.String getEnumName();
+
+    @Override
+    public Function<Integer, ProtocolMessageEnum> getForNumber() {
+      throw new UnsupportedOperationException(
+          "Cannot get forNumber for unresolved enum type. Please use the getProtoEnum overload that takes a forNumber function.");
+    }
+
+    @Override
+    public Code getCode() {
+      return Code.ENUM;
+    }
+
+    @Override
+    public java.lang.String toString() {
+      return getCode().name() + "{enumName=" + getEnumName() + "}";
+    }
+  }
+
   // Implementation detail to make singleton instances private without referencing the concrete
   // autovalue generated class from the abstract base classes.
   @InternalApi
   class DefaultInstances {
+
     private static final Bytes BYTES = new AutoValue_Type_Bytes();
     private static final String STRING = new AutoValue_Type_String();
     private static final Int64 INT64 = new AutoValue_Type_Int64();

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReader.java

@@ -23,13 +23,17 @@
 import com.google.cloud.bigtable.data.v2.models.sql.Struct;
 import com.google.cloud.bigtable.data.v2.models.sql.StructReader;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 @InternalApi
 public abstract class AbstractProtoStructReader implements StructReader {
@@ -220,6 +224,12 @@ public <ElemType> List<ElemType> getList(int columnIndex, SqlType.Array<ElemType
     SqlType<?> actualType = getColumnType(columnIndex);
     checkNonNullOfType(columnIndex, arrayType, actualType, columnIndex);
     Value value = values().get(columnIndex);
+    // If the element type is proto/enum, we should use the user passed type, which contains the
+    // schema. Otherwise, we should use the type from metadata.
+    SqlType<?> elementType = arrayType.getElementType();
+    if (elementType.getCode() == SqlType.Code.PROTO || elementType.getCode() == SqlType.Code.ENUM) {
+      return (List<ElemType>) decodeValue(value, arrayType);
+    }
     return (List<ElemType>) decodeValue(value, actualType);
   }
 
@@ -231,6 +241,12 @@ public <ElemType> List<ElemType> getList(String columnName, SqlType.Array<ElemTy
     SqlType<?> actualType = getColumnType(columnIndex);
     checkNonNullOfType(columnIndex, arrayType, actualType, columnName);
     Value value = values().get(columnIndex);
+    // If the element type is proto/enum, we should use the user passed type, which contains the
+    // schema. Otherwise, we should use the type from metadata.
+    SqlType<?> elementType = arrayType.getElementType();
+    if (elementType.getCode() == SqlType.Code.PROTO || elementType.getCode() == SqlType.Code.ENUM) {
+      return (List<ElemType>) decodeValue(value, arrayType);
+    }
     return (List<ElemType>) decodeValue(value, actualType);
   }
 
@@ -241,6 +257,12 @@ public <K, V> Map<K, V> getMap(int columnIndex, SqlType.Map<K, V> mapType) {
     SqlType<?> actualType = getColumnType(columnIndex);
     checkNonNullOfType(columnIndex, mapType, actualType, columnIndex);
     Value value = values().get(columnIndex);
+    // If the value type is proto/enum, we should use the user passed type, which contains the
+    // schema. Otherwise, we should use the type from metadata.
+    SqlType<?> valueType = mapType.getValueType();
+    if (valueType.getCode() == SqlType.Code.PROTO || valueType.getCode() == SqlType.Code.ENUM) {
+      return (Map<K, V>) decodeValue(value, mapType);
+    }
     return (Map<K, V>) decodeValue(value, actualType);
   }
 
@@ -252,9 +274,61 @@ public <K, V> Map<K, V> getMap(String columnName, SqlType.Map<K, V> mapType) {
     SqlType<?> actualType = getColumnType(columnIndex);
     checkNonNullOfType(columnIndex, mapType, actualType, columnName);
     Value value = values().get(columnIndex);
+    // If the value type is proto/enum, we should use the user passed type, which contains the
+    // schema. Otherwise, we should use the type from metadata.
+    SqlType<?> valueType = mapType.getValueType();
+    if (valueType.getCode() == SqlType.Code.PROTO || valueType.getCode() == SqlType.Code.ENUM) {
+      return (Map<K, V>) decodeValue(value, mapType);
+    }
     return (Map<K, V>) decodeValue(value, actualType);
   }
 
+  @Override
+  public <MsgType extends AbstractMessage> MsgType getProtoMessage(
+      int columnIndex, MsgType message) {
+    // Note it is import that we use the user passed message object to decode, because the type in
+    // the corresponding column metadata only have a message name and doesn't have schemas
+    SqlType.Proto<MsgType> actualType = SqlType.protoOf(message);
+    checkNonNullOfType(columnIndex, getColumnType(columnIndex), actualType, columnIndex);
+    Value value = values().get(columnIndex);
+    return (MsgType) decodeValue(value, actualType);
+  }
+
+  @Override
+  public <MsgType extends AbstractMessage> MsgType getProtoMessage(
+      String columnName, MsgType message) {
+    int columnIndex = getColumnIndex(columnName);
+    // Note it is import that we use the user passed message object to decode, because the type in
+    // the corresponding column metadata only have a message name and doesn't have schemas
+    SqlType.Proto<MsgType> actualType = SqlType.protoOf(message);
+    checkNonNullOfType(columnIndex, getColumnType(columnIndex), actualType, columnName);
+    Value value = values().get(columnIndex);
+    return (MsgType) decodeValue(value, actualType);
+  }
+
+  @Override
+  public <EnumType extends ProtocolMessageEnum> EnumType getProtoEnum(
+      int columnIndex, Function<Integer, EnumType> forNumber) {
+    // Note it is import that we use the user passed function to decode, because the type in
+    // the corresponding column metadata only have an enum message name and doesn't have schemas
+    SqlType.Enum<EnumType> actualType = SqlType.enumOf(forNumber);
+    checkNonNullOfType(columnIndex, getColumnType(columnIndex), actualType, columnIndex);
+    Value value = values().get(columnIndex);
+    return (EnumType) decodeValue(value, actualType);
+  }
+
+  @Override
+  public <EnumType extends ProtocolMessageEnum> EnumType getProtoEnum(
+      String columnName, Function<Integer, EnumType> forNumber) {
+    int columnIndex = getColumnIndex(columnName);
+    // Note it is import that we use the user passed function to decode, because the type in
+    // the corresponding column metadata only have an enum message name and doesn't have schemas
+    SqlType.Enum<EnumType> actualType = SqlType.enumOf(forNumber);
+    checkNonNullOfType(columnIndex, getColumnType(columnIndex), actualType, columnName);
+    Value value = values().get(columnIndex);
+    return (EnumType) decodeValue(value, actualType);
+  }
+
   Object decodeValue(Value value, SqlType<?> type) {
     if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) {
       return null;
@@ -281,6 +355,15 @@ Object decodeValue(Value value, SqlType<?> type) {
         SqlType.Struct schema = (SqlType.Struct) type;
         // A struct value is represented as an array
         return ProtoStruct.create(schema, value.getArrayValue());
+      case PROTO:
+        try {
+          SqlType.Proto protoType = (SqlType.Proto) type;
+          return protoType.getParserForType().parseFrom(value.getBytesValue());
+        } catch (InvalidProtocolBufferException e) {
+          throw new IllegalStateException("Unable to parse value to proto " + type, e);
+        }
+      case ENUM:
+        return ((SqlType.Enum<?>) type).getForNumber().apply((int) value.getIntValue());
       case ARRAY:
         ArrayList<Object> listBuilder = new ArrayList<>();
         SqlType.Array<?> arrayType = (SqlType.Array<?>) type;