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 b8d2315b98..3b6f00480b 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 @@ -47,6 +47,7 @@ import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import com.google.spanner.v1.PartialResultSet; import com.google.spanner.v1.ReadRequest; +import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; @@ -557,6 +558,14 @@ QueryOptions buildQueryOptions(QueryOptions requestOptions) { return builder.build(); } + RequestOptions buildRequestOptions(Options options) { + RequestOptions.Builder builder = RequestOptions.newBuilder(); + if (options.hasPriority()) { + builder.setPriority(options.priority()); + } + return builder.build(); + } + ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( Statement statement, QueryMode queryMode, Options options, boolean withTransactionSelector) { ExecuteSqlRequest.Builder builder = @@ -580,6 +589,7 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( } builder.setSeqno(getSeqNo()); builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions())); + builder.setRequestOptions(buildRequestOptions(options)); return builder; } @@ -610,6 +620,7 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder( builder.setTransaction(selector); } builder.setSeqno(getSeqNo()); + builder.setRequestOptions(buildRequestOptions(options)); return builder; } @@ -760,6 +771,7 @@ CloseableIterator startStream(@Nullable ByteString resumeToken if (selector != null) { builder.setTransaction(selector); } + builder.setRequestOptions(buildRequestOptions(readOptions)); SpannerRpc.StreamingCall call = rpc.read(builder.build(), stream.consumer(), session.getOptions()); call.request(prefetchChunks); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index bd42aa307f..215d86c825 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.Options.UpdateOption; @@ -75,9 +76,21 @@ public interface DatabaseClient { * .set("LastName") * .to("Joel") * .build(); - * dbClient.writeWithOptions(Collections.singletonList(mutation)); + * dbClient.writeWithOptions( + * Collections.singletonList(mutation), + * Options.priority(RpcPriority.HIGH)); * } * + * Options for a transaction can include: + * + * + * * @return a response with the timestamp at which the write was committed */ CommitResponse writeWithOptions(Iterable mutations, TransactionOption... options) @@ -138,9 +151,21 @@ CommitResponse writeWithOptions(Iterable mutations, TransactionOption. * .set("LastName") * .to("Joel") * .build(); - * dbClient.writeAtLeastOnce(Collections.singletonList(mutation)); + * dbClient.writeAtLeastOnceWithOptions( + * Collections.singletonList(mutation), + * Options.priority(RpcPriority.LOW)); * } * + * Options for a transaction can include: + * + *
    + *
  • {@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link + * RpcPriority} to use for the commit request of the transaction. The priority will not be + * applied to any other requests on the transaction. + *
  • {@link Options#commitStats()}: Request that the server includes commit statistics in the + * {@link CommitResponse}. + *
+ * * @return a response with the timestamp at which the write was committed */ CommitResponse writeAtLeastOnceWithOptions( @@ -308,6 +333,16 @@ CommitResponse writeAtLeastOnceWithOptions( * } * }); * + * + * Options for a transaction can include: + * + *
    + *
  • {@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link + * RpcPriority} to use for the commit request of the transaction. The priority will not be + * applied to any other requests on the transaction. + *
  • {@link Options#commitStats()}: Request that the server includes commit statistics in the + * {@link CommitResponse}. + *
*/ TransactionRunner readWriteTransaction(TransactionOption... options); @@ -338,6 +373,16 @@ CommitResponse writeAtLeastOnceWithOptions( * } * } * } + * + * Options for a transaction can include: + * + *
    + *
  • {@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link + * RpcPriority} to use for the commit request of the transaction. The priority will not be + * applied to any other requests on the transaction. + *
  • {@link Options#commitStats()}: Request that the server includes commit statistics in the + * {@link CommitResponse}. + *
*/ TransactionManager transactionManager(TransactionOption... options); @@ -371,6 +416,16 @@ CommitResponse writeAtLeastOnceWithOptions( * }, * executor); * + * + * Options for a transaction can include: + * + *
    + *
  • {@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link + * RpcPriority} to use for the commit request of the transaction. The priority will not be + * applied to any other requests on the transaction. + *
  • {@link Options#commitStats()}: Request that the server includes commit statistics in the + * {@link CommitResponse}. + *
*/ AsyncRunner runAsync(TransactionOption... options); @@ -459,6 +514,18 @@ CommitResponse writeAtLeastOnceWithOptions( * } * } * } + * + * Options for a transaction can include: + * + *

