Skip to content

Commit 26c6c63

Browse files
arpan14arpan14
authored
feat: add support for Directed Read options (#2766)
* fix: prevent illegal negative timeout values into thread sleep() method while retrying exceptions in unit tests. * For details on issue see - #2206 * Fixing lint issues. * feat: add support for Directed Read options. * chore: fix lint issues. * test: add unit tests for options class. * test: add tests using mock spanner. * test: add unit test for partitioned read. * test: add unit test for partitioned read. * chore: adding option in spanner options. * chore: fix NPE. * chore: disabling test on emulator. * chore: adding test for query in RW transaction. * chore: adding IT for transaction manager interface. * chore: disable IT for emulator. * chore: PR comments. * chore: address PR comments.
1 parent 74a586f commit 26c6c63

12 files changed

+740
-14
lines changed

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

+19
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import com.google.common.util.concurrent.MoreExecutors;
4343
import com.google.protobuf.ByteString;
4444
import com.google.spanner.v1.BeginTransactionRequest;
45+
import com.google.spanner.v1.DirectedReadOptions;
4546
import com.google.spanner.v1.ExecuteBatchDmlRequest;
4647
import com.google.spanner.v1.ExecuteSqlRequest;
4748
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -72,6 +73,7 @@ abstract static class Builder<B extends Builder<?, T>, T extends AbstractReadCon
7273
private Span span = Tracing.getTracer().getCurrentSpan();
7374
private int defaultPrefetchChunks = SpannerOptions.Builder.DEFAULT_PREFETCH_CHUNKS;
7475
private QueryOptions defaultQueryOptions = SpannerOptions.Builder.DEFAULT_QUERY_OPTIONS;
76+
private DirectedReadOptions defaultDirectedReadOption;
7577
private ExecutorProvider executorProvider;
7678
private Clock clock = new Clock();
7779

@@ -117,6 +119,11 @@ B setClock(Clock clock) {
117119
return self();
118120
}
119121

122+
B setDefaultDirectedReadOptions(DirectedReadOptions directedReadOptions) {
123+
this.defaultDirectedReadOption = directedReadOptions;
124+
return self();
125+
}
126+
120127
abstract T build();
121128
}
122129

@@ -399,6 +406,7 @@ void initTransaction() {
399406
private final int defaultPrefetchChunks;
400407
private final QueryOptions defaultQueryOptions;
401408

409+
private final DirectedReadOptions defaultDirectedReadOptions;
402410
private final Clock clock;
403411

404412
@GuardedBy("lock")
@@ -423,6 +431,7 @@ void initTransaction() {
423431
this.rpc = builder.rpc;
424432
this.defaultPrefetchChunks = builder.defaultPrefetchChunks;
425433
this.defaultQueryOptions = builder.defaultQueryOptions;
434+
this.defaultDirectedReadOptions = builder.defaultDirectedReadOption;
426435
this.span = builder.span;
427436
this.executorProvider = builder.executorProvider;
428437
this.clock = builder.clock;
@@ -623,6 +632,11 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
623632
if (options.hasDataBoostEnabled()) {
624633
builder.setDataBoostEnabled(options.dataBoostEnabled());
625634
}
635+
if (options.hasDirectedReadOptions()) {
636+
builder.setDirectedReadOptions(options.directedReadOptions());
637+
} else if (defaultDirectedReadOptions != null) {
638+
builder.setDirectedReadOptions(defaultDirectedReadOptions);
639+
}
626640
builder.setSeqno(getSeqNo());
627641
builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
628642
builder.setRequestOptions(buildRequestOptions(options));
@@ -811,6 +825,11 @@ ResultSet readInternalWithOptions(
811825
if (readOptions.hasDataBoostEnabled()) {
812826
builder.setDataBoostEnabled(readOptions.dataBoostEnabled());
813827
}
828+
if (readOptions.hasDirectedReadOptions()) {
829+
builder.setDirectedReadOptions(readOptions.directedReadOptions());
830+
} else if (defaultDirectedReadOptions != null) {
831+
builder.setDirectedReadOptions(defaultDirectedReadOptions);
832+
}
814833
final int prefetchChunks =
815834
readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
816835
ResumableStreamIterator stream =

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

+6-2
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,9 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
6060
.setDefaultQueryOptions(
6161
sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId()))
6262
.setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider())
63-
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()),
63+
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks())
64+
.setDefaultDirectedReadOptions(
65+
sessionClient.getSpanner().getOptions().getDirectedReadOptions()),
6466
checkNotNull(bound));
6567
}
6668

