feat: support Directed Read in Connection API (#2855)

* feat: support Directed Read in Connection API

Adds support for setting Directed Read options in the Connection API.
The value must be a valid JSON representation of a DirectedReadOptions
protobuf instance.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* chore: address review comments

* fix: verify that DirectedRead is not used for DML

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: olavloite

google-cloud-spanner/clirr-ignored-differences.xml

@@ -540,4 +540,17 @@
     <className>com/google/cloud/spanner/Dialect</className>
     <method>java.lang.String getDefaultSchema()</method>
   </difference>
+  
+  <!-- Added DirectedReadOptions -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>com.google.spanner.v1.DirectedReadOptions getDirectedRead()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void setDirectedRead(com.google.spanner.v1.DirectedReadOptions)</method>
+  </difference>
+  
 </differences>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java

@@ -18,6 +18,7 @@
 
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TimestampBound.Mode;
@@ -27,6 +28,7 @@
 import com.google.common.base.Preconditions;
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.Durations;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.RequestOptions.Priority;
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -306,6 +308,48 @@ public TimestampBound convert(String value) {
       return null;
     }
   }
+  /**
+   * Converter from string to possible values for {@link com.google.spanner.v1.DirectedReadOptions}.
+   */
+  static class DirectedReadOptionsConverter
+      implements ClientSideStatementValueConverter<DirectedReadOptions> {
+    private final Pattern allowedValues;
+
+    public DirectedReadOptionsConverter(String allowedValues) {
+      // Remove the single quotes at the beginning and end.
+      this.allowedValues =
+          Pattern.compile(
+              "(?is)\\A" + allowedValues.substring(1, allowedValues.length() - 1) + "\\z");
+    }
+
+    @Override
+    public Class<DirectedReadOptions> getParameterClass() {
+      return DirectedReadOptions.class;
+    }
+
+    @Override
+    public DirectedReadOptions convert(String value) {
+      Matcher matcher = allowedValues.matcher(value);
+      if (matcher.find()) {
+        try {
+          return DirectedReadOptionsUtil.parse(value);
+        } catch (SpannerException spannerException) {
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.INVALID_ARGUMENT,
+              String.format(
+                  "Failed to parse '%s' as a valid value for DIRECTED_READ.\n"
+                      + "The value should be a JSON string like this: '%s'.\n"
+                      + "You can generate a valid JSON string from a DirectedReadOptions instance by calling %s.%s",
+                  value,
+                  "{\"includeReplicas\":{\"replicaSelections\":[{\"location\":\"eu-west1\",\"type\":\"READ_ONLY\"}]}}",
+                  DirectedReadOptionsUtil.class.getName(),
+                  "toString(DirectedReadOptions directedReadOptions)"),
+              spannerException);
+        }
+      }
+      return null;
+    }
+  }
 
   /** Converter for converting strings to {@link AutocommitDmlMode} values. */
   static class AutocommitDmlModeConverter

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java

@@ -38,6 +38,7 @@
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ResultSetStats;
 import java.util.Iterator;
