Skip to content

Commit b1d44ce

Browse files
rajatbhattarajatbhatta
authored
Merge branch 'main' into lar
2 parents 90d455c + fc721c4 commit b1d44ce

File tree

13 files changed

+1338
-54
lines changed

13 files changed

+1338
-54
lines changed

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

Lines changed: 116 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,14 @@
3232
import com.google.cloud.spanner.spi.v1.SpannerRpc;
3333
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
3434
import com.google.common.annotations.VisibleForTesting;
35+
import com.google.common.base.Preconditions;
3536
import com.google.common.collect.AbstractIterator;
3637
import com.google.common.collect.ImmutableMap;
3738
import com.google.common.collect.Lists;
3839
import com.google.common.util.concurrent.Uninterruptibles;
3940
import com.google.protobuf.ByteString;
4041
import com.google.protobuf.ListValue;
42+
import com.google.protobuf.NullValue;
4143
import com.google.protobuf.Value.KindCase;
4244
import com.google.spanner.v1.PartialResultSet;
4345
import com.google.spanner.v1.ResultSetMetadata;
@@ -55,23 +57,29 @@
5557
import java.math.BigDecimal;
5658
import java.util.AbstractList;
5759
import java.util.ArrayList;
60+
import java.util.Base64;
5861
import java.util.BitSet;
5962
import java.util.Collections;
6063
import java.util.Iterator;
6164
import java.util.LinkedList;
6265
import java.util.List;
66+
import java.util.Objects;
6367
import java.util.concurrent.BlockingQueue;
6468
import java.util.concurrent.CountDownLatch;
6569
import java.util.concurrent.Executor;
6670
import java.util.concurrent.LinkedBlockingQueue;
6771
import java.util.concurrent.TimeUnit;
6872
import java.util.logging.Level;
6973
import java.util.logging.Logger;
74+
import java.util.stream.Collectors;
75+
import javax.annotation.Nonnull;
7076
import javax.annotation.Nullable;
7177

