Merge branch 'main' into mux-rw-it

Author: harshachinta

.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.42.0
+        uses: googleapis/sdk-platform-java/java-shared-dependencies/unmanaged-dependency-check@google-cloud-shared-dependencies/v3.43.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.42.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_b:3.43.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.42.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_a:3.43.0"
 }
 
 env_vars: {

google-cloud-spanner-bom/pom.xml

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

google-cloud-spanner-executor/pom.xml

@@ -46,6 +46,16 @@
       <groupId>com.google.cloud.opentelemetry</groupId>
       <artifactId>exporter-trace</artifactId>
       <version>0.33.0</version>
+      <exclusions>
+        <exclusion>
+          <groupId>io.opentelemetry.semconv</groupId>
+          <artifactId>opentelemetry-semconv</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry.semconv</groupId>
+      <artifactId>opentelemetry-semconv</artifactId>
     </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
@@ -266,7 +276,7 @@
           <groupId>org.apache.maven.plugins</groupId>
           <artifactId>maven-dependency-plugin</artifactId>
           <configuration>
-            <ignoredDependencies> com.google.api:gax,org.apache.maven.surefire:surefire-junit4 </ignoredDependencies>
+            <ignoredDependencies> com.google.api:gax,org.apache.maven.surefire:surefire-junit4,io.opentelemetry.semconv:opentelemetry-semconv </ignoredDependencies>
           </configuration>
         </plugin>
       </plugins>

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

@@ -22,6 +22,7 @@
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.ReadOption;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.Struct;
@@ -35,6 +36,7 @@
 import java.time.Instant;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantLock;
 import javax.annotation.Nullable;
@@ -59,6 +61,13 @@ public class BatchClientImpl implements BatchClient {
   @GuardedBy("multiplexedSessionLock")
   private final AtomicReference<SessionImpl> multiplexedSessionReference;
 
+  /**
+   * This flag is set to true if the server return UNIMPLEMENTED when partitioned transaction is
+   * executed on a multiplexed session. TODO: Remove once this is guaranteed to be available.
+   */
+  @VisibleForTesting
+  static final AtomicBoolean unimplementedForPartitionedOps = new AtomicBoolean(false);
+
   BatchClientImpl(SessionClient sessionClient, boolean isMultiplexedSessionEnabled) {
     this.sessionClient = checkNotNull(sessionClient);
     this.isMultiplexedSessionEnabled = isMultiplexedSessionEnabled;
@@ -85,7 +94,7 @@ public String getDatabaseRole() {
   @Override
   public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
     SessionImpl session;
-    if (isMultiplexedSessionEnabled) {
+    if (canUseMultiplexedSession()) {
       session = getMultiplexedSession();
     } else {
       session = sessionClient.createSession();
@@ -131,6 +140,10 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
         batchTransactionId);
   }
 
+  private boolean canUseMultiplexedSession() {
+    return isMultiplexedSessionEnabled && !unimplementedForPartitionedOps.get();
+  }
+
   private SessionImpl getMultiplexedSession() {
     this.multiplexedSessionLock.lock();
     try {
@@ -216,15 +229,26 @@ public List<Partition> partitionReadUsingIndex(
       builder.setPartitionOptions(pbuilder.build());
 
       final PartitionReadRequest request = builder.build();
-      PartitionResponse response = rpc.partitionRead(request, options);
-      ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
-      for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
-        Partition partition =
-            Partition.createReadPartition(
-                p.getPartitionToken(), partitionOptions, table, index, keys, columns, readOptions);
-        partitions.add(partition);
+      try {
+        PartitionResponse response = rpc.partitionRead(request, options);
+        ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
+        for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
+          Partition partition =
+              Partition.createReadPartition(
+                  p.getPartitionToken(),
+                  partitionOptions,
+                  table,
+                  index,
+                  keys,
+                  columns,
+                  readOptions);
+          partitions.add(partition);
+        }
+        return partitions.build();
+      } catch (SpannerException e) {
+        maybeMarkUnimplementedForPartitionedOps(e);
+        throw e;
       }
-      return partitions.build();
     }
 
     @Override
@@ -256,15 +280,27 @@ public List<Partition> partitionQuery(
       builder.setPartitionOptions(pbuilder.build());
 
       final PartitionQueryRequest request = builder.build();
-      PartitionResponse response = rpc.partitionQuery(request, options);
-      ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
-      for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
-        Partition partition =
-            Partition.createQueryPartition(
-                p.getPartitionToken(), partitionOptions, statement, queryOptions);
-        partitions.add(partition);
+      try {
+        PartitionResponse response = rpc.partitionQuery(request, options);
+        ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
+        for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
+          Partition partition =
+              Partition.createQueryPartition(
+                  p.getPartitionToken(), partitionOptions, statement, queryOptions);
+          partitions.add(partition);
+        }
+        return partitions.build();
+      } catch (SpannerException e) {
+        maybeMarkUnimplementedForPartitionedOps(e);
+        throw e;
+      }
+    }
+
+    void maybeMarkUnimplementedForPartitionedOps(SpannerException spannerException) {
+      if (MultiplexedSessionDatabaseClient.verifyErrorMessage(
+          spannerException, "Partitioned operations are not supported with multiplexed sessions")) {
+        unimplementedForPartitionedOps.set(true);
       }
-      return partitions.build();
     }
 
     @Override

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

@@ -124,6 +124,12 @@ private boolean canUseMultiplexedSessionsForRW() {
         && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsForRWSupported();
   }
 
+  private boolean canUseMultiplexedSessionsForPartitionedOps() {
+    return this.useMultiplexedSessionPartitionedOps
+        && this.multiplexedSessionDatabaseClient != null
+        && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsForPartitionedOpsSupported();
+  }
+
   @Override
   public Dialect getDialect() {
     return pool.getDialect();
@@ -323,8 +329,15 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
 
   @Override
   public long executePartitionedUpdate(final Statement stmt, final UpdateOption... options) {
-    if (useMultiplexedSessionPartitionedOps) {
-      return getMultiplexedSession().executePartitionedUpdate(stmt, options);
+
+    if (canUseMultiplexedSessionsForPartitionedOps()) {
+      try {
+        return getMultiplexedSession().executePartitionedUpdate(stmt, options);
+      } catch (SpannerException e) {
+        if (!multiplexedSessionDatabaseClient.maybeMarkUnimplementedForPartitionedOps(e)) {
+          throw e;
+        }
+      }
     }
     return executePartitionedUpdateWithPooledSession(stmt, options);
   }

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

@@ -104,6 +104,11 @@ void onError(SpannerException spannerException) {
       // UNIMPLEMENTED with error message "Transaction type read_write not supported with
       // multiplexed sessions" is returned.
       this.client.maybeMarkUnimplementedForRW(spannerException);
+      // Mark multiplexed sessions for Partitioned Ops as unimplemented and fall back to regular
+      // sessions if
+      // UNIMPLEMENTED with error message "Partitioned operations are not supported with multiplexed
+      // sessions".
+      this.client.maybeMarkUnimplementedForPartitionedOps(spannerException);
     }
 
     @Override
@@ -214,6 +219,12 @@ public void close() {
    */
   @VisibleForTesting final AtomicBoolean unimplementedForRW = new AtomicBoolean(false);
 
+  /**
+   * This flag is set to true if the server return UNIMPLEMENTED when partitioned transaction is
+   * executed on a multiplexed session. TODO: Remove once this is guaranteed to be available.
+   */
+  @VisibleForTesting final AtomicBoolean unimplementedForPartitionedOps = new AtomicBoolean(false);
+
   MultiplexedSessionDatabaseClient(SessionClient sessionClient) {
     this(sessionClient, Clock.systemUTC());
   }
@@ -316,7 +327,18 @@ && verifyErrorMessage(
     }
   }
 
-  private boolean verifyErrorMessage(SpannerException spannerException, String message) {
+  boolean maybeMarkUnimplementedForPartitionedOps(SpannerException spannerException) {
+    if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED
+        && verifyErrorMessage(
+            spannerException,
+            "Transaction type partitioned_dml not supported with multiplexed sessions")) {
+      unimplementedForPartitionedOps.set(true);
+      return true;
+    }
+    return false;
+  }
+
+  static boolean verifyErrorMessage(SpannerException spannerException, String message) {
     if (spannerException.getCause() == null) {
       return false;
     }
@@ -391,6 +413,10 @@ boolean isMultiplexedSessionsForRWSupported() {
     return !this.unimplementedForRW.get();
   }
 
+  boolean isMultiplexedSessionsForPartitionedOpsSupported() {
+    return !this.unimplementedForPartitionedOps.get();
+  }
+
   void close() {
     synchronized (this) {
       if (!this.isClosed) {

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

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2025 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.InternalApi;
+import com.google.common.annotations.VisibleForTesting;
+import java.math.BigInteger;
+import java.security.SecureRandom;
+import java.util.Objects;
+
+@InternalApi
+public class XGoogSpannerRequestId {
+  // 1. Generate the random process Id singleton.
+  @VisibleForTesting
+  static final String RAND_PROCESS_ID = XGoogSpannerRequestId.generateRandProcessId();
+
+  @VisibleForTesting
+  static final long VERSION = 1; // The version of the specification being implemented.
+
+  private final long nthClientId;
+  private final long nthChannelId;
+  private final long nthRequest;
+  private long attempt;
+
+  XGoogSpannerRequestId(long nthClientId, long nthChannelId, long nthRequest, long attempt) {
+    this.nthClientId = nthClientId;
+    this.nthChannelId = nthChannelId;
+    this.nthRequest = nthRequest;
+    this.attempt = attempt;
+  }
+
+  public static XGoogSpannerRequestId of(
+      long nthClientId, long nthChannelId, long nthRequest, long attempt) {
+    return new XGoogSpannerRequestId(nthClientId, nthChannelId, nthRequest, attempt);
+  }
+
+  private static String generateRandProcessId() {
+    // Expecting to use 64-bits of randomness to avoid clashes.
+    BigInteger bigInt = new BigInteger(64, new SecureRandom());
+    return String.format("%016x", bigInt);
+  }
+
+  @Override
+  public String toString() {
+    return String.format(
+        "%d.%s.%d.%d.%d.%d",
+        XGoogSpannerRequestId.VERSION,
+        XGoogSpannerRequestId.RAND_PROCESS_ID,
+        this.nthClientId,
+        this.nthChannelId,
+        this.nthRequest,
+        this.attempt);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    // instanceof for a null object returns false.
+    if (!(other instanceof XGoogSpannerRequestId)) {
+      return false;
+    }
+
+    XGoogSpannerRequestId otherReqId = (XGoogSpannerRequestId) (other);
+
+    return Objects.equals(this.nthClientId, otherReqId.nthClientId)
+        && Objects.equals(this.nthChannelId, otherReqId.nthChannelId)
+        && Objects.equals(this.nthRequest, otherReqId.nthRequest)
+        && Objects.equals(this.attempt, otherReqId.attempt);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(this.nthClientId, this.nthChannelId, this.nthRequest, this.attempt);
+  }
+}

google-cloud-spanner/src/main/resources/META-INF/native-image/native-image.properties

@@ -1,4 +1,5 @@
 Args = --initialize-at-build-time=com.google.cloud.spanner.IntegrationTestEnv,\
   org.junit.experimental.categories.CategoryValidator,\
-  org.junit.validator.AnnotationValidator \
+  org.junit.validator.AnnotationValidator,\
+  java.lang.annotation.Annotation \
   --features=com.google.cloud.spanner.nativeimage.SpannerFeature