Merge branch 'main' into main

Author: pratickchokhani

.github/workflows/hermetic_library_generation.yaml

@@ -37,7 +37,7 @@ jobs:
       with:
         fetch-depth: 0
         token: ${{ secrets.CLOUD_JAVA_BOT_TOKEN }}
-    - uses: googleapis/sdk-platform-java/.github/scripts@v2.52.0
+    - uses: googleapis/sdk-platform-java/.github/scripts@v2.53.0
       if: env.SHOULD_RUN == 'true'
       with:
         base_ref: ${{ github.base_ref }}

README.md

@@ -19,7 +19,7 @@ If you are using Maven with [BOM][libraries-bom], add this to your pom.xml file:
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>libraries-bom</artifactId>
-      <version>26.53.0</version>
+      <version>26.54.0</version>
       <type>pom</type>
       <scope>import</scope>
     </dependency>
@@ -49,7 +49,7 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.53.0')
+implementation platform('com.google.cloud:libraries-bom:26.54.0')
 
 implementation 'com.google.cloud:google-cloud-spanner'
 ```

generation_config.yaml

@@ -1,6 +1,6 @@
-gapic_generator_version: 2.52.0
-googleapis_commitish: 3cf61b2df20eace09e6336c23f9e08859c0d87ae
-libraries_bom_version: 26.53.0
+gapic_generator_version: 2.53.0
+googleapis_commitish: 9605bff3d36fbdb1227b26bce68258c5f00815e4
+libraries_bom_version: 26.54.0
 libraries:
   - api_shortname: spanner
     name_pretty: Cloud Spanner

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

@@ -698,6 +698,9 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
     if (!isReadOnly()) {
       builder.setSeqno(getSeqNo());
     }
+    if (options.hasLastStatement()) {
+      builder.setLastStatement(options.isLastStatement());
+    }
     builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
     builder.setRequestOptions(buildRequestOptions(options));
     return builder;
@@ -743,6 +746,9 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(
     if (selector != null) {
       builder.setTransaction(selector);
     }
+    if (options.hasLastStatement()) {
+      builder.setLastStatements(options.isLastStatement());
+    }
     builder.setSeqno(getSeqNo());
     builder.setRequestOptions(buildRequestOptions(options));
     return builder;

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

@@ -55,7 +55,7 @@ public void setSpan(ISpan span) {
 
   @Override
   public void close() {
-    closeAsync();
+    SpannerApiFutures.get(closeAsync());
   }
 
   @Override
@@ -183,6 +183,10 @@ public ApiFuture<Void> rollbackAsync() {
 
   @Override
   public TransactionContextFuture resetForRetryAsync() {
+    if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) {
+      throw new IllegalStateException(
+          "resetForRetry can only be called if the previous attempt aborted");
+    }
     return new TransactionContextFutureImpl(this, internalBeginAsync(false));
   }
 

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/Options.java

@@ -108,6 +108,9 @@ public interface ReadQueryUpdateTransactionOption
   /** Marker interface to mark options applicable to Update and Write operations */
   public interface UpdateTransactionOption extends UpdateOption, TransactionOption {}
 
+  /** Marker interface for options that can be used with both executeQuery and executeUpdate. */
+  public interface QueryUpdateOption extends QueryOption, UpdateOption {}
+
   /**
    * Marker interface to mark options applicable to Create, Update and Delete operations in admin
    * API.
@@ -236,6 +239,20 @@ public static DataBoostQueryOption dataBoostEnabled(Boolean dataBoostEnabled) {
     return new DataBoostQueryOption(dataBoostEnabled);
   }
 
+  /**
+   * If set to true, this option marks the end of the transaction. The transaction should be
+   * committed or aborted after this statement executes, and attempts to execute any other requests
+   * against this transaction (including reads and queries) will be rejected. Mixing mutations with
+   * statements that are marked as the last statement is not allowed.
+   *
+   * <p>For DML statements, setting this option may cause some error reporting to be deferred until
+   * commit time (e.g. validation of unique constraints). Given this, successful execution of a DML
+   * statement should not be assumed until the transaction commits.
+   */
+  public static QueryUpdateOption lastStatement() {
+    return new LastStatementUpdateOption();
+  }
+
   /**
    * Specifying this will cause the list operation to start fetching the record from this onwards.
    */
@@ -494,6 +511,7 @@ void appendToOptions(Options options) {
   private DecodeMode decodeMode;
   private RpcOrderBy orderBy;
   private RpcLockHint lockHint;
+  private Boolean lastStatement;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -630,6 +648,14 @@ OrderBy orderBy() {
     return orderBy == null ? null : orderBy.proto;
   }
 
+  boolean hasLastStatement() {
+    return lastStatement != null;
+  }
+
+  Boolean isLastStatement() {
+    return lastStatement;
+  }
+
   boolean hasLockHint() {
     return lockHint != null;
   }
@@ -694,6 +720,9 @@ public String toString() {
     if (orderBy != null) {
       b.append("orderBy: ").append(orderBy).append(' ');
     }
+    if (lastStatement != null) {
+      b.append("lastStatement: ").append(lastStatement).append(' ');
+    }
     if (lockHint != null) {
       b.append("lockHint: ").append(lockHint).append(' ');
     }
@@ -737,6 +766,7 @@ public boolean equals(Object o) {
         && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
         && Objects.equals(directedReadOptions(), that.directedReadOptions())
         && Objects.equals(orderBy(), that.orderBy())
+        && Objects.equals(isLastStatement(), that.isLastStatement())
         && Objects.equals(lockHint(), that.lockHint());
   }
 
@@ -797,6 +827,9 @@ public int hashCode() {
     if (orderBy != null) {
       result = 31 * result + orderBy.hashCode();
     }
+    if (lastStatement != null) {
+      result = 31 * result + lastStatement.hashCode();
+    }
     if (lockHint != null) {
       result = 31 * result + lockHint.hashCode();
     }
@@ -965,4 +998,24 @@ public boolean equals(Object o) {
       return Objects.equals(filter, ((FilterOption) o).filter);
     }
   }
+
+  static final class LastStatementUpdateOption extends InternalOption implements QueryUpdateOption {
+
+    LastStatementUpdateOption() {}
+
+    @Override
+    void appendToOptions(Options options) {
+      options.lastStatement = true;
+    }
+
+    @Override
+    public int hashCode() {
+      return LastStatementUpdateOption.class.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      return o instanceof LastStatementUpdateOption;
+    }
+  }
 }