diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index a424a93115a..92ebf006d29 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/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) { 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 index 58123ae36b8..9c3257586fb 100644 --- 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 @@ -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; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java index 65f27d55810..ce7d6b300d1 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java +++ b/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()) + .setOrderByValue(2) + .build(); + assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy()); + } + @Test public void testGetExecuteBatchDmlRequestBuilderWithPriority() { ExecuteBatchDmlRequest.Builder request = diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 6443c904b86..3dacd43a55e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/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 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 = diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java index 8c9a5d957e8..38b7a121731 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java +++ b/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,7 +82,8 @@ 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(); @@ -87,6 +91,7 @@ public void allOptionsPresent() { 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;