@@ -77,7 +79,9 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
7779
.setDefaultQueryOptions(
7880
sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId()))
7981
.setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider())
80-
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()),
82+
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks())
83+
.setDefaultDirectedReadOptions(
84+
sessionClient.getSpanner().getOptions().getDirectedReadOptions()),
8185
batchTransactionId);
8286
}
8387

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

+45-1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.spanner;
1818

1919
import com.google.common.base.Preconditions;
20+
import com.google.spanner.v1.DirectedReadOptions;
2021
import com.google.spanner.v1.RequestOptions.Priority;
2122
import java.io.Serializable;
2223
import java.util.Objects;
@@ -224,6 +225,18 @@ public static CreateUpdateDeleteAdminApiOption validateOnly(Boolean validateOnly
224225
return new ValidateOnlyOption(validateOnly);
225226
}
226227

228+
/**
229+
* Option to request DirectedRead for ReadOnlyTransaction and SingleUseTransaction.
230+
*
231+
* <p>The DirectedReadOptions can be used to indicate which replicas or regions should be used for
232+
* non-transactional reads or queries. Not all requests can be sent to non-leader replicas. In
233+
* particular, some requests such as reads within read-write transactions must be sent to a
234+
* designated leader replica. These requests ignore DirectedReadOptions.
235+
*/
236+
public static ReadAndQueryOption directedRead(DirectedReadOptions directedReadOptions) {
237+
return new DirectedReadOption(directedReadOptions);
238+
}
239+
227240
/** Option to request {@link CommitStats} for read/write transactions. */
228241
static final class CommitStatsOption extends InternalOption implements TransactionOption {
229242
@Override
@@ -325,6 +338,21 @@ void appendToOptions(Options options) {
325338
}
326339
}
327340

