Merge branch 'main' into generate-libraries-main

Author: rahul2393

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/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;
+    }
+  }
 }

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

@@ -265,6 +265,9 @@ static ErrorDetails extractErrorDetails(Throwable cause) {
       if (cause instanceof ApiException) {
         return ((ApiException) cause).getErrorDetails();
       }
+      if (cause instanceof SpannerException) {
+        return ((SpannerException) cause).getErrorDetails();
+      }
       prevCause = cause;
       cause = cause.getCause();
     }