feat: Adding support for databoost enabled in PartitionedRead and PartitionedQuery

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

@@ -595,6 +595,9 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
         builder.setTransaction(selector);
       }
     }
+    if (options.hasDataBoostEnabled()) {
+      builder.setDataBoostEnabled(options.dataBoostEnabled());
+    }
     builder.setSeqno(getSeqNo());
     builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
     builder.setRequestOptions(buildRequestOptions(options));
@@ -773,6 +776,9 @@ ResultSet readInternalWithOptions(
     if (partitionToken != null) {
       builder.setPartitionToken(partitionToken);
     }
+    if (readOptions.hasDataBoostEnabled()) {
+      builder.setDataBoostEnabled(readOptions.dataBoostEnabled());
+    }
     final int prefetchChunks =
         readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
     ResumableStreamIterator stream =

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner;
 
+import com.google.api.core.BetaApi;
 import com.google.common.base.Preconditions;
 import com.google.spanner.v1.RequestOptions.Priority;
 import java.io.Serializable;
@@ -146,6 +147,15 @@ public static ListOption pageSize(int pageSize) {
     return new PageSizeOption(pageSize);
   }
 
+  /**
+   * If this is for a partitioned read & query and this field is set to `true`, the request will be
+   * executed via Spanner independent compute resources.
+   */
+  @BetaApi
+  public static DataBoostQueryOption dataBoostEnabled(Boolean dataBoostEnabled) {
+    return new DataBoostQueryOption(dataBoostEnabled);
+  }
+
   /**
    * Specifying this will cause the list operation to start fetching the record from this onwards.
    */
@@ -321,6 +331,7 @@ void appendToOptions(Options options) {
   private String etag;
   private Boolean validateOnly;
   private Boolean withOptimisticLock;
+  private Boolean dataBoostEnabled;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -413,6 +424,14 @@ Boolean withOptimisticLock() {
     return withOptimisticLock;
   }
 
+  boolean hasDataBoostEnabled() {
+    return dataBoostEnabled != null;
+  }
+
+  Boolean dataBoostEnabled() {
+    return dataBoostEnabled;
+  }
+
   @Override
   public String toString() {
     StringBuilder b = new StringBuilder();
@@ -440,6 +459,9 @@ public String toString() {
     if (tag != null) {
       b.append("tag: ").append(tag).append(' ');
     }
+    if (dataBoostEnabled != null) {
+      b.append("dataBoostEnabled: ").append(dataBoostEnabled).append(' ');
+    }
     if (etag != null) {
       b.append("etag: ").append(etag).append(' ');
     }
@@ -482,6 +504,7 @@ public boolean equals(Object o) {
         && Objects.equals(priority(), that.priority())
         && Objects.equals(tag(), that.tag())
         && Objects.equals(etag(), that.etag())
+        && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
         && Objects.equals(validateOnly(), that.validateOnly())
         && Objects.equals(withOptimisticLock(), that.withOptimisticLock());
   }
@@ -525,6 +548,9 @@ public int hashCode() {
     if (withOptimisticLock != null) {
       result = 31 * result + withOptimisticLock.hashCode();
     }
+    if (dataBoostEnabled != null) {
+      result = 31 * result + dataBoostEnabled.hashCode();
+    }
     return result;
   }
 
@@ -605,6 +631,20 @@ void appendToOptions(Options options) {
     }
   }
 
+  static final class DataBoostQueryOption extends InternalOption implements ReadAndQueryOption {
+
+    private final Boolean dataBoostEnabled;
+
+    DataBoostQueryOption(Boolean dataBoostEnabled) {
+      this.dataBoostEnabled = dataBoostEnabled;
+    }
+
+    @Override
+    void appendToOptions(Options options) {
+      options.dataBoostEnabled = dataBoostEnabled;
+    }
+  }
+
   static class PageSizeOption extends InternalOption implements ListOption {
     private final int pageSize;
 

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

@@ -31,6 +31,7 @@
 import com.google.cloud.spanner.IntegrationTestEnv;
 import com.google.cloud.spanner.KeySet;
 import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.ParallelIntegrationTest;
 import com.google.cloud.spanner.Partition;
 import com.google.cloud.spanner.PartitionOptions;
@@ -238,6 +239,27 @@ public void readUsingIndex() {
     assertThat(numRowsRead).isEqualTo(numRows);
   }
 
+  @Test
+  public void dataBoostRead() {
+    assumeFalse(
+        "PostgreSQL does not support the PartitionRead RPC", dialect.dialect == Dialect.POSTGRESQL);
+    assumeFalse("Emulator does not support data boost read", isUsingEmulator());
+
+    BitSet seenRows = new BitSet(numRows);
+    TimestampBound bound = getRandomBound();
+    PartitionOptions partitionParams = getRandomPartitionOptions();
+    batchTxn = getBatchClient().batchReadOnlyTransaction(bound);
+    List<Partition> partitions =
+        batchTxn.partitionRead(
+            partitionParams,
+            TABLE_NAME,
+            KeySet.all(),
+            Arrays.asList("Key", "Data", "Fingerprint", "Size"),
+            Options.dataBoostEnabled(true));
+    BatchTransactionId txnID = batchTxn.getBatchTransactionId();
+    fetchAndValidateRows(partitions, txnID, seenRows);
+  }
+
   @After
   public void tearDown() {
     if (batchTxn != null) {
@@ -273,6 +295,22 @@ private PartitionOptions getRandomPartitionOptions() {
     return parameters;
   }
 
+  @Test
+  public void dataBoostQuery() {
+    assumeFalse("Emulator does not support data boost query", isUsingEmulator());
+    BitSet seenRows = new BitSet(numRows);
+    TimestampBound bound = getRandomBound();
+    PartitionOptions partitionParams = getRandomPartitionOptions();
+    batchTxn = getBatchClient().batchReadOnlyTransaction(bound);
+    List<Partition> partitions =
+        batchTxn.partitionQuery(
+            partitionParams,
+            Statement.of("SELECT Key, Data, Fingerprint, Size FROM " + TABLE_NAME),
+            Options.dataBoostEnabled(true));
+    BatchTransactionId txnID = batchTxn.getBatchTransactionId();
+    fetchAndValidateRows(partitions, txnID, seenRows);
+  }
+
   private TimestampBound getRandomBound() {
     Date date = new Date();
     switch (RANDOM.nextInt(3)) {