341+
static final class DirectedReadOption extends InternalOption implements ReadAndQueryOption {
342+
private final DirectedReadOptions directedReadOptions;
343+
344+
DirectedReadOption(DirectedReadOptions directedReadOptions) {
345+
this.directedReadOptions =
346+
Preconditions.checkNotNull(directedReadOptions, "DirectedReadOptions cannot be null");
347+
;
348+
}
349+
350+
@Override
351+
void appendToOptions(Options options) {
352+
options.directedReadOptions = directedReadOptions;
353+
}
354+
}
355+
328356
private boolean withCommitStats;
329357
private Long limit;
330358
private Integer prefetchChunks;
@@ -338,6 +366,7 @@ void appendToOptions(Options options) {
338366
private Boolean validateOnly;
339367
private Boolean withOptimisticLock;
340368
private Boolean dataBoostEnabled;
369+
private DirectedReadOptions directedReadOptions;
341370

342371
// Construction is via factory methods below.
343372
private Options() {}
@@ -438,6 +467,14 @@ Boolean dataBoostEnabled() {
438467
return dataBoostEnabled;
439468
}
440469

470+
boolean hasDirectedReadOptions() {
471+
return directedReadOptions != null;
472+
}
473+
474+
DirectedReadOptions directedReadOptions() {
475+
return directedReadOptions;
476+
}
477+
441478
@Override
442479
public String toString() {
443480
StringBuilder b = new StringBuilder();
@@ -477,6 +514,9 @@ public String toString() {
477514
if (dataBoostEnabled != null) {
478515
b.append("dataBoostEnabled: ").append(dataBoostEnabled).append(' ');
479516
}
517+
if (directedReadOptions != null) {
518+
b.append("directedReadOptions: ").append(directedReadOptions).append(' ');
519+
}
480520
return b.toString();
481521
}
482522

@@ -512,7 +552,8 @@ public boolean equals(Object o) {
512552
&& Objects.equals(etag(), that.etag())
513553
&& Objects.equals(validateOnly(), that.validateOnly())
514554
&& Objects.equals(withOptimisticLock(), that.withOptimisticLock())
515-
&& Objects.equals(dataBoostEnabled(), that.dataBoostEnabled());
555+
&& Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
556+
&& Objects.equals(directedReadOptions(), that.directedReadOptions());
516557
}
517558

518559
@Override
@@ -557,6 +598,9 @@ public int hashCode() {
557598
if (dataBoostEnabled != null) {
558599
result = 31 * result + dataBoostEnabled.hashCode();
559600
}
601+
if (directedReadOptions != null) {
602+
result = 31 * result + directedReadOptions.hashCode();
603+
}
560604
return result;
561605
}
562606

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

+3
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ public ReadContext singleUse(TimestampBound bound) {
255255
.setRpc(spanner.getRpc())
256256
.setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId))
257257
.setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks())
258+
.setDefaultDirectedReadOptions(spanner.getOptions().getDirectedReadOptions())
258259
.setSpan(currentSpan)
259260
.setExecutorProvider(spanner.getAsyncExecutorProvider())
260261
.build());
@@ -274,6 +275,7 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) {
274275
.setRpc(spanner.getRpc())
275276
.setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId))
276277
.setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks())
278+
.setDefaultDirectedReadOptions(spanner.getOptions().getDirectedReadOptions())
277279
.setSpan(currentSpan)
278280
.setExecutorProvider(spanner.getAsyncExecutorProvider())
279281
.buildSingleUseReadOnlyTransaction());
@@ -293,6 +295,7 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) {
293295
.setRpc(spanner.getRpc())
294296
.setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId))
295297
.setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks())
298+
.setDefaultDirectedReadOptions(spanner.getOptions().getDirectedReadOptions())
296299
.setSpan(currentSpan)
297300
.setExecutorProvider(spanner.getAsyncExecutorProvider())
298301
.build());

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

+36
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import com.google.cloud.TransportOptions;
3434
import com.google.cloud.grpc.GcpManagedChannelOptions;
3535
import com.google.cloud.grpc.GrpcTransportOptions;
36+
import com.google.cloud.spanner.Options.DirectedReadOption;
3637
import com.google.cloud.spanner.Options.QueryOption;
3738
import com.google.cloud.spanner.Options.UpdateOption;
3839
import com.google.cloud.spanner.admin.database.v1.DatabaseAdminSettings;
@@ -50,6 +51,7 @@
5051
import com.google.common.collect.ImmutableMap;
5152
import com.google.common.collect.ImmutableSet;
5253
import com.google.common.util.concurrent.ThreadFactoryBuilder;
54+
import com.google.spanner.v1.DirectedReadOptions;
5355
import com.google.spanner.v1.ExecuteSqlRequest;
5456
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
5557
import com.google.spanner.v1.SpannerGrpc;
@@ -137,6 +139,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
137139
private final String compressorName;
138140
private final boolean leaderAwareRoutingEnabled;
139141
private final boolean attemptDirectPath;
142+
private final DirectedReadOptions directedReadOptions;
140143