@@ -489,6 +490,22 @@ default String getStatementTag() {
    */
   TimestampBound getReadOnlyStaleness();
 
+  /**
+   * Sets the {@link DirectedReadOptions} to use for both single-use and multi-use read-only
+   * transactions on this connection.
+   */
+  default void setDirectedRead(DirectedReadOptions directedReadOptions) {
+    throw new UnsupportedOperationException("Unimplemented");
+  }
+
+  /**
+   * Returns the {@link DirectedReadOptions} that are used for both single-use and multi-use
+   * read-only transactions on this connection.
+   */
+  default DirectedReadOptions getDirectedRead() {
+    throw new UnsupportedOperationException("Unimplemented");
+  }
+
   /**
    * Sets the query optimizer version to use for this connection.
    *

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -53,6 +53,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.ResultSetStats;
 import java.util.ArrayList;
@@ -236,6 +237,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
    */
   private int maxPartitionedParallelism;
 
+  private DirectedReadOptions directedReadOptions = null;
   private QueryOptions queryOptions = QueryOptions.getDefaultInstance();
   private RpcPriority rpcPriority = null;
   private SavepointSupport savepointSupport = SavepointSupport.FAIL_AFTER_ROLLBACK;
@@ -510,6 +512,21 @@ public TimestampBound getReadOnlyStaleness() {
     return this.readOnlyStaleness;
   }
 
+  @Override
+  public void setDirectedRead(DirectedReadOptions directedReadOptions) {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !isTransactionStarted(),
+        "Cannot set directed read options when a transaction has been started");
+    this.directedReadOptions = directedReadOptions;
+  }
+
+  @Override
+  public DirectedReadOptions getDirectedRead() {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    return this.directedReadOptions;
+  }
+
   @Override
   public void setOptimizerVersion(String optimizerVersion) {
     Preconditions.checkNotNull(optimizerVersion);
@@ -1131,7 +1148,8 @@ public ResultSet partitionQuery(
             CallType.SYNC,
             parsedStatement,
             getEffectivePartitionOptions(partitionOptions),
-            mergeDataBoost(mergeQueryRequestOptions(mergeQueryStatementTag(options)))));
+            mergeDataBoost(
+                mergeQueryRequestOptions(parsedStatement, mergeQueryStatementTag(options)))));
   }
 
   private PartitionOptions getEffectivePartitionOptions(
@@ -1427,41 +1445,38 @@ private List<ParsedStatement> parseUpdateStatements(Iterable<Statement> updates)
 
   private QueryOption[] mergeDataBoost(QueryOption... options) {
     if (this.dataBoostEnabled) {
-
-      // Shortcut for the most common scenario.
-      if (options == null || options.length == 0) {
-        options = new QueryOption[] {Options.dataBoostEnabled(true)};
-      } else {
-        options = Arrays.copyOf(options, options.length + 1);
-        options[options.length - 1] = Options.dataBoostEnabled(true);
-      }
+      options = appendQueryOption(options, Options.dataBoostEnabled(true));
     }
     return options;
   }
 
   private QueryOption[] mergeQueryStatementTag(QueryOption... options) {
     if (this.statementTag != null) {
-      // Shortcut for the most common scenario.
-      if (options == null || options.length == 0) {
-        options = new QueryOption[] {Options.tag(statementTag)};
-      } else {
-        options = Arrays.copyOf(options, options.length + 1);
-        options[options.length - 1] = Options.tag(statementTag);
-      }
+      options = appendQueryOption(options, Options.tag(statementTag));
       this.statementTag = null;
     }
     return options;
   }
 
-  private QueryOption[] mergeQueryRequestOptions(QueryOption... options) {
+  private QueryOption[] mergeQueryRequestOptions(
+      ParsedStatement parsedStatement, QueryOption... options) {
     if (this.rpcPriority != null) {
-      // Shortcut for the most common scenario.
-      if (options == null || options.length == 0) {
-        options = new QueryOption[] {Options.priority(this.rpcPriority)};
-      } else {
-        options = Arrays.copyOf(options, options.length + 1);
-        options[options.length - 1] = Options.priority(this.rpcPriority);
-      }
+      options = appendQueryOption(options, Options.priority(this.rpcPriority));
+    }
+    if (this.directedReadOptions != null
+        && currentUnitOfWork != null
+        && currentUnitOfWork.supportsDirectedReads(parsedStatement)) {
+      options = appendQueryOption(options, Options.directedRead(this.directedReadOptions));
+    }
+    return options;
+  }
+
+  private QueryOption[] appendQueryOption(QueryOption[] options, QueryOption append) {
+    if (options == null || options.length == 0) {
+      options = new QueryOption[] {append};
+    } else {
+      options = Arrays.copyOf(options, options.length + 1);
+      options[options.length - 1] = append;
     }
     return options;
   }
@@ -1516,7 +1531,7 @@ private ResultSet internalExecuteQuery(
             callType,
             statement,
             analyzeMode,
-            mergeQueryRequestOptions(mergeQueryStatementTag(options))));
+            mergeQueryRequestOptions(statement, mergeQueryStatementTag(options))));
   }
 
   private AsyncResultSet internalExecuteQueryAsync(
@@ -1538,7 +1553,7 @@ private AsyncResultSet internalExecuteQueryAsync(
             callType,
             statement,
             analyzeMode,
-            mergeQueryRequestOptions(mergeQueryStatementTag(options))),
+            mergeQueryRequestOptions(statement, mergeQueryStatementTag(options))),
         spanner.getAsyncExecutorProvider(),
         options);
   }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java

