feat: support Read RPC OrderBy (#3180)

Author: mayurkale22

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -902,6 +902,9 @@ ResultSet readInternalWithOptions(
     if (readOptions.hasDataBoostEnabled()) {
       builder.setDataBoostEnabled(readOptions.dataBoostEnabled());
     }
+    if (readOptions.hasOrderBy()) {
+      builder.setOrderBy(readOptions.orderBy());
+    }
     if (readOptions.hasDirectedReadOptions()) {
       builder.setDirectedReadOptions(readOptions.directedReadOptions());
     } else if (defaultDirectedReadOptions != null) {

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java

@@ -18,6 +18,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.spanner.v1.DirectedReadOptions;
+import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import java.io.Serializable;
 import java.time.Duration;
@@ -51,6 +52,29 @@ public static RpcPriority fromProto(Priority proto) {
     }
   }
 
+  /**
+   * OrderBy for an RPC invocation. The default orderby is {@link #PRIMARY_KEY}. This enum can be
+   * used to control the order in which rows are returned from a read.
+   */
+  public enum RpcOrderBy {
+    PRIMARY_KEY(OrderBy.ORDER_BY_PRIMARY_KEY),
+    NO_ORDER(OrderBy.ORDER_BY_NO_ORDER),
+    UNSPECIFIED(OrderBy.ORDER_BY_UNSPECIFIED);
+
+    private final OrderBy proto;
+
+    RpcOrderBy(OrderBy proto) {
+      this.proto = Preconditions.checkNotNull(proto);
+    }
+
+    public static RpcOrderBy fromProto(OrderBy proto) {
+      for (RpcOrderBy e : RpcOrderBy.values()) {
+        if (e.proto.equals(proto)) return e;
+      }
+      return RpcOrderBy.UNSPECIFIED;
+    }
+  }
+
   /** Marker interface to mark options applicable to both Read and Query operations */
   public interface ReadAndQueryOption extends ReadOption, QueryOption {}
 
@@ -131,6 +155,11 @@ public static ReadOption limit(long limit) {
     return new LimitOption(limit);
   }
 
+  /** Specifies the order_by to use for the RPC. */
+  public static ReadOption orderBy(RpcOrderBy orderBy) {
+    return new OrderByOption(orderBy);
+  }
+
   /**
    * Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
    * PartialResultSet} chunks for read and query. The data size of each chunk depends on the server
@@ -439,6 +468,7 @@ void appendToOptions(Options options) {
   private Boolean dataBoostEnabled;
   private DirectedReadOptions directedReadOptions;
   private DecodeMode decodeMode;
+  private RpcOrderBy orderBy;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -567,6 +597,14 @@ DecodeMode decodeMode() {
     return decodeMode;
   }
 
+  boolean hasOrderBy() {
+    return orderBy != null;
+  }
+
+  OrderBy orderBy() {
+    return orderBy == null ? null : orderBy.proto;
+  }
+
   @Override
   public String toString() {
     StringBuilder b = new StringBuilder();
@@ -620,6 +658,9 @@ public String toString() {
     if (decodeMode != null) {
       b.append("decodeMode: ").append(decodeMode).append(' ');
     }
+    if (orderBy != null) {
+      b.append("orderBy: ").append(orderBy).append(' ');
+    }
     return b.toString();
   }
 
@@ -658,7 +699,8 @@ public boolean equals(Object o) {
         && Objects.equals(withOptimisticLock(), that.withOptimisticLock())
         && Objects.equals(withExcludeTxnFromChangeStreams(), that.withExcludeTxnFromChangeStreams())
         && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
-        && Objects.equals(directedReadOptions(), that.directedReadOptions());
+        && Objects.equals(directedReadOptions(), that.directedReadOptions())
+        && Objects.equals(orderBy(), that.orderBy());
   }
 
   @Override
@@ -715,6 +757,9 @@ public int hashCode() {
     if (decodeMode != null) {
       result = 31 * result + decodeMode.hashCode();
     }
+    if (orderBy != null) {
+      result = 31 * result + orderBy.hashCode();
+    }
     return result;
   }
 
@@ -795,6 +840,19 @@ void appendToOptions(Options options) {
     }
   }
 
+  static class OrderByOption extends InternalOption implements ReadOption {
+    private final RpcOrderBy orderBy;
+
+    OrderByOption(RpcOrderBy orderBy) {
+      this.orderBy = orderBy;
+    }
+
+    @Override
+    void appendToOptions(Options options) {
+      options.orderBy = orderBy;
+    }
+  }
+
   static final class DataBoostQueryOption extends InternalOption implements ReadAndQueryOption {
 
     private final Boolean dataBoostEnabled;

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java

@@ -33,8 +33,11 @@
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions;
 import com.google.spanner.v1.RequestOptions.Priority;
+import com.google.spanner.v1.SessionName;
 import com.google.spanner.v1.TransactionSelector;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -223,6 +226,21 @@ public void testGetExecuteSqlRequestBuilderWithDataBoost() {
     assertTrue(request.getDataBoostEnabled());
   }
 
+  @Test
+  public void testGetReadRequestBuilderWithOrderBy() {
+    ReadRequest request =
+        ReadRequest.newBuilder()
+            .setSession(
+                SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]").toString())
+            .setTransaction(TransactionSelector.newBuilder().build())
+            .setTable("table110115790")
+            .setIndex("index100346066")
+            .addAllColumns(new ArrayList<String>())
+            .setOrderByValue(2)
+            .build();
+    assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
+  }
+
   @Test
   public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
     ExecuteBatchDmlRequest.Builder request =

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -52,6 +52,7 @@
 import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Options.RpcOrderBy;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
@@ -89,6 +90,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
@@ -1722,6 +1724,27 @@ public void testExecuteReadWithTag() {
     assertThat(request.getRequestOptions().getTransactionTag()).isEmpty();
   }
 
+  @Test
+  public void testExecuteReadWithOrderByOption() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    try (ResultSet resultSet =
+        client
+            .singleUse()
+            .read(
+                READ_TABLE_NAME,
+                KeySet.singleKey(Key.of(1L)),
+                READ_COLUMN_NAMES,
+                Options.orderBy(RpcOrderBy.NO_ORDER))) {
+      consumeResults(resultSet);
+    }
+
+    List<ReadRequest> requests = mockSpanner.getRequestsOfType(ReadRequest.class);
+    assertThat(requests).hasSize(1);
+    ReadRequest request = requests.get(0);
+    assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
+  }
+
   @Test
   public void testExecuteReadWithDirectedReadOptions() {
     DatabaseClient client =

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java

@@ -18,16 +18,19 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.Options.RpcOrderBy;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
 import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
+import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -79,14 +82,16 @@ public void allOptionsPresent() {
             Options.limit(10),
             Options.prefetchChunks(1),
             Options.dataBoostEnabled(true),
-            Options.directedRead(DIRECTED_READ_OPTIONS));
+            Options.directedRead(DIRECTED_READ_OPTIONS),
+            Options.orderBy(RpcOrderBy.NO_ORDER));
     assertThat(options.hasLimit()).isTrue();
     assertThat(options.limit()).isEqualTo(10);
     assertThat(options.hasPrefetchChunks()).isTrue();
     assertThat(options.prefetchChunks()).isEqualTo(1);
     assertThat(options.hasDataBoostEnabled()).isTrue();
     assertTrue(options.dataBoostEnabled());
     assertTrue(options.hasDirectedReadOptions());
+    assertTrue(options.hasOrderBy());
     assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
   }
 
@@ -101,6 +106,7 @@ public void allOptionsAbsent() {
     assertThat(options.hasTag()).isFalse();
     assertThat(options.hasDataBoostEnabled()).isFalse();
     assertThat(options.hasDirectedReadOptions()).isFalse();
+    assertThat(options.hasOrderBy()).isFalse();
     assertNull(options.withExcludeTxnFromChangeStreams());
     assertThat(options.toString()).isEqualTo("");
     assertThat(options.equals(options)).isTrue();
@@ -182,7 +188,8 @@ public void readOptionsTest() {
             Options.limit(limit),
             Options.tag(tag),
             Options.dataBoostEnabled(true),
-            Options.directedRead(DIRECTED_READ_OPTIONS));
+            Options.directedRead(DIRECTED_READ_OPTIONS),
+            Options.orderBy(RpcOrderBy.NO_ORDER));
 
     assertThat(options.toString())
         .isEqualTo(
@@ -197,10 +204,14 @@ public void readOptionsTest() {
                 + " "
                 + "directedReadOptions: "
                 + DIRECTED_READ_OPTIONS
+                + " "
+                + "orderBy: "
+                + RpcOrderBy.NO_ORDER
                 + " ");
     assertThat(options.tag()).isEqualTo(tag);
     assertEquals(dataBoost, options.dataBoostEnabled());
     assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
+    assertEquals(OrderBy.ORDER_BY_NO_ORDER, options.orderBy());
   }
 
   @Test
@@ -354,6 +365,24 @@ public void testTransactionOptionsPriority() {
     assertEquals("priority: " + priority + " ", options.toString());
   }
 
+  @Test
+  public void testReadOptionsOrderBy() {
+    RpcOrderBy orderBy = RpcOrderBy.NO_ORDER;
+    Options options = Options.fromReadOptions(Options.orderBy(orderBy));
+    assertTrue(options.hasOrderBy());
+    assertEquals("orderBy: " + orderBy + " ", options.toString());
+  }
+
+  @Test
+  public void testReadOptionsWithOrderByEquality() {
+    Options optionsWithNoOrderBy1 = Options.fromReadOptions(Options.orderBy(RpcOrderBy.NO_ORDER));
+    Options optionsWithNoOrderBy2 = Options.fromReadOptions(Options.orderBy(RpcOrderBy.NO_ORDER));
+    assertTrue(optionsWithNoOrderBy1.equals(optionsWithNoOrderBy2));
+
+    Options optionsWithPkOrder = Options.fromReadOptions(Options.orderBy(RpcOrderBy.PRIMARY_KEY));
+    assertFalse(optionsWithNoOrderBy1.equals(optionsWithPkOrder));
+  }
+
   @Test
   public void testQueryOptionsPriority() {
     RpcPriority priority = RpcPriority.MEDIUM;