Merge branch 'main' of https://github.com/googleapis/java-spanner into grpc-telemetry-client-interceptor

Author: nareshz

.github/workflows/unmanaged_dependency_check.yaml

@@ -17,6 +17,6 @@ jobs:
           # repository
           .kokoro/build.sh
       - name: Unmanaged dependency check
-        uses: googleapis/sdk-platform-java/java-shared-dependencies/unmanaged-dependency-check@google-cloud-shared-dependencies/v3.33.0
+        uses: googleapis/sdk-platform-java/java-shared-dependencies/unmanaged-dependency-check@google-cloud-shared-dependencies/v3.34.0
         with:
           bom-path: google-cloud-spanner-bom/pom.xml

.kokoro/presubmit/graalvm-native-17.cfg

@@ -3,7 +3,7 @@
 # Configure the docker image for kokoro-trampoline.
 env_vars: {
   key: "TRAMPOLINE_IMAGE"
-  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_b:3.33.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_b:3.34.0"
 }
 
 env_vars: {

.kokoro/presubmit/graalvm-native.cfg

@@ -3,7 +3,7 @@
 # Configure the docker image for kokoro-trampoline.
 env_vars: {
   key: "TRAMPOLINE_IMAGE"
-  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_a:3.33.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_a:3.34.0"
 }
 
 env_vars: {

benchmarks/pom.xml

@@ -34,7 +34,7 @@
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
     <junixsocket.version>2.10.0</junixsocket.version>
-    <opentelemetry.version>1.40.0</opentelemetry.version>
+    <opentelemetry.version>1.41.0</opentelemetry.version>
   </properties>
 
   <dependencies>
@@ -85,19 +85,19 @@
     <dependency>
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-bom</artifactId>
-      <version>1.40.0</version>
+      <version>1.41.0</version>
       <type>pom</type>
       <scope>import</scope>
     </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner</artifactId>
-      <version>6.71.0</version>
+      <version>6.72.0</version>
     </dependency>
     <dependency>
       <groupId>commons-cli</groupId>
       <artifactId>commons-cli</artifactId>
-      <version>1.8.0</version>
+      <version>1.9.0</version>
     </dependency>
     <dependency>
       <groupId>com.google.auto.value</groupId>
@@ -118,7 +118,7 @@
     <dependency>
       <groupId>commons-cli</groupId>
       <artifactId>commons-cli</artifactId>
-      <version>1.8.0</version>
+      <version>1.9.0</version>
     </dependency>
 
     <dependency>
@@ -133,7 +133,7 @@
       <plugin>
         <groupId>org.codehaus.mojo</groupId>
         <artifactId>exec-maven-plugin</artifactId>
-        <version>3.4.0</version>
+        <version>3.4.1</version>
         <configuration>
           <mainClass>com.google.cloud.spanner.benchmark.LatencyBenchmark</mainClass>
           <cleanupDaemonThreads>false</cleanupDaemonThreads>

google-cloud-spanner-bom/pom.xml

@@ -8,7 +8,7 @@
   <parent>
     <groupId>com.google.cloud</groupId>
     <artifactId>sdk-platform-java-config</artifactId>
-    <version>3.33.0</version>
+    <version>3.34.0</version>
   </parent>
 
   <name>Google Cloud Spanner BOM</name>

google-cloud-spanner-executor/pom.xml

@@ -129,7 +129,7 @@
     <dependency>
       <groupId>commons-cli</groupId>
       <artifactId>commons-cli</artifactId>
-      <version>1.8.0</version>
+      <version>1.9.0</version>
     </dependency>
     <dependency>
       <groupId>commons-io</groupId>
@@ -188,7 +188,7 @@
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-failsafe-plugin</artifactId>
-        <version>3.3.1</version>
+        <version>3.4.0</version>
       </plugin>
     </plugins>
     <pluginManagement>

google-cloud-spanner/pom.xml

@@ -359,6 +359,12 @@
       <version>${graal-sdk.version}</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>org.graalvm.sdk</groupId>
+      <artifactId>nativeimage</artifactId>
+      <version>${graal-sdk.version}</version>
+      <scope>provided</scope>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
@@ -411,7 +417,7 @@
     <dependency>
       <groupId>org.hamcrest</groupId>
       <artifactId>hamcrest</artifactId>
-      <version>2.2</version>
+      <version>3.0</version>
       <scope>test</scope>
     </dependency>
     <!-- Benchmarking dependencies -->

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

@@ -36,6 +36,7 @@
 import com.google.cloud.spanner.SessionClient.SessionOption;
 import com.google.cloud.spanner.SessionImpl.SessionTransaction;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
+import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
@@ -45,7 +46,6 @@
 import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
-import com.google.spanner.v1.ExecuteSqlRequest.Builder;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.PartialResultSet;
@@ -69,6 +69,7 @@ abstract class AbstractReadContext
 
   abstract static class Builder<B extends Builder<?, T>, T extends AbstractReadContext> {
     private SessionImpl session;
+    private boolean cancelQueryWhenClientIsClosed;
     private SpannerRpc rpc;
     private ISpan span;
     private TraceWrapper tracer;
@@ -91,6 +92,11 @@ B setSession(SessionImpl session) {
       return self();
     }
 
+    B setCancelQueryWhenClientIsClosed(boolean cancelQueryWhenClientIsClosed) {
+      this.cancelQueryWhenClientIsClosed = cancelQueryWhenClientIsClosed;
+      return self();
+    }
+
     B setRpc(SpannerRpc rpc) {
       this.rpc = rpc;
       return self();
@@ -184,7 +190,7 @@ static Builder newBuilder() {
     @GuardedBy("lock")
     private boolean used;
 
-    private final Map<SpannerRpc.Option, ?> channelHint;
+    private Map<SpannerRpc.Option, ?> channelHint;
 
     private SingleReadContext(Builder builder) {
       super(builder);
@@ -227,6 +233,16 @@ TransactionSelector getTransactionSelector() {
     Map<SpannerRpc.Option, ?> getTransactionChannelHint() {
       return channelHint;
     }
+
+    @Override
+    boolean prepareRetryOnDifferentGrpcChannel() {
+      if (session.getIsMultiplexed() && channelHint.get(Option.CHANNEL_HINT) != null) {
+        long channelHintForTransaction = Option.CHANNEL_HINT.getLong(channelHint) + 1L;
+        channelHint = optionMap(SessionOption.channelHint(channelHintForTransaction));
+        return true;
+      }
+      return super.prepareRetryOnDifferentGrpcChannel();
+    }
   }
 
   private static void assertTimestampAvailable(boolean available) {
@@ -440,6 +456,7 @@ void initTransaction() {
 
   final Object lock = new Object();
   final SessionImpl session;
+  final boolean cancelQueryWhenClientIsClosed;
   final SpannerRpc rpc;
   final ExecutorProvider executorProvider;
   ISpan span;
@@ -469,6 +486,7 @@ void initTransaction() {
 
   AbstractReadContext(Builder<?, ?> builder) {
     this.session = builder.session;
+    this.cancelQueryWhenClientIsClosed = builder.cancelQueryWhenClientIsClosed;
     this.rpc = builder.rpc;
     this.defaultPrefetchChunks = builder.defaultPrefetchChunks;
     this.defaultQueryOptions = builder.defaultQueryOptions;
@@ -745,11 +763,13 @@ ResultSet executeQueryInternalWithOptions(
             span,
             tracer,
             tracer.createStatementAttributes(statement, options),
+            session.getErrorHandler(),
             rpc.getExecuteQueryRetrySettings(),
             rpc.getExecuteQueryRetryableCodes()) {
           @Override
           CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
-            GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
+            GrpcStreamIterator stream =
+                new GrpcStreamIterator(statement, prefetchChunks, cancelQueryWhenClientIsClosed);
             if (partitionToken != null) {
               request.setPartitionToken(partitionToken);
             }
@@ -774,6 +794,11 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             stream.setCall(call, request.getTransaction().hasBegin());
             return stream;
           }
+
+          @Override
+          boolean prepareIteratorForRetryOnDifferentGrpcChannel() {
+            return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel();
+          }
         };
     return new GrpcResultSet(
         stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode);
@@ -840,6 +865,10 @@ public void close() {
    */
   abstract Map<SpannerRpc.Option, ?> getTransactionChannelHint();
 
+  boolean prepareRetryOnDifferentGrpcChannel() {
+    return false;
+  }
+
   /**
    * Returns the transaction tag for this {@link AbstractReadContext} or <code>null</code> if this
    * {@link AbstractReadContext} does not have a transaction tag.
@@ -918,11 +947,13 @@ ResultSet readInternalWithOptions(
             SpannerImpl.READ,
             span,
             tracer,
+            session.getErrorHandler(),
             rpc.getReadRetrySettings(),
             rpc.getReadRetryableCodes()) {
           @Override
           CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
-            GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
+            GrpcStreamIterator stream =
+                new GrpcStreamIterator(prefetchChunks, cancelQueryWhenClientIsClosed);
             TransactionSelector selector = null;
             if (resumeToken != null) {
               builder.setResumeToken(resumeToken);
@@ -945,6 +976,11 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());
             return stream;
           }
+
+          @Override
+          boolean prepareIteratorForRetryOnDifferentGrpcChannel() {
+            return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel();
+          }
         };
     return new GrpcResultSet(
         stream, this, readOptions.hasDecodeMode() ? readOptions.decodeMode() : defaultDecodeMode);

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

@@ -54,6 +54,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
     return new BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.newBuilder()
             .setSession(session)
+            .setCancelQueryWhenClientIsClosed(true)
             .setRpc(sessionClient.getSpanner().getRpc())
             .setTimestampBound(bound)
             .setDefaultQueryOptions(
@@ -75,6 +76,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
     return new BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.newBuilder()
             .setSession(session)
+            .setCancelQueryWhenClientIsClosed(true)
             .setRpc(sessionClient.getSpanner().getRpc())
             .setTransactionId(batchTransactionId.getTransactionId())
             .setTimestamp(batchTransactionId.getTimestamp())

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

@@ -0,0 +1,50 @@
+/*
+ * 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;
+
+import com.google.api.core.BetaApi;
+import javax.annotation.Nonnull;
+
+/**
+ * The {@link ErrorHandler} interface can be used to implement custom error and retry handling for
+ * specific cases. The default implementation does nothing and falls back to the standard error and
+ * retry handling in Gax and the Spanner client.
+ */
+@BetaApi
+interface ErrorHandler {
+  @Nonnull
+  Throwable translateException(@Nonnull Throwable exception);
+
+  int getMaxAttempts();
+
+  class DefaultErrorHandler implements ErrorHandler {
+    static final DefaultErrorHandler INSTANCE = new DefaultErrorHandler();
+
+    private DefaultErrorHandler() {}
+
+    @Nonnull
+    @Override
+    public Throwable translateException(@Nonnull Throwable exception) {
+      return exception;
+    }
+
+    @Override
+    public int getMaxAttempts() {
+      return 0;
+    }
+  }
+}

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

@@ -38,7 +38,7 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
   private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
   private static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
 
-  private final ConsumerImpl consumer = new ConsumerImpl();
+  private final ConsumerImpl consumer;
   private final BlockingQueue<PartialResultSet> stream;
   private final Statement statement;
 
@@ -49,13 +49,15 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
   private SpannerException error;
 
   @VisibleForTesting
-  GrpcStreamIterator(int prefetchChunks) {
-    this(null, prefetchChunks);
+  GrpcStreamIterator(int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
+    this(null, prefetchChunks, cancelQueryWhenClientIsClosed);
   }
 
   @VisibleForTesting
-  GrpcStreamIterator(Statement statement, int prefetchChunks) {
+  GrpcStreamIterator(
+      Statement statement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
     this.statement = statement;
+    this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed);
     // One extra to allow for END_OF_STREAM message.
     this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
   }
@@ -136,6 +138,12 @@ private void addToStream(PartialResultSet results) {
   }
 
   private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
+    private final boolean cancelQueryWhenClientIsClosed;
+
+    ConsumerImpl(boolean cancelQueryWhenClientIsClosed) {
+      this.cancelQueryWhenClientIsClosed = cancelQueryWhenClientIsClosed;
+    }
+
     @Override
     public void onPartialResultSet(PartialResultSet results) {
       addToStream(results);
@@ -168,5 +176,10 @@ public void onError(SpannerException e) {
       error = e;
       addToStream(END_OF_STREAM);
     }
+
+    @Override
+    public boolean cancelQueryWhenClientIsClosed() {
+      return this.cancelQueryWhenClientIsClosed;
+    }
   }
 }