7278
/** Implementation of {@link ResultSet}. */
7379
abstract class AbstractResultSet<R> extends AbstractStructReader implements ResultSet {
7480
private static final Tracer tracer = Tracing.getTracer();
81+
private static final com.google.protobuf.Value NULL_VALUE =
82+
com.google.protobuf.Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build();
7583

7684
interface Listener {
7785
/**
@@ -353,6 +361,79 @@ private boolean isMergeable(KindCase kind) {
353361
}
354362
}
355363

364+
static final class LazyByteArray implements Serializable {
365+
private static final Base64.Encoder ENCODER = Base64.getEncoder();
366+
private static final Base64.Decoder DECODER = Base64.getDecoder();
367+
private final String base64String;
368+
private transient AbstractLazyInitializer<ByteArray> byteArray;
369+
370+
LazyByteArray(@Nonnull String base64String) {
371+
this.base64String = Preconditions.checkNotNull(base64String);
372+
this.byteArray = defaultInitializer();
373+
}
374+
375+
LazyByteArray(@Nonnull ByteArray byteArray) {
376+
this.base64String =
377+
ENCODER.encodeToString(Preconditions.checkNotNull(byteArray).toByteArray());
378+
this.byteArray =
379+
new AbstractLazyInitializer<ByteArray>() {
380+
@Override
381+
protected ByteArray initialize() {
382+
return byteArray;
383+
}
384+
};
385+
}
386+
387+
private AbstractLazyInitializer<ByteArray> defaultInitializer() {
388+
return new AbstractLazyInitializer<ByteArray>() {
389+
@Override
390+
protected ByteArray initialize() {
391+
return ByteArray.copyFrom(DECODER.decode(base64String));
392+
}
393+
};
394+
}
395+
396+
private void readObject(java.io.ObjectInputStream in)
397+
throws IOException, ClassNotFoundException {
398+
in.defaultReadObject();
399+
byteArray = defaultInitializer();
400+
}
401+
402+
ByteArray getByteArray() {
403+
try {
404+
return byteArray.get();
405+
} catch (Throwable t) {
406+
throw SpannerExceptionFactory.asSpannerException(t);
407+
}
408+
}
409+
410+
String getBase64String() {
411+
return base64String;
412+
}
413+
414+
@Override
415+
public String toString() {
416+
return getBase64String();
417+
}
418+
419+
@Override
420+
public int hashCode() {
421+
return base64String.hashCode();
422+
}
423+
424+
@Override
425+
public boolean equals(Object o) {
426+
if (o instanceof LazyByteArray) {
427+
return lazyByteArraysEqual((LazyByteArray) o);
428+
}
429+
return false;
430+
}
431+
432+
private boolean lazyByteArraysEqual(LazyByteArray other) {
433+
return Objects.equals(getBase64String(), other.getBase64String());
434+
}
435+
}
436+
356437
static class GrpcStruct extends Struct implements Serializable {
357438
private final Type type;
358439
private final List<Object> rowData;
@@ -395,7 +476,11 @@ private Object writeReplace() {
395476
builder.set(fieldName).to(Value.pgJsonb((String) value));
396477
break;
397478
case BYTES:
398-
builder.set(fieldName).to((ByteArray) value);
479+
builder
480+
.set(fieldName)
481+
.to(
482+
Value.bytesFromBase64(
483+
value == null ? null : ((LazyByteArray) value).getBase64String()));
399484
break;
400485
case TIMESTAMP:
401486
builder.set(fieldName).to((Timestamp) value);
@@ -431,7 +516,17 @@ private Object writeReplace() {
431516
builder.set(fieldName).toPgJsonbArray((Iterable<String>) value);
432517
break;
433518
case BYTES:
434-
builder.set(fieldName).toBytesArray((Iterable<ByteArray>) value);
519+
builder
520+
.set(fieldName)
521+
.toBytesArrayFromBase64(
522+
value == null
523+
? null
524+
: ((List<LazyByteArray>) value)
525+
.stream()
526+
.map(
527+
element ->
528+
element == null ? null : element.getBase64String())
529+
.collect(Collectors.toList()));
435530
break;
436531
case TIMESTAMP:
437532
builder.set(fieldName).toTimestampArray((Iterable<Timestamp>) value);
@@ -511,7 +606,7 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
511606
return proto.getStringValue();
512607
case BYTES:
513608
checkType(fieldType, proto, KindCase.STRING_VALUE);
514-
return ByteArray.fromBase64(proto.getStringValue());
609+
return new LazyByteArray(proto.getStringValue());
515610
case TIMESTAMP:
516611
checkType(fieldType, proto, KindCase.STRING_VALUE);
517612
return Timestamp.parseTimestamp(proto.getStringValue());
@@ -526,6 +621,8 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
526621
checkType(fieldType, proto, KindCase.LIST_VALUE);
527622
ListValue structValue = proto.getListValue();
528623
return decodeStructValue(fieldType, structValue);
624+
case UNRECOGNIZED:
625+
return proto;
529626
default:
530627
throw new AssertionError("Unhandled type code: " + fieldType.getCode());
531628
}
@@ -634,7 +731,11 @@ protected String getPgJsonbInternal(int columnIndex) {
634731

635732
@Override
636733
protected ByteArray getBytesInternal(int columnIndex) {
637-
return (ByteArray) rowData.get(columnIndex);
734+
return getLazyBytesInternal(columnIndex).getByteArray();
735+
}
736+
737+
LazyByteArray getLazyBytesInternal(int columnIndex) {
738+
return (LazyByteArray) rowData.get(columnIndex);
638739
}
639740

640741
@Override
@@ -647,6 +748,10 @@ protected Date getDateInternal(int columnIndex) {
647748
return (Date) rowData.get(columnIndex);
648749
}
649750

751+
protected com.google.protobuf.Value getProtoValueInternal(int columnIndex) {
752+
return (com.google.protobuf.Value) rowData.get(columnIndex);
753+
}
754+
650755
@Override
651756
protected Value getValueInternal(int columnIndex) {
652757
final List<Type.StructField> structFields = getType().getStructFields();
@@ -671,13 +776,16 @@ protected Value getValueInternal(int columnIndex) {
671776
case PG_JSONB:
672777
return Value.pgJsonb(isNull ? null : getPgJsonbInternal(columnIndex));
673778
case BYTES:
674-
return Value.bytes(isNull ? null : getBytesInternal(columnIndex));
779+
return Value.internalBytes(isNull ? null : getLazyBytesInternal(columnIndex));
675780
case TIMESTAMP:
676781
return Value.timestamp(isNull ? null : getTimestampInternal(columnIndex));
677782
case DATE:
678783
return Value.date(isNull ? null : getDateInternal(columnIndex));
679784
case STRUCT:
680785
return Value.struct(isNull ? null : getStructInternal(columnIndex));
786+
case UNRECOGNIZED:
787+
return Value.unrecognized(
788+
isNull ? NULL_VALUE : getProtoValueInternal(columnIndex), columnType);
681789
case ARRAY:
682790
final Type elementType = columnType.getArrayElementType();
683791
switch (elementType.getCode()) {
@@ -785,9 +893,10 @@ protected List<String> getPgJsonbListInternal(int columnIndex) {
785893
}
786894

787895
@Override
788-
@SuppressWarnings("unchecked") // We know ARRAY<BYTES> produces a List<ByteArray>.
896+
@SuppressWarnings("unchecked") // We know ARRAY<BYTES> produces a List<LazyByteArray>.
789897
protected List<ByteArray> getBytesListInternal(int columnIndex) {
790-
return Collections.unmodifiableList((List<ByteArray>) rowData.get(columnIndex));
898+
return Lists.transform(
899+
(List<LazyByteArray>) rowData.get(columnIndex), l -> l == null ? null : l.getByteArray());
791900
}
792901

793902
@Override

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,6 @@ public Date getDate(String columnName) {
251251

252252
@Override
253253
public Value getValue(int columnIndex) {
254-
checkNonNull(columnIndex, columnIndex);
255254
return getValueInternal(columnIndex);
256255
}
257256

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/Type.java

Lines changed: 51 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.Map.Entry;
3434
import java.util.Objects;
3535
import java.util.TreeMap;
36+
import javax.annotation.Nonnull;
3637
import javax.annotation.Nullable;
3738
import javax.annotation.concurrent.Immutable;
3839

@@ -71,6 +72,10 @@ public final class Type implements Serializable {
7172
private static final int AMBIGUOUS_FIELD = -1;
7273
private static final long serialVersionUID = -3076152125004114582L;
7374

75+
static Type unrecognized(com.google.spanner.v1.Type proto) {
76+
return new Type(proto);
77+
}
78+
7479
/** Returns the descriptor for the {@code BOOL type}. */
7580
public static Type bool() {
7681
return TYPE_BOOL;
@@ -190,6 +195,7 @@ public static Type struct(StructField... fields) {
190195
return new Type(Code.STRUCT, null, ImmutableList.copyOf(fields));
191196
}
192197

198+
private final com.google.spanner.v1.Type proto;
193199
private final Code code;
194200
private final Type arrayElementType;
195201
private final ImmutableList<StructField> structFields;
@@ -201,16 +207,34 @@ public static Type struct(StructField... fields) {
201207
private Map<String, Integer> fieldsByName;
202208

203209
private Type(
204-
Code code,
210+
@Nonnull Code code,
205211
@Nullable Type arrayElementType,
206212
@Nullable ImmutableList<StructField> structFields) {
213+
this(null, Preconditions.checkNotNull(code), arrayElementType, structFields);
214+
}
215+
216+
private Type(@Nonnull com.google.spanner.v1.Type proto) {
217+
this(
218+
Preconditions.checkNotNull(proto),
219+
Code.UNRECOGNIZED,
220+
proto.hasArrayElementType() ? new Type(proto.getArrayElementType()) : null,
221+
null);
222+
}
223+
224+
private Type(
225+
com.google.spanner.v1.Type proto,
226+
Code code,
227+
Type arrayElementType,
228+
ImmutableList<StructField> structFields) {
229+
this.proto = proto;
207230
this.code = code;
208231
this.arrayElementType = arrayElementType;
209232
this.structFields = structFields;
210233
}
211234

212235
/** Enumerates the categories of types. */
213236
public enum Code {
237+
UNRECOGNIZED(TypeCode.UNRECOGNIZED),
214238
BOOL(TypeCode.BOOL),
215239
INT64(TypeCode.INT64),
216240
NUMERIC(TypeCode.NUMERIC),
@@ -258,8 +282,7 @@ TypeAnnotationCode getTypeAnnotationCode() {
258282

259283
static Code fromProto(TypeCode typeCode, TypeAnnotationCode typeAnnotationCode) {
260284
Code code = protoToCode.get(new SimpleEntry<>(typeCode, typeAnnotationCode));
261-
checkArgument(code != null, "Invalid code: %s<%s>", typeCode, typeAnnotationCode);
262-
return code;
285+
return code == null ? Code.UNRECOGNIZED : code;
263286
}
264287

265288
@Override
@@ -325,7 +348,7 @@ public Code getCode() {
325348
* @throws IllegalStateException if {@code code() != Code.ARRAY}
326349
*/
327350
public Type getArrayElementType() {
328-
Preconditions.checkState(code == Code.ARRAY, "Illegal call for non-ARRAY type");
351+
Preconditions.checkState(arrayElementType != null, "Illegal call for non-ARRAY type");
329352
return arrayElementType;
330353
}
331354

@@ -378,8 +401,14 @@ public int getFieldIndex(String fieldName) {
378401
}
379402

380403
void toString(StringBuilder b) {
381-
if (code == Code.ARRAY) {
382-
b.append("ARRAY<");
404+
if (code == Code.ARRAY || (proto != null && proto.hasArrayElementType())) {
405+
if (code == Code.ARRAY) {
406+
b.append("ARRAY<");
407+
} else {
408+
// This is very unlikely to happen. It would mean that we have introduced a type that
409+
// is not an ARRAY, but does have an array element type.
410+
b.append("UNRECOGNIZED<");
411+
}
383412
arrayElementType.toString(b);
384413
b.append('>');
385414
} else if (code == Code.STRUCT) {
@@ -393,6 +422,11 @@ void toString(StringBuilder b) {
393422
f.getType().toString(b);
394423
}
395424
b.append('>');
425+
} else if (proto != null) {
426+
b.append(proto.getCode().name());
427+
if (proto.getTypeAnnotation() != TYPE_ANNOTATION_CODE_UNSPECIFIED) {
428+
b.append("<").append(proto.getTypeAnnotation().name()).append(">");
429+
}
396430
} else {
397431
b.append(code.toString());
398432
}
@@ -414,17 +448,26 @@ public boolean equals(Object o) {
414448
return false;
415449
}
416450
Type that = (Type) o;
451+
if (proto != null) {
452+
return Objects.equals(proto, that.proto);
453+
}
417454
return code == that.code
418455
&& Objects.equals(arrayElementType, that.arrayElementType)
419456
&& Objects.equals(structFields, that.structFields);
420457
}
421458

422459
@Override
423460
public int hashCode() {
461+
if (proto != null) {
462+
return proto.hashCode();
463+
}
424464
return Objects.hash(code, arrayElementType, structFields);
425465
}
426466

427467
com.google.spanner.v1.Type toProto() {
468+
if (proto != null) {
469+
return proto;
470+
}
428471
com.google.spanner.v1.Type.Builder proto = com.google.spanner.v1.Type.newBuilder();
429472
proto.setCode(code.getTypeCode());
430473
proto.setTypeAnnotation(code.getTypeAnnotationCode());
@@ -490,8 +533,9 @@ static Type fromProto(com.google.spanner.v1.Type proto) {
490533
fields.add(StructField.of(name, fromProto(field.getType())));
491534
}
492535
return struct(fields);
536+
case UNRECOGNIZED:
493537
default:
494-
throw new AssertionError("Unimplemented case: " + type);
538+
return unrecognized(proto);
495539
}
496540
}
497541
}

0 commit comments

Comments
 (0)