feat: Add support for Proto and Enum types

Change-Id: I27c4d06a3e29bb52b0e5391acba9730c05997164

Author: trollyxia

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,6 +391,184 @@ 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 message();
+
+    @Override
+    public Code getCode() {
+      return Code.PROTO;
+    }
+
+    @Nonnull
+    @Override
+    public Parser<T> getParserForType() {
+      return (Parser<T>) message().getParserForType();
+    }
+
+    @Override
+    public java.lang.String getMessageName() {
+      return message().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 Code getCode() {
+      return Code.ENUM;
+    }
+
+    @Override
+    public java.lang.String toString() {
+      T thisEnum = getForNumber().apply(0);
+      if (thisEnum == null) {
+        return getCode().name() + "{function=" + getForNumber() + "}";
+      }
+      return getCode().name() + "{enum=" + thisEnum.getDescriptorForType().getFullName() + "}";
+    }
+
+    @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 in the {@link
+   * com.google.cloud.bigtable.data.v2.models.sql.StructReader} getters that require types. We don't
+   * want users to need to specify the proto schema when the schema will be validated on calls to
+   * {@link com.google.cloud.bigtable.data.v2.models.sql.StructReader} methods on the struct.
+   *
+   * <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);
+    }
+
+    abstract java.lang.String messageName();
+
+    @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 java.lang.String getMessageName() {
+      return messageName();
+    }
+
+    @Override
+    public Code getCode() {
+      return Code.PROTO;
+    }
+
+    @Override
+    public java.lang.String toString() {
+      return getCode().name() + "{messageName=" + messageName() + "}";
+    }
+  }
+
+  /**
+   * This is a special version of enum that is intended to only be used in the {@link
+   * com.google.cloud.bigtable.data.v2.models.sql.StructReader} getters that require types. We don't
+   * want users to need to specify the enum schema when the schema will be validated on calls to
+   * {@link com.google.cloud.bigtable.data.v2.models.sql.StructReader} methods on the struct.
+   *
+   * <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);
+    }
+
+    abstract java.lang.String enumName();
+
+    @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=" + enumName() + "}";
+    }
+  }
+
   // Implementation detail to make singleton instances private without referencing the concrete
   // autovalue generated class from the abstract base classes.
   @InternalApi

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,14 @@ 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 {
+          return ((SqlType.Proto<?>) type).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;

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

@@ -27,11 +27,14 @@
 import com.google.cloud.bigtable.data.v2.models.sql.StructReader;
 import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.ProtocolMessageEnum;
 import java.time.Instant;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * The primary implementation of a ResultSet.
@@ -215,4 +218,28 @@ public <K, V> Map<K, V> getMap(int columnIndex, SqlType.Map<K, V> mapType) {
   public <K, V> Map<K, V> getMap(String columnName, SqlType.Map<K, V> mapType) {
     return getCurrentRow().getMap(columnName, mapType);
   }
+
+  @Override
+  public <EnumType extends ProtocolMessageEnum> EnumType getProtoEnum(
+      int columnIndex, Function<Integer, EnumType> forNumber) {
+    return getCurrentRow().getProtoEnum(columnIndex, forNumber);
+  }
+
+  @Override
+  public <EnumType extends ProtocolMessageEnum> EnumType getProtoEnum(
+      String columnName, Function<Integer, EnumType> forNumber) {
+    return getCurrentRow().getProtoEnum(columnName, forNumber);
+  }
+
+  @Override
+  public <MsgType extends AbstractMessage> MsgType getProtoMessage(
+      int columnIndex, MsgType message) {
+    return getCurrentRow().getProtoMessage(columnIndex, message);
+  }
+
+  @Override
+  public <MsgType extends AbstractMessage> MsgType getProtoMessage(
+      String columnName, MsgType message) {
+    return getCurrentRow().getProtoMessage(columnName, message);
+  }
 }