141144
/** Interface that can be used to provide {@link CallCredentials} to {@link SpannerOptions}. */
142145
public interface CallCredentialsProvider {
@@ -627,6 +630,7 @@ private SpannerOptions(Builder builder) {
627630
compressorName = builder.compressorName;
628631
leaderAwareRoutingEnabled = builder.leaderAwareRoutingEnabled;
629632
attemptDirectPath = builder.attemptDirectPath;
633+
directedReadOptions = builder.directedReadOptions;
630634
}
631635

632636
/**
@@ -729,6 +733,7 @@ public static class Builder
729733
private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
730734
private boolean leaderAwareRoutingEnabled = true;
731735
private boolean attemptDirectPath = true;
736+
private DirectedReadOptions directedReadOptions;
732737

733738
private static String createCustomClientLibToken(String token) {
734739
return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -789,6 +794,7 @@ private Builder() {
789794
this.channelConfigurator = options.channelConfigurator;
790795
this.interceptorProvider = options.interceptorProvider;
791796
this.attemptDirectPath = options.attemptDirectPath;
797+
this.directedReadOptions = options.directedReadOptions;
792798
}
793799

794800
@Override
@@ -1153,6 +1159,32 @@ public Builder setAsyncExecutorProvider(CloseableExecutorProvider provider) {
11531159
return this;
11541160
}
11551161

1162+
/**
1163+
* Sets the {@link DirectedReadOption} that specify which replicas or regions should be used for
1164+
* non-transactional reads or queries.
1165+
*
1166+
* <p>DirectedReadOptions set at the request level will take precedence over the options set
1167+
* using this method.
1168+
*
1169+
* <p>An example below of how {@link DirectedReadOptions} can be constructed by including a
1170+
* replica.
1171+
*
1172+
* <pre><code>
1173+
* DirectedReadOptions.newBuilder()
1174+
* .setIncludeReplicas(
1175+
* IncludeReplicas.newBuilder()
1176+
* .addReplicaSelections(
1177+
* ReplicaSelection.newBuilder().setLocation("us-east1").build()))
1178+
* .build();
1179+
* }
1180+
* </code></pre>
1181+
*/
1182+
public Builder setDirectedReadOptions(DirectedReadOptions directedReadOptions) {
1183+
this.directedReadOptions =
1184+
Preconditions.checkNotNull(directedReadOptions, "DirectedReadOptions cannot be null");
1185+
return this;
1186+
}
1187+
11561188
/**
11571189
* Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
11581190
* PartialResultSet} chunks for each read and query. The data size of each chunk depends on the
@@ -1371,6 +1403,10 @@ public boolean isLeaderAwareRoutingEnabled() {
13711403
return leaderAwareRoutingEnabled;
13721404
}
13731405

1406+
public DirectedReadOptions getDirectedReadOptions() {
1407+
return directedReadOptions;
1408+
}
1409+
13741410
@BetaApi
13751411
public boolean isAttemptDirectPath() {
13761412
return attemptDirectPath;

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

+22
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@
2525
import com.google.api.gax.core.ExecutorProvider;
2626
import com.google.cloud.spanner.Options.RpcPriority;
2727
import com.google.cloud.spanner.spi.v1.SpannerRpc;
28+
import com.google.spanner.v1.DirectedReadOptions;
29+
import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
30+
import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
2831
import com.google.spanner.v1.ExecuteBatchDmlRequest;
2932
import com.google.spanner.v1.ExecuteSqlRequest;
3033
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -45,6 +48,14 @@
4548

4649
@RunWith(Parameterized.class)
4750
public class AbstractReadContextTest {
51+
private static final DirectedReadOptions DIRECTED_READ_OPTIONS =
52+
DirectedReadOptions.newBuilder()
53+
.setIncludeReplicas(
54+
IncludeReplicas.newBuilder()
55+
.addReplicaSelections(
56+
ReplicaSelection.newBuilder().setLocation("us-west1").build()))
57+
.build();
58+
4859
@Parameter(0)
4960
public QueryOptions defaultQueryOptions;
5061

@@ -250,4 +261,15 @@ public void executeSqlRequestBuilderWithRequestOptionsWithTxnTag() {
250261
.isEqualTo("app=spanner,env=test,action=query");
251262
assertThat(request.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test");
252263
}
264+
265+
@Test
266+
public void testGetExecuteSqlRequestBuilderWithDirectedReadOptions() {
267+
ExecuteSqlRequest.Builder request =
268+
context.getExecuteSqlRequestBuilder(
269+
Statement.of("SELECT * FROM FOO"),
270+
QueryMode.NORMAL,
271+
Options.fromQueryOptions(Options.directedRead(DIRECTED_READ_OPTIONS)),
272+
false);
273+
assertEquals(DIRECTED_READ_OPTIONS, request.getDirectedReadOptions());
274+
}
253275
}

0 commit comments

Comments
 (0)