@@ -20,6 +20,7 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.PgTransactionMode.IsolationLevel;
 import com.google.protobuf.Duration;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.RequestOptions.Priority;
 
 /**
@@ -65,6 +66,10 @@ interface ConnectionStatementExecutor {
 
   StatementResult statementShowReadOnlyStaleness();
 
+  StatementResult statementSetDirectedRead(DirectedReadOptions directedReadOptions);
+
+  StatementResult statementShowDirectedRead();
+
   StatementResult statementSetOptimizerVersion(String optimizerVersion);
 
   StatementResult statementShowOptimizerVersion();

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java

@@ -28,6 +28,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DATA_BOOST_ENABLED;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DEFAULT_TRANSACTION_ISOLATION;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DIRECTED_READ;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_MAX_PARTITIONED_PARALLELISM;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_MAX_PARTITIONS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_OPTIMIZER_STATISTICS_PACKAGE;
@@ -49,6 +50,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_COMMIT_TIMESTAMP;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DATA_BOOST_ENABLED;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DIRECTED_READ;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_MAX_PARTITIONED_PARALLELISM;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_MAX_PARTITIONS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_OPTIMIZER_STATISTICS_PACKAGE;
@@ -91,6 +93,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Duration;
+import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.PlanNode;
 import com.google.spanner.v1.QueryPlan;
 import com.google.spanner.v1.RequestOptions;
@@ -283,6 +286,21 @@ public StatementResult statementShowReadOnlyStaleness() {
         SHOW_READ_ONLY_STALENESS);
   }
 
+  @Override
+  public StatementResult statementSetDirectedRead(DirectedReadOptions directedReadOptions) {
+    getConnection().setDirectedRead(directedReadOptions);
+    return noResult(SET_DIRECTED_READ);
+  }
+
+  @Override
+  public StatementResult statementShowDirectedRead() {
+    DirectedReadOptions directedReadOptions = getConnection().getDirectedRead();
+    return resultSet(
+        String.format("%sDIRECTED_READ", getNamespace(connection.getDialect())),
+        DirectedReadOptionsUtil.toString(directedReadOptions),
+        SHOW_DIRECTED_READ);
+  }
+
   @Override
   public StatementResult statementSetOptimizerVersion(String optimizerVersion) {
     getConnection().setOptimizerVersion(optimizerVersion);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DirectedReadOptionsUtil.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.connection;
+
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.common.base.Strings;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import com.google.spanner.v1.DirectedReadOptions;
+
+public class DirectedReadOptionsUtil {
+
+  /**
+   * Generates a valid JSON string for the given {@link DirectedReadOptions} that can be used with
+   * the JDBC driver.
+   */
+  public static String toString(DirectedReadOptions directedReadOptions) {
+    if (directedReadOptions == null
+        || DirectedReadOptions.getDefaultInstance().equals(directedReadOptions)) {
+      return "";
+    }
+    try {
+      return JsonFormat.printer().omittingInsignificantWhitespace().print(directedReadOptions);
+    } catch (InvalidProtocolBufferException invalidProtocolBufferException) {
+      throw SpannerExceptionFactory.asSpannerException(invalidProtocolBufferException);
+    }
+  }
+
+  static DirectedReadOptions parse(String json) {
+    if (Strings.isNullOrEmpty(json)) {
+      return DirectedReadOptions.getDefaultInstance();
+    }
+    DirectedReadOptions.Builder builder = DirectedReadOptions.newBuilder();
+    try {
+      JsonFormat.parser().merge(json, builder);
+      return builder.build();
+    } catch (InvalidProtocolBufferException invalidProtocolBufferException) {
+      throw SpannerExceptionFactory.asSpannerException(invalidProtocolBufferException);
+    }
+  }
+}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadOnlyTransaction.java

@@ -107,6 +107,11 @@ public boolean isReadOnly() {
     return true;
   }
 
+  @Override
+  public boolean supportsDirectedReads(ParsedStatement ignore) {
+    return true;
+  }
+
   @Override
   void checkAborted() {
     // No-op for read-only transactions as they cannot abort.

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java

@@ -188,6 +188,11 @@ public boolean isReadOnly() {
     return readOnly;
   }
 
+  @Override
+  public boolean supportsDirectedReads(ParsedStatement parsedStatement) {
+    return parsedStatement.isQuery();
+  }
+
   private void checkAndMarkUsed() {
     Preconditions.checkState(!used, "This single-use transaction has already been used");
     used = true;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java

@@ -63,6 +63,8 @@ enum ClientSideStatementType {
     SHOW_COMMIT_RESPONSE,
     SHOW_READ_ONLY_STALENESS,
     SET_READ_ONLY_STALENESS,
+    SHOW_DIRECTED_READ,
+    SET_DIRECTED_READ,
     SHOW_OPTIMIZER_VERSION,
     SET_OPTIMIZER_VERSION,
     SHOW_OPTIMIZER_STATISTICS_PACKAGE,