googleapis · olavloite · Feb 9, 2024 · Jan 27, 2024 · Jan 27, 2024 · Jan 28, 2024
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java
@@ -23,7 +23,7 @@
 import com.google.spanner.v1.ResultSetStats;
 
 /** Forwarding implementation of ResultSet that forwards all calls to a delegate. */
-public class ForwardingResultSet extends ForwardingStructReader implements ResultSet {
+public class ForwardingResultSet extends ForwardingStructReader implements ProtobufResultSet {
 
   private Supplier<ResultSet> delegate;
 
@@ -55,6 +55,22 @@ public boolean next() throws SpannerException {
     return delegate.get().next();
   }
 
+  @Override
+  public boolean canGetProtobufValue(int columnIndex) {
+    ResultSet resultSetDelegate = delegate.get();
+    return (resultSetDelegate instanceof ProtobufResultSet)
+        && ((ProtobufResultSet) resultSetDelegate).canGetProtobufValue(columnIndex);
+  }
+
+  @Override
+  public com.google.protobuf.Value getProtobufValue(int columnIndex) {
+    ResultSet resultSetDelegate = delegate.get();
+    Preconditions.checkState(
+        resultSetDelegate instanceof ProtobufResultSet,
+        "The result set does not support protobuf values");
+    return ((ProtobufResultSet) resultSetDelegate).getProtobufValue(columnIndex);
+  }
+
   @Override
   public Struct getCurrentRowAsStruct() {
     return delegate.get().getCurrentRowAsStruct();

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java
@@ -20,6 +20,7 @@
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.Value;
 import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
@@ -28,7 +29,7 @@
 import javax.annotation.Nullable;
 
 @VisibleForTesting
-class GrpcResultSet extends AbstractResultSet<List<Object>> {
+class GrpcResultSet extends AbstractResultSet<List<Object>> implements ProtobufResultSet {
   private final GrpcValueIterator iterator;
   private final Listener listener;
   private final DecodeMode decodeMode;
@@ -49,6 +50,18 @@ class GrpcResultSet extends AbstractResultSet<List<Object>> {
     this.decodeMode = decodeMode;
   }
 
+  @Override
+  public boolean canGetProtobufValue(int columnIndex) {
+    return !closed && currRow != null && currRow.canGetProtoValue(columnIndex);
+  }
+
+  @Override
+  public Value getProtobufValue(int columnIndex) {
+    checkState(!closed, "ResultSet is closed");
+    checkState(currRow != null, "next() call required");
+    return currRow.getProtoValueInternal(columnIndex);
+  }
+
   @Override
   protected GrpcStruct currRow() {
     checkState(!closed, "ResultSet is closed");

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStruct.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStruct.java
@@ -27,6 +27,7 @@
 import com.google.cloud.spanner.AbstractResultSet.Float64Array;
 import com.google.cloud.spanner.AbstractResultSet.Int64Array;
 import com.google.cloud.spanner.AbstractResultSet.LazyByteArray;
+import com.google.cloud.spanner.Type.Code;
 import com.google.cloud.spanner.Type.StructField;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
@@ -459,10 +460,37 @@ protected Date getDateInternal(int columnIndex) {
     return (Date) rowData.get(columnIndex);
   }
 
+  private boolean isUnrecognizedType(int columnIndex) {
+    return type.getStructFields().get(columnIndex).getType().getCode() == Code.UNRECOGNIZED;
+  }
+
+  boolean canGetProtoValue(int columnIndex) {
+    return isUnrecognizedType(columnIndex)
+        || (decodeMode == DecodeMode.LAZY_PER_ROW && !rowDecoded)
+        || (decodeMode == DecodeMode.LAZY_PER_COL && !colDecoded.get(columnIndex));
+  }
+
   protected com.google.protobuf.Value getProtoValueInternal(int columnIndex) {
+    checkProtoValueSupported(columnIndex);
     return (com.google.protobuf.Value) rowData.get(columnIndex);
   }
 
+  private void checkProtoValueSupported(int columnIndex) {
+    // Unrecognized types are returned as protobuf values.
+    if (isUnrecognizedType(columnIndex)) {
+      return;
+    }
+    Preconditions.checkState(
+        decodeMode != DecodeMode.DIRECT,
+        "Getting proto value is not supported when DecodeMode#DIRECT is used.");
+    Preconditions.checkState(
+        !(decodeMode == DecodeMode.LAZY_PER_ROW && rowDecoded),
+        "Getting proto value after the row has been decoded is not supported.");
+    Preconditions.checkState(
+        !(decodeMode == DecodeMode.LAZY_PER_COL && colDecoded.get(columnIndex)),
+        "Getting proto value after the column has been decoded is not supported.");
+  }
+
   private void ensureDecoded(int columnIndex) {
     if (decodeMode == DecodeMode.LAZY_PER_ROW && !rowDecoded) {
       for (int i = 0; i < rowData.size(); i++) {
@@ -496,7 +524,8 @@ protected Value getValueInternal(int columnIndex) {
       case INT64:
         return Value.int64(isNull ? null : getLongInternal(columnIndex));
       case ENUM:
-        return Value.protoEnum(getLongInternal(columnIndex), columnType.getProtoTypeFqn());
+        return Value.protoEnum(
+            isNull ? null : getLongInternal(columnIndex), columnType.getProtoTypeFqn());
       case NUMERIC:
         return Value.numeric(isNull ? null : getBigDecimalInternal(columnIndex));
       case PG_NUMERIC:
@@ -512,7 +541,8 @@ protected Value getValueInternal(int columnIndex) {
       case BYTES:
         return Value.internalBytes(isNull ? null : getLazyBytesInternal(columnIndex));
       case PROTO:
-        return Value.protoMessage(getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
+        return Value.protoMessage(
+            isNull ? null : getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
       case TIMESTAMP:
         return Value.timestamp(isNull ? null : getTimestampInternal(columnIndex));
       case DATE:

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java
@@ -383,7 +383,6 @@ static final class DecodeOption extends InternalOption implements ReadAndQueryOp
 
     DecodeOption(DecodeMode decodeMode) {
       this.decodeMode = Preconditions.checkNotNull(decodeMode, "DecodeMode cannot be null");
-      ;
     }
 
     @Override

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ProtobufResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ProtobufResultSet.java
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2024 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.api.core.InternalApi;
+import com.google.protobuf.Value;
+
+/** Interface for {@link ResultSet}s that can return a protobuf value. */
+@InternalApi
+public interface ProtobufResultSet extends ResultSet {
+
+  /** Returns true if the protobuf value for the given column is still available. */
+  boolean canGetProtobufValue(int columnIndex);
+
+  /**
+   * Returns the column value as a protobuf value.
+   *
+   * <p>This is an internal method not intended for external usage.
+   *
+   * <p>This method may only be called before the column value has been decoded to a plain Java
+   * object. This means that the {@link DecodeMode} that is used for the {@link ResultSet} must be
+   * one of {@link DecodeMode#LAZY_PER_ROW} and {@link DecodeMode#LAZY_PER_COL}, and that the
+   * corresponding {@link ResultSet#getValue(int)}, {@link ResultSet#getBoolean(int)}, ... method
+   * may not yet have been called for the column.
+   */
+  @InternalApi
+  Value getProtobufValue(int columnIndex);
+}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java
@@ -109,7 +109,7 @@ public ResultSet get() {
     }
   }
 
-  private static class PrePopulatedResultSet implements ResultSet {
+  private static class PrePopulatedResultSet implements ProtobufResultSet {
     private final List<Struct> rows;
     private final Type type;
     private int index = -1;
@@ -137,6 +137,19 @@ public boolean next() throws SpannerException {
       return ++index < rows.size();
     }
 
+    @Override
+    public boolean canGetProtobufValue(int columnIndex) {
+      return !closed && index >= 0 && index < rows.size();
+    }
+
+    @Override
+    public com.google.protobuf.Value getProtobufValue(int columnIndex) {
+      Preconditions.checkState(!closed, "ResultSet is closed");
+      Preconditions.checkState(index >= 0, "Must be preceded by a next() call");
+      Preconditions.checkElementIndex(index, rows.size(), "All rows have been yielded");
+      return getValue(columnIndex).toProto();
+    }
+
     @Override
     public Struct getCurrentRowAsStruct() {
       Preconditions.checkState(!closed, "ResultSet is closed");
-Original file line number
+Diff line change
@@ Expand Up @@
         DecodeOption(DecodeMode decodeMode) {
           this.decodeMode = Preconditions.checkNotNull(decodeMode, "DecodeMode cannot be null");
-          ;
         }
         @Override
@@ Expand Down @@