Options for a transaction can include: + * + *

    + *
  • {@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link + * RpcPriority} to use for the commit request of the transaction. The priority will not be + * applied to any other requests on the transaction. + *
  • {@link Options#commitStats()}: Request that the server includes commit statistics in the + * {@link CommitResponse}. + *
*/ AsyncTransactionManager transactionManagerAsync(TransactionOption... options); 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 38932bf8a8..8435680eab 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 @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.common.base.Preconditions; +import com.google.spanner.v1.RequestOptions.Priority; import java.io.Serializable; import java.util.Objects; @@ -24,6 +25,22 @@ public final class Options implements Serializable { private static final long serialVersionUID = 8067099123096783941L; + /** + * Priority for an RPC invocation. The default priority is {@link #HIGH}. This enum can be used to + * set a lower priority for a specific RPC invocation. + */ + public enum RpcPriority { + LOW(Priority.PRIORITY_LOW), + MEDIUM(Priority.PRIORITY_MEDIUM), + HIGH(Priority.PRIORITY_HIGH); + + private final Priority proto; + + private RpcPriority(Priority proto) { + this.proto = Preconditions.checkNotNull(proto); + } + } + /** Marker interface to mark options applicable to both Read and Query operations */ public interface ReadAndQueryOption extends ReadOption, QueryOption {} @@ -79,6 +96,11 @@ public static ReadAndQueryOption bufferRows(int bufferRows) { return new BufferRowsOption(bufferRows); } + /** Specifies the priority to use for the RPC. */ + public static ReadQueryUpdateTransactionOption priority(RpcPriority priority) { + return new PriorityOption(priority); + } + /** * Specifying this will cause the list operations to fetch at most this many records in a page. */ @@ -158,6 +180,20 @@ void appendToOptions(Options options) { } } + static final class PriorityOption extends InternalOption + implements ReadQueryUpdateTransactionOption { + private final RpcPriority priority; + + PriorityOption(RpcPriority priority) { + this.priority = priority; + } + + @Override + void appendToOptions(Options options) { + options.priority = priority; + } + } + private boolean withCommitStats; private Long limit; private Integer prefetchChunks; @@ -165,6 +201,7 @@ void appendToOptions(Options options) { private Integer pageSize; private String pageToken; private String filter; + private RpcPriority priority; // Construction is via factory methods below. private Options() {} @@ -221,6 +258,14 @@ String filter() { return filter; } + boolean hasPriority() { + return priority != null; + } + + Priority priority() { + return priority == null ? null : priority.proto; + } + @Override public String toString() { StringBuilder b = new StringBuilder(); @@ -242,6 +287,9 @@ public String toString() { if (filter != null) { b.append("filter: ").append(filter).append(' '); } + if (priority != null) { + b.append("priority: ").append(priority).append(' '); + } return b.toString(); } @@ -271,7 +319,8 @@ public boolean equals(Object o) { && (!hasPageSize() && !that.hasPageSize() || hasPageSize() && that.hasPageSize() && Objects.equals(pageSize(), that.pageSize())) && Objects.equals(pageToken(), that.pageToken()) - && Objects.equals(filter(), that.filter()); + && Objects.equals(filter(), that.filter()) + && Objects.equals(priority(), that.priority()); } @Override @@ -298,6 +347,9 @@ public int hashCode() { if (filter != null) { result = 31 * result + filter.hashCode(); } + if (priority != null) { + result = 31 * result + priority.hashCode(); + } return result; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java index b3aa4dea3d..697f3e6351 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java @@ -26,12 +26,14 @@ import com.google.api.gax.rpc.UnavailableException; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Stopwatch; import com.google.common.base.Ticker; import com.google.protobuf.ByteString; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; @@ -162,8 +164,8 @@ private ExecuteSqlRequest resumeOrRestartRequest( } } - private ExecuteSqlRequest newTransactionRequestFrom( - final Statement statement, final Options options) { + @VisibleForTesting + ExecuteSqlRequest newTransactionRequestFrom(final Statement statement, final Options options) { ByteString transactionId = initTransaction(); final TransactionSelector transactionSelector = @@ -179,6 +181,11 @@ private ExecuteSqlRequest newTransactionRequestFrom( builder.setResumeToken(ByteString.EMPTY); + if (options.hasPriority()) { + builder.setRequestOptions( + RequestOptions.newBuilder().setPriority(options.priority()).build()); + } + return builder.build(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 51ec7b4d50..d122310eb7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -37,6 +37,7 @@ import com.google.protobuf.Empty; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; import io.opencensus.common.Scope; @@ -160,9 +161,10 @@ public CommitResponse writeAtLeastOnceWithOptions( Iterable mutations, TransactionOption... transactionOptions) throws SpannerException { setActive(null); + Options commitRequestOptions = Options.fromTransactionOptions(transactionOptions); List mutationsProto = new ArrayList<>(); Mutation.toProto(mutations, mutationsProto); - final CommitRequest request = + final CommitRequest.Builder requestBuilder = CommitRequest.newBuilder() .setSession(name) .setReturnCommitStats( @@ -170,12 +172,15 @@ public CommitResponse writeAtLeastOnceWithOptions( .addAllMutations(mutationsProto) .setSingleUseTransaction( TransactionOptions.newBuilder() - .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())) - .build(); + .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())); + if (commitRequestOptions.hasPriority()) { + requestBuilder.setRequestOptions( + RequestOptions.newBuilder().setPriority(commitRequestOptions.priority()).build()); + } Span span = tracer.spanBuilder(SpannerImpl.COMMIT).startSpan(); try (Scope s = tracer.withSpan(span)) { com.google.spanner.v1.CommitResponse response = - spanner.getRpc().commit(request, this.options); + spanner.getRpc().commit(requestBuilder.build(), this.options); return new CommitResponse(response); } catch (RuntimeException e) { TraceUtil.setWithFailure(span, e); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index daa87fea02..860348adff 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -43,6 +43,7 @@ import com.google.spanner.v1.ExecuteBatchDmlResponse; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.ResultSet; import com.google.spanner.v1.RollbackRequest; import com.google.spanner.v1.Transaction; @@ -298,6 +299,10 @@ ApiFuture commitAsync() { CommitRequest.newBuilder() .setSession(session.getName()) .setReturnCommitStats(options.withCommitStats()); + if (options.hasPriority()) { + builder.setRequestOptions( + RequestOptions.newBuilder().setPriority(options.priority()).build()); + } synchronized (lock) { if (transactionIdFuture == null && transactionId == null && runningAsyncOperations == 0) { finishOps = SettableApiFuture.create(); @@ -344,6 +349,10 @@ public void run() { requestBuilder.setTransactionId( transactionId == null ? transactionIdFuture.get() : transactionId); } + if (options.hasPriority()) { + requestBuilder.setRequestOptions( + RequestOptions.newBuilder().setPriority(options.priority()).build()); + } final CommitRequest commitRequest = requestBuilder.build(); span.addAnnotation("Starting Commit"); final Span opSpan = 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 bc0e475f44..795a26fff7 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 @@ -17,17 +17,23 @@ package com.google.cloud.spanner; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import com.google.api.gax.core.ExecutorProvider; +import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.spanner.v1.ExecuteBatchDmlRequest; 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.RequestOptions; +import com.google.spanner.v1.RequestOptions.Priority; import com.google.spanner.v1.TransactionSelector; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.List; import org.junit.Before; import org.junit.Test; @@ -114,4 +120,46 @@ public void executeSqlRequestBuilderWithQueryOptions() { assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR"); assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("2.0"); } + + @Test + public void testBuildRequestionOptions() { + RequestOptions requestOptions = context.buildRequestOptions(Options.fromQueryOptions()); + assertEquals(RequestOptions.Priority.PRIORITY_UNSPECIFIED, requestOptions.getPriority()); + } + + @Test + public void testBuildRequestOptionsWithPriority() { + RequestOptions requestOptionsHighPriority = + context.buildRequestOptions(Options.fromQueryOptions(Options.priority(RpcPriority.HIGH))); + assertEquals(RequestOptions.Priority.PRIORITY_HIGH, requestOptionsHighPriority.getPriority()); + + RequestOptions requestOptionsMediumPriority = + context.buildRequestOptions(Options.fromQueryOptions(Options.priority(RpcPriority.MEDIUM))); + assertEquals( + RequestOptions.Priority.PRIORITY_MEDIUM, requestOptionsMediumPriority.getPriority()); + + RequestOptions requestOptionsLowPriority = + context.buildRequestOptions(Options.fromQueryOptions(Options.priority(RpcPriority.LOW))); + assertEquals(RequestOptions.Priority.PRIORITY_LOW, requestOptionsLowPriority.getPriority()); + } + + @Test + public void testGetExecuteSqlRequestBuilderWithPriority() { + ExecuteSqlRequest.Builder request = + context.getExecuteSqlRequestBuilder( + Statement.of("SELECT * FROM FOO"), + QueryMode.NORMAL, + Options.fromQueryOptions(Options.priority(RpcPriority.MEDIUM)), + false); + assertEquals(Priority.PRIORITY_MEDIUM, request.getRequestOptions().getPriority()); + } + + @Test + public void testGetExecuteBatchDmlRequestBuilderWithPriority() { + ExecuteBatchDmlRequest.Builder request = + context.getExecuteBatchDmlRequestBuilder( + Collections.singleton(Statement.of("SELECT * FROM FOO")), + Options.fromQueryOptions(Options.priority(RpcPriority.LOW))); + assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority()); + } } 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 5cfdb5ea62..867d8fb242 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 @@ -16,12 +16,17 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_COLUMN_NAMES; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_RESULTSET; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_TABLE_NAME; import static com.google.cloud.spanner.MockSpannerTestUtil.SELECT1; import static com.google.cloud.spanner.SpannerApiFutures.get; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -37,8 +42,11 @@ import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.AsyncRunner.AsyncWork; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; +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.RpcPriority; import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.SessionPool.PooledSessionFuture; @@ -48,9 +56,13 @@ import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.SettableFuture; import com.google.protobuf.AbstractMessage; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteBatchDmlRequest; 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.RequestOptions.Priority; import io.grpc.Context; import io.grpc.Server; import io.grpc.Status; @@ -106,6 +118,8 @@ public static void startStaticServer() throws IOException { mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); mockSpanner.putStatementResult( StatementResult.query(SELECT1, MockSpannerTestUtil.SELECT1_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); mockSpanner.putStatementResult( StatementResult.exception( INVALID_UPDATE_STATEMENT, @@ -168,6 +182,28 @@ public void testWrite() { Arrays.asList( Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build())); assertNotNull(timestamp); + + List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(commitRequests).hasSize(1); + CommitRequest commit = commitRequests.get(0); + assertNotNull(commit.getRequestOptions()); + assertEquals(Priority.PRIORITY_UNSPECIFIED, commit.getRequestOptions().getPriority()); + } + + @Test + public void testWriteWithOptions() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + client.writeWithOptions( + Arrays.asList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()), + Options.priority(RpcPriority.HIGH)); + + List commits = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(commits).hasSize(1); + CommitRequest commit = commits.get(0); + assertNotNull(commit.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, commit.getRequestOptions().getPriority()); } @Test @@ -207,6 +243,32 @@ public void testWriteAtLeastOnceWithCommitStats() { assertNotNull(response); assertNotNull(response.getCommitTimestamp()); assertNotNull(response.getCommitStats()); + + List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(commitRequests).hasSize(1); + CommitRequest commit = commitRequests.get(0); + assertNotNull(commit.getSingleUseTransaction()); + assertTrue(commit.getSingleUseTransaction().hasReadWrite()); + assertNotNull(commit.getRequestOptions()); + assertEquals(Priority.PRIORITY_UNSPECIFIED, commit.getRequestOptions().getPriority()); + } + + @Test + public void testWriteAtLeastOnceWithOptions() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + client.writeAtLeastOnceWithOptions( + Arrays.asList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()), + Options.priority(RpcPriority.LOW)); + + List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(commitRequests).hasSize(1); + CommitRequest commit = commitRequests.get(0); + assertNotNull(commit.getSingleUseTransaction()); + assertTrue(commit.getSingleUseTransaction().hasReadWrite()); + assertNotNull(commit.getRequestOptions()); + assertEquals(Priority.PRIORITY_LOW, commit.getRequestOptions().getPriority()); } @Test @@ -1642,4 +1704,238 @@ public void testTransactionManagerAsync_usesOptions() { verify(session).transactionManagerAsync(option); } + + @Test + public void testExecuteQueryWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet resultSet = + client.singleUse().executeQuery(SELECT1, Options.priority(RpcPriority.HIGH))) { + while (resultSet.next()) {} + } + + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertThat(requests).hasSize(1); + ExecuteSqlRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testExecuteReadWithPriority() { + 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.priority(RpcPriority.HIGH))) { + while (resultSet.next()) {} + } + + List requests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertThat(requests).hasSize(1); + ReadRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testReadWriteExecuteQueryWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.priority(RpcPriority.HIGH))) { + while (resultSet.next()) {} + } + return null; + } + }); + + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertThat(requests).hasSize(1); + ExecuteSqlRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testReadWriteExecuteReadWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + try (ResultSet resultSet = + transaction.read( + READ_TABLE_NAME, + KeySet.singleKey(Key.of(1L)), + READ_COLUMN_NAMES, + Options.priority(RpcPriority.HIGH))) { + while (resultSet.next()) {} + } + return null; + } + }); + + List requests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertThat(requests).hasSize(1); + ReadRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testExecuteUpdateWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + runner.run( + new TransactionCallable() { + @Override + public Long run(TransactionContext transaction) throws Exception { + return transaction.executeUpdate(UPDATE_STATEMENT, Options.priority(RpcPriority.HIGH)); + } + }); + + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertThat(requests).hasSize(1); + ExecuteSqlRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testBatchUpdateWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + runner.run( + new TransactionCallable() { + @Override + public long[] run(TransactionContext transaction) throws Exception { + return transaction.batchUpdate( + Arrays.asList(UPDATE_STATEMENT), Options.priority(RpcPriority.HIGH)); + } + }); + + List requests = + mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class); + assertThat(requests).hasSize(1); + ExecuteBatchDmlRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testPartitionedDMLWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + client.executePartitionedUpdate(UPDATE_STATEMENT, Options.priority(RpcPriority.HIGH)); + + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertThat(requests).hasSize(1); + ExecuteSqlRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testCommitWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(Options.priority(RpcPriority.HIGH)); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.buffer(Mutation.delete("TEST", KeySet.all())); + return null; + } + }); + + List requests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(requests).hasSize(1); + CommitRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testTransactionManagerCommitWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionManager manager = client.transactionManager(Options.priority(RpcPriority.HIGH)); + TransactionContext transaction = manager.begin(); + transaction.buffer(Mutation.delete("TEST", KeySet.all())); + manager.commit(); + + List requests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(requests).hasSize(1); + CommitRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testAsyncRunnerCommitWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + AsyncRunner runner = client.runAsync(Options.priority(RpcPriority.HIGH)); + get( + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.buffer(Mutation.delete("TEST", KeySet.all())); + return ApiFutures.immediateFuture(null); + } + }, + executor)); + + List requests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(requests).hasSize(1); + CommitRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } + + @Test + public void testAsyncTransactionManagerCommitWithPriority() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (AsyncTransactionManager manager = + client.transactionManagerAsync(Options.priority(RpcPriority.HIGH))) { + TransactionContextFuture transaction = manager.beginAsync(); + get( + transaction + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + txn.buffer(Mutation.delete("TEST", KeySet.all())); + return ApiFutures.immediateFuture(null); + } + }, + executor) + .commitAsync()); + } + + List requests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertThat(requests).hasSize(1); + CommitRequest request = requests.get(0); + assertNotNull(request.getRequestOptions()); + assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority()); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java index 00c2f1ab47..d8a780ef7d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java @@ -1408,7 +1408,7 @@ public Long run(TransactionContext transaction) throws Exception { assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(2L); assertThat(mockSpanner.countRequestsOfType(CommitRequest.class)).isEqualTo(1L); assertThat(countTransactionsStarted()).isEqualTo(1); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); assertThat(requests.get(0)).isInstanceOf(ExecuteSqlRequest.class); assertThat(((ExecuteSqlRequest) requests.get(0)).getSql()) .isEqualTo(UPDATE_STATEMENT.getSql()); @@ -1439,7 +1439,7 @@ public Void run(TransactionContext transaction) throws Exception { assertThat(countRequests(CommitRequest.class)).isEqualTo(1); assertThat(countTransactionsStarted()).isEqualTo(1); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); assertThat(requests.get(0)).isInstanceOf(ExecuteSqlRequest.class); ExecuteSqlRequest request1 = (ExecuteSqlRequest) requests.get(0); assertThat(request1.getSql()).isEqualTo(SELECT1_UNION_ALL_SELECT2.getSql()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java index 0b68829604..f09d86a429 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java @@ -1941,14 +1941,15 @@ public void clearRequests() { this.requests.clear(); } - public List getRequestsOfType(Class type) { - List res = new ArrayList<>(); - for (AbstractMessage m : this.requests) { - if (m.getClass().equals(type)) { - res.add(m); + @SuppressWarnings("unchecked") + public List getRequestsOfType(Class type) { + List result = new ArrayList<>(); + for (AbstractMessage message : this.requests) { + if (message.getClass().equals(type)) { + result.add((T) message); } } - return res; + return result; } public Iterable> getRequestTypes() { 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 db642aedd1..4ab08a4785 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,10 +18,14 @@ 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.assertTrue; import static org.junit.Assert.fail; +import com.google.cloud.spanner.Options.RpcPriority; +import com.google.spanner.v1.RequestOptions.Priority; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -86,6 +90,7 @@ public void allOptionsAbsent() { assertThat(options.hasPrefetchChunks()).isFalse(); assertThat(options.hasFilter()).isFalse(); assertThat(options.hasPageToken()).isFalse(); + assertThat(options.hasPriority()).isFalse(); assertThat(options.toString()).isEqualTo(""); assertThat(options.equals(options)).isTrue(); assertThat(options.equals(null)).isFalse(); @@ -277,21 +282,232 @@ public void testTransactionOptions_withCommitStatsAndOtherOptionHashCode() { } @Test - public void testFromUpdateOptions() { - Options opts = Options.fromUpdateOptions(); - assertThat(opts.toString()).isEqualTo(""); + public void testTransactionOptionsPriority() { + RpcPriority priority = RpcPriority.HIGH; + Options options = Options.fromTransactionOptions(Options.priority(priority)); + assertTrue(options.hasPriority()); + assertEquals("priority: " + priority + " ", options.toString()); + } + + @Test + public void testQueryOptionsPriority() { + RpcPriority priority = RpcPriority.MEDIUM; + Options options = Options.fromQueryOptions(Options.priority(priority)); + assertTrue(options.hasPriority()); + assertEquals("priority: " + priority + " ", options.toString()); + } + + @Test + public void testReadOptionsPriority() { + RpcPriority priority = RpcPriority.LOW; + Options options = Options.fromReadOptions(Options.priority(priority)); + assertTrue(options.hasPriority()); + assertEquals("priority: " + priority + " ", options.toString()); + } + + @Test + public void testUpdateOptionsPriority() { + RpcPriority priority = RpcPriority.LOW; + Options options = Options.fromUpdateOptions(Options.priority(priority)); + assertTrue(options.hasPriority()); + assertEquals("priority: " + priority + " ", options.toString()); + } + + @Test + public void testTransactionOptionsHashCode() { + Options option1 = Options.fromTransactionOptions(); + Options option2 = Options.fromTransactionOptions(); + assertEquals(option1.hashCode(), option2.hashCode()); + } + + @Test + public void testTransactionOptionsWithPriorityEquality() { + Options optionsWithHighPriority1 = + Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = + Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + assertTrue(optionsWithHighPriority1.equals(optionsWithHighPriority2)); + + Options optionsWithMediumPriority = + Options.fromTransactionOptions(Options.priority(RpcPriority.MEDIUM)); + assertFalse(optionsWithHighPriority1.equals(optionsWithMediumPriority)); + } + + @Test + public void testTransactionOptionsWithPriorityHashCode() { + Options optionsWithHighPriority1 = + Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = + Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + assertEquals(optionsWithHighPriority1.hashCode(), optionsWithHighPriority2.hashCode()); + + Options optionsWithMediumPriority = + Options.fromTransactionOptions(Options.priority(RpcPriority.MEDIUM)); + assertNotEquals(optionsWithHighPriority1.hashCode(), optionsWithMediumPriority.hashCode()); } @Test public void testUpdateOptionsEquality() { - Options o1; - Options o2; + Options option1 = Options.fromUpdateOptions(); + Options option2 = Options.fromUpdateOptions(); + assertTrue(option1.equals(option2)); + } - o1 = Options.fromUpdateOptions(); - o2 = Options.fromUpdateOptions(); - assertThat(o1.equals(o2)).isTrue(); + @Test + public void testUpdateOptionsHashCode() { + Options option1 = Options.fromUpdateOptions(); + Options option2 = Options.fromUpdateOptions(); + assertEquals(option1.hashCode(), option2.hashCode()); + } - o2 = Options.fromReadOptions(Options.prefetchChunks(1)); - assertThat(o1.equals(o2)).isFalse(); + @Test + public void testUpdateOptionsWithPriorityEquality() { + Options optionsWithHighPriority1 = + Options.fromUpdateOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = + Options.fromUpdateOptions(Options.priority(RpcPriority.HIGH)); + assertTrue(optionsWithHighPriority1.equals(optionsWithHighPriority2)); + + Options optionsWithMediumPriority = + Options.fromUpdateOptions(Options.priority(RpcPriority.MEDIUM)); + assertFalse(optionsWithHighPriority1.equals(optionsWithMediumPriority)); + } + + @Test + public void testUpdateOptionsWithPriorityHashCode() { + Options optionsWithHighPriority1 = + Options.fromUpdateOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = + Options.fromUpdateOptions(Options.priority(RpcPriority.HIGH)); + assertEquals(optionsWithHighPriority1.hashCode(), optionsWithHighPriority2.hashCode()); + + Options optionsWithMediumPriority = + Options.fromUpdateOptions(Options.priority(RpcPriority.MEDIUM)); + assertNotEquals(optionsWithHighPriority1.hashCode(), optionsWithMediumPriority.hashCode()); + } + + @Test + public void testQueryOptionsEquality() { + Options option1 = Options.fromQueryOptions(); + Options option2 = Options.fromQueryOptions(); + assertTrue(option1.equals(option2)); + } + + @Test + public void testQueryOptionsHashCode() { + Options option1 = Options.fromQueryOptions(); + Options option2 = Options.fromQueryOptions(); + assertEquals(option1.hashCode(), option2.hashCode()); + } + + @Test + public void testQueryOptionsWithPriorityEquality() { + Options optionsWithHighPriority1 = Options.fromQueryOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = Options.fromQueryOptions(Options.priority(RpcPriority.HIGH)); + assertTrue(optionsWithHighPriority1.equals(optionsWithHighPriority2)); + + Options optionsWithMediumPriority = + Options.fromQueryOptions(Options.priority(RpcPriority.MEDIUM)); + assertFalse(optionsWithHighPriority1.equals(optionsWithMediumPriority)); + + Options optionsWithHighPriorityAndBufferRows = + Options.fromQueryOptions(Options.priority(RpcPriority.HIGH), Options.bufferRows(10)); + assertFalse(optionsWithHighPriorityAndBufferRows.equals(optionsWithHighPriority1)); + } + + @Test + public void testQueryOptionsWithPriorityHashCode() { + Options optionsWithHighPriority1 = Options.fromQueryOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = Options.fromQueryOptions(Options.priority(RpcPriority.HIGH)); + assertEquals(optionsWithHighPriority1.hashCode(), optionsWithHighPriority2.hashCode()); + + Options optionsWithMediumPriority = + Options.fromQueryOptions(Options.priority(RpcPriority.MEDIUM)); + assertNotEquals(optionsWithHighPriority1.hashCode(), optionsWithMediumPriority.hashCode()); + + Options optionsWithHighPriorityAndBufferRows = + Options.fromQueryOptions(Options.priority(RpcPriority.HIGH), Options.bufferRows(10)); + assertNotEquals( + optionsWithHighPriorityAndBufferRows.hashCode(), optionsWithHighPriority1.hashCode()); + } + + @Test + public void testReadOptionsEquality() { + Options option1 = Options.fromReadOptions(); + Options option2 = Options.fromReadOptions(); + assertTrue(option1.equals(option2)); + } + + @Test + public void testReadOptionsHashCode() { + Options option1 = Options.fromReadOptions(); + Options option2 = Options.fromReadOptions(); + assertEquals(option1.hashCode(), option2.hashCode()); + } + + @Test + public void testReadOptionsWithPriorityEquality() { + Options optionsWithHighPriority1 = Options.fromReadOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = Options.fromReadOptions(Options.priority(RpcPriority.HIGH)); + assertTrue(optionsWithHighPriority1.equals(optionsWithHighPriority2)); + + Options optionsWithMediumPriority = + Options.fromReadOptions(Options.priority(RpcPriority.MEDIUM)); + assertFalse(optionsWithHighPriority1.equals(optionsWithMediumPriority)); + + Options optionsWithHighPriorityAndBufferRows = + Options.fromReadOptions(Options.priority(RpcPriority.HIGH), Options.bufferRows(10)); + assertFalse(optionsWithHighPriorityAndBufferRows.equals(optionsWithHighPriority1)); + } + + @Test + public void testReadOptionsWithPriorityHashCode() { + Options optionsWithHighPriority1 = Options.fromReadOptions(Options.priority(RpcPriority.HIGH)); + Options optionsWithHighPriority2 = Options.fromReadOptions(Options.priority(RpcPriority.HIGH)); + assertEquals(optionsWithHighPriority1.hashCode(), optionsWithHighPriority2.hashCode()); + + Options optionsWithMediumPriority = + Options.fromReadOptions(Options.priority(RpcPriority.MEDIUM)); + assertNotEquals(optionsWithHighPriority1.hashCode(), optionsWithMediumPriority.hashCode()); + + Options optionsWithHighPriorityAndBufferRows = + Options.fromReadOptions(Options.priority(RpcPriority.HIGH), Options.bufferRows(10)); + assertNotEquals( + optionsWithHighPriorityAndBufferRows.hashCode(), optionsWithHighPriority1.hashCode()); + } + + @Test + public void testFromUpdateOptions() { + Options options = Options.fromUpdateOptions(); + assertThat(options.toString()).isEqualTo(""); + } + + @Test + public void testTransactionOptions() { + RpcPriority prio = RpcPriority.HIGH; + Options opts = Options.fromTransactionOptions(Options.priority(prio)); + assertThat(opts.toString()).isEqualTo("priority: " + prio + " "); + assertThat(opts.priority()).isEqualTo(Priority.PRIORITY_HIGH); + } + + @Test + public void testTransactionOptionsDefaultEqual() { + Options options1 = Options.fromTransactionOptions(); + Options options2 = Options.fromTransactionOptions(); + assertEquals(options1, options2); + } + + @Test + public void testTransactionOptionsPriorityEquality() { + Options options1 = Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + Options options2 = Options.fromTransactionOptions(Options.priority(RpcPriority.HIGH)); + Options options3 = Options.fromTransactionOptions(); + Options options4 = Options.fromTransactionOptions(Options.priority(RpcPriority.LOW)); + + assertEquals(options1, options2); + assertNotEquals(options1, options3); + assertNotEquals(options1, options4); + assertNotEquals(options2, options3); + assertNotEquals(options2, options4); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java index a38cd71baa..4048d4b37c 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; import static org.mockito.Matchers.any; import static org.mockito.Matchers.anyMap; @@ -27,12 +28,14 @@ import com.google.api.gax.rpc.InternalException; import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.UnavailableException; +import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.base.Ticker; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import com.google.spanner.v1.*; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.RequestOptions.Priority; import io.grpc.Status.Code; import java.util.Collections; import java.util.Iterator; @@ -335,4 +338,21 @@ public void testExecuteStreamingPartitionedUpdateGenericInternalException() { Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); } } + + @Test + public void testRequestWithoutPriority() { + ExecuteSqlRequest request = + tx.newTransactionRequestFrom( + Statement.of("UPDATE FOO SET BAR=1 WHERE TRUE"), Options.fromUpdateOptions()); + assertEquals(Priority.PRIORITY_UNSPECIFIED, request.getRequestOptions().getPriority()); + } + + @Test + public void testRequestWithPriority() { + ExecuteSqlRequest request = + tx.newTransactionRequestFrom( + Statement.of("UPDATE FOO SET BAR=1 WHERE TRUE"), + Options.fromUpdateOptions(Options.priority(RpcPriority.LOW))); + assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority()); + } }