Skip to content

Commit f39e4a3

Browse files
gauravpurohit06gauravpurohit06
authored
feat: Adding support for databoost enabled in PartitionedRead and PartitionedQuery (#2316)
* feat: Adding support for spanner serverless analytics * feat: Adding BetaAPI annotation to flag out changes are under development * test: Adding Integration test * style: formatting * feat: Using databoost field instead of serverless analytics * feat: Integration test and sample. * refactor: method name change * refactor: minor refactoring * refactor: Adding more junit and java docs
1 parent 6159d7e commit f39e4a3

File tree

5 files changed

+143
-8
lines changed

5 files changed

+143
-8
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,9 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
595595
builder.setTransaction(selector);
596596
}
597597
}
598+
if (options.hasDataBoostEnabled()) {
599+
builder.setDataBoostEnabled(options.dataBoostEnabled());
600+
}
598601
builder.setSeqno(getSeqNo());
599602
builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
600603
builder.setRequestOptions(buildRequestOptions(options));
@@ -773,6 +776,9 @@ ResultSet readInternalWithOptions(
773776
if (partitionToken != null) {
774777
builder.setPartitionToken(partitionToken);
775778
}
779+
if (readOptions.hasDataBoostEnabled()) {
780+
builder.setDataBoostEnabled(readOptions.dataBoostEnabled());
781+
}
776782
final int prefetchChunks =
777783
readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
778784
ResumableStreamIterator stream =

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

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.spanner;
1818

19+
import com.google.api.core.BetaApi;
1920
import com.google.common.base.Preconditions;
2021
import com.google.spanner.v1.RequestOptions.Priority;
2122
import java.io.Serializable;
@@ -154,6 +155,16 @@ public static ListOption pageSize(int pageSize) {
154155
return new PageSizeOption(pageSize);
155156
}
156157

158+
/**
159+
* If this is for a partitioned read & query and this field is set to `true`, the request will be
160+
* executed via Spanner independent compute resources. The method is available in Beta mode (and
161+
* is not generally available now).
162+
*/
163+
@BetaApi
164+
public static DataBoostQueryOption dataBoostEnabled(Boolean dataBoostEnabled) {
165+
return new DataBoostQueryOption(dataBoostEnabled);
166+
}
167+
157168
/**
158169
* Specifying this will cause the list operation to start fetching the record from this onwards.
159170
*/
@@ -329,6 +340,7 @@ void appendToOptions(Options options) {
329340
private String etag;
330341
private Boolean validateOnly;
331342
private Boolean withOptimisticLock;
343+
private Boolean dataBoostEnabled;
332344

333345
// Construction is via factory methods below.
334346
private Options() {}
@@ -421,6 +433,14 @@ Boolean withOptimisticLock() {
421433
return withOptimisticLock;
422434
}
423435

436+
boolean hasDataBoostEnabled() {
437+
return dataBoostEnabled != null;
438+
}
439+
440+
Boolean dataBoostEnabled() {
441+
return dataBoostEnabled;
442+
}
443+
424444
@Override
425445
public String toString() {
426446
StringBuilder b = new StringBuilder();
@@ -457,6 +477,9 @@ public String toString() {
457477
if (withOptimisticLock != null) {
458478
b.append("withOptimisticLock: ").append(withOptimisticLock).append(' ');
459479
}
480+
if (dataBoostEnabled != null) {
481+
b.append("dataBoostEnabled: ").append(dataBoostEnabled).append(' ');
482+
}
460483
return b.toString();
461484
}
462485

@@ -491,7 +514,8 @@ public boolean equals(Object o) {
491514
&& Objects.equals(tag(), that.tag())
492515
&& Objects.equals(etag(), that.etag())
493516
&& Objects.equals(validateOnly(), that.validateOnly())
494-
&& Objects.equals(withOptimisticLock(), that.withOptimisticLock());
517+
&& Objects.equals(withOptimisticLock(), that.withOptimisticLock())
518+
&& Objects.equals(dataBoostEnabled(), that.dataBoostEnabled());
495519
}
496520

497521
@Override
@@ -533,6 +557,9 @@ public int hashCode() {
533557
if (withOptimisticLock != null) {
534558
result = 31 * result + withOptimisticLock.hashCode();
535559
}
560+
if (dataBoostEnabled != null) {
561+
result = 31 * result + dataBoostEnabled.hashCode();
562+
}
536563
return result;
537564
}
538565

@@ -613,6 +640,20 @@ void appendToOptions(Options options) {
613640
}
614641
}
615642

643+
static final class DataBoostQueryOption extends InternalOption implements ReadAndQueryOption {
644+
645+
private final Boolean dataBoostEnabled;
646+
647+
DataBoostQueryOption(Boolean dataBoostEnabled) {
648+
this.dataBoostEnabled = dataBoostEnabled;
649+
}
650+
651+
@Override
652+
void appendToOptions(Options options) {
653+
options.dataBoostEnabled = dataBoostEnabled;
654+
}
655+
}
656+
616657
static class PageSizeOption extends InternalOption implements ListOption {
617658
private final int pageSize;
618659

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.truth.Truth.assertThat;
2020
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertTrue;
2122
import static org.mockito.Mockito.mock;
2223
import static org.mockito.Mockito.when;
2324

@@ -188,6 +189,17 @@ public void testGetExecuteSqlRequestBuilderWithPriority() {
188189
assertEquals(Priority.PRIORITY_MEDIUM, request.getRequestOptions().getPriority());
189190
}
190191

192+
@Test
193+
public void testGetExecuteSqlRequestBuilderWithDataBoost() {
194+
ExecuteSqlRequest.Builder request =
195+
context.getExecuteSqlRequestBuilder(
196+
Statement.of("SELECT * FROM FOO"),
197+
QueryMode.NORMAL,
198+
Options.fromQueryOptions(Options.dataBoostEnabled(true)),
199+
false);
200+
assertTrue(request.getDataBoostEnabled());
201+
}
202+
191203
@Test
192204
public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
193205
ExecuteBatchDmlRequest.Builder request =

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

Lines changed: 47 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,15 @@ public void zeroPrefetchChunksNotAllowed() {
6363

6464
@Test
6565
public void allOptionsPresent() {
66-
Options options = Options.fromReadOptions(Options.limit(10), Options.prefetchChunks(1));
66+
Options options =
67+
Options.fromReadOptions(
68+
Options.limit(10), Options.prefetchChunks(1), Options.dataBoostEnabled(true));
6769
assertThat(options.hasLimit()).isTrue();
6870
assertThat(options.limit()).isEqualTo(10);
6971
assertThat(options.hasPrefetchChunks()).isTrue();
7072
assertThat(options.prefetchChunks()).isEqualTo(1);
73+
assertThat(options.hasDataBoostEnabled()).isTrue();
74+
assertTrue(options.dataBoostEnabled());
7175
}
7276

7377
@Test
@@ -79,6 +83,7 @@ public void allOptionsAbsent() {
7983
assertThat(options.hasPageToken()).isFalse();
8084
assertThat(options.hasPriority()).isFalse();
8185
assertThat(options.hasTag()).isFalse();
86+
assertThat(options.hasDataBoostEnabled()).isFalse();
8287
assertThat(options.toString()).isEqualTo("");
8388
assertThat(options.equals(options)).isTrue();
8489
assertThat(options.equals(null)).isFalse();
@@ -153,11 +158,17 @@ public void listEquality() {
153158
public void readOptionsTest() {
154159
int limit = 3;
155160
String tag = "app=spanner,env=test,action=read";
156-
Options options = Options.fromReadOptions(Options.limit(limit), Options.tag(tag));
161+
boolean dataBoost = true;
162+
Options options =
163+
Options.fromReadOptions(
164+
Options.limit(limit), Options.tag(tag), Options.dataBoostEnabled(true));
157165

158-
assertThat(options.toString()).isEqualTo("limit: " + limit + " " + "tag: " + tag + " ");
166+
assertThat(options.toString())
167+
.isEqualTo(
168+
"limit: " + limit + " " + "tag: " + tag + " " + "dataBoostEnabled: " + dataBoost + " ");
159169
assertThat(options.tag()).isEqualTo(tag);
160-
assertThat(options.hashCode()).isEqualTo(-1111478426);
170+
assertEquals(dataBoost, options.dataBoostEnabled());
171+
assertThat(options.hashCode()).isEqualTo(-96091607);
161172
}
162173

163174
@Test
@@ -185,12 +196,41 @@ public void readEquality() {
185196
public void queryOptionsTest() {
186197
int chunks = 3;
187198
String tag = "app=spanner,env=test,action=query";
188-
Options options = Options.fromQueryOptions(Options.prefetchChunks(chunks), Options.tag(tag));
199+
boolean dataBoost = true;
200+
Options options =
201+
Options.fromQueryOptions(
202+
Options.prefetchChunks(chunks), Options.tag(tag), Options.dataBoostEnabled(true));
189203
assertThat(options.toString())
190-
.isEqualTo("prefetchChunks: " + chunks + " " + "tag: " + tag + " ");
204+
.isEqualTo(
205+
"prefetchChunks: "
206+
+ chunks
207+
+ " "
208+
+ "tag: "
209+
+ tag
210+
+ " "
211+
+ "dataBoostEnabled: "
212+
+ dataBoost
213+
+ " ");
191214
assertThat(options.prefetchChunks()).isEqualTo(chunks);
192215
assertThat(options.tag()).isEqualTo(tag);
193-
assertThat(options.hashCode()).isEqualTo(-97431824);
216+
assertEquals(dataBoost, options.dataBoostEnabled());
217+
assertThat(options.hashCode()).isEqualTo(1274581983);
218+
}
219+
220+
@Test
221+
public void testReadOptionsDataBoost() {
222+
boolean dataBoost = true;
223+
Options options = Options.fromReadOptions(Options.dataBoostEnabled(true));
224+
assertTrue(options.hasDataBoostEnabled());
225+
assertEquals("dataBoostEnabled: " + dataBoost + " ", options.toString());
226+
}
227+
228+
@Test
229+
public void testQueryOptionsDataBoost() {
230+
boolean dataBoost = true;
231+
Options options = Options.fromQueryOptions(Options.dataBoostEnabled(true));
232+
assertTrue(options.hasDataBoostEnabled());
233+
assertEquals("dataBoostEnabled: " + dataBoost + " ", options.toString());
194234
}
195235

196236
@Test

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchReadTest.java

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.google.cloud.spanner.IntegrationTestEnv;
3232
import com.google.cloud.spanner.KeySet;
3333
import com.google.cloud.spanner.Mutation;
34+
import com.google.cloud.spanner.Options;
3435
import com.google.cloud.spanner.ParallelIntegrationTest;
3536
import com.google.cloud.spanner.Partition;
3637
import com.google.cloud.spanner.PartitionOptions;
@@ -238,6 +239,25 @@ public void readUsingIndex() {
238239
assertThat(numRowsRead).isEqualTo(numRows);
239240
}
240241

242+
@Test
243+
public void dataBoostRead() {
244+
assumeFalse("Emulator does not support data boost read", isUsingEmulator());
245+
246+
BitSet seenRows = new BitSet(numRows);
247+
TimestampBound bound = getRandomBound();
248+
PartitionOptions partitionParams = getRandomPartitionOptions();
249+
batchTxn = getBatchClient().batchReadOnlyTransaction(bound);
250+
List<Partition> partitions =
251+
batchTxn.partitionRead(
252+
partitionParams,
253+
TABLE_NAME,
254+
KeySet.all(),
255+
Arrays.asList("Key", "Data", "Fingerprint", "Size"),
256+
Options.dataBoostEnabled(true));
257+
BatchTransactionId txnID = batchTxn.getBatchTransactionId();
258+
fetchAndValidateRows(partitions, txnID, seenRows);
259+
}
260+
241261
@After
242262
public void tearDown() {
243263
if (batchTxn != null) {
@@ -273,6 +293,22 @@ private PartitionOptions getRandomPartitionOptions() {
273293
return parameters;
274294
}
275295

296+
@Test
297+
public void dataBoostQuery() {
298+
assumeFalse("Emulator does not support data boost query", isUsingEmulator());
299+
BitSet seenRows = new BitSet(numRows);
300+
TimestampBound bound = getRandomBound();
301+
PartitionOptions partitionParams = getRandomPartitionOptions();
302+
batchTxn = getBatchClient().batchReadOnlyTransaction(bound);
303+
List<Partition> partitions =
304+
batchTxn.partitionQuery(
305+
partitionParams,
306+
Statement.of("SELECT Key, Data, Fingerprint, Size FROM " + TABLE_NAME),
307+
Options.dataBoostEnabled(true));
308+
BatchTransactionId txnID = batchTxn.getBatchTransactionId();
309+
fetchAndValidateRows(partitions, txnID, seenRows);
310+
}
311+
276312
private TimestampBound getRandomBound() {
277313
Date date = new Date();
278314
switch (RANDOM.nextInt(3)) {

0 commit comments

Comments
 (0)