feat: Support for Proto Messages & Enums #2155
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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: | ||
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| builder | ||
| .set(fieldName) | ||
| .to(Value.protoMessage((ByteArray) value, fieldType.getProtoTypeFqn())); | ||
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| break; | ||
| case ENUM: | ||
| builder.set(fieldName).to(Value.protoEnum((Long) value, fieldType.getProtoTypeFqn())); | ||
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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: | ||
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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: | ||
| 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."); | ||
gauravpurohit06 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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."); | ||
|
There was a problem hiding this comment. Is it possible to call it for the user in Java? There was a problem hiding this comment. @Shua1: I prefer getProtoEnumListInternal(idx, my_method) here. Adding With getProtoEnumListInternal(idx), the Java client will not have any way to convert |
||
|
|
||
| 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); | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think these clirr changes are no longer needed, since they were made non-mandatory a few months back. Any particular reason these are failing?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As far as I can tell, the check is no longer required, but it is still active. Which means that if you don't add this, the build of the PR will be marked red (and some of us have brains that are not able to handle that, and need a green checkmark to be able to sleep at night ;-) )