feat: add support for BatchWriteAtLeastOnce

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

@@ -16,10 +16,12 @@
 
 package com.google.cloud.spanner;
 
+import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
+import com.google.spanner.v1.BatchWriteResponse;
 
 /**
  * Interface for all the APIs that are used to read/write data into a Cloud Spanner database. An
@@ -191,6 +193,77 @@ CommitResponse writeWithOptions(Iterable<Mutation> mutations, TransactionOption.
   CommitResponse writeAtLeastOnceWithOptions(
       Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException;
 
+  /**
+   * Applies batch of mutation groups in a collection of efficient transactions. The mutation groups
+   * are applied non-atomically in an unspecified order and thus, they must be independent of each
+   * other. Partial failure is possible, i.e., some mutation groups may have been applied
+   * successfully, while some may have failed. The results of individual batches are streamed into
+   * the response as and when the batches are applied.
+   *
+   * <p>Since this method does not feature replay protection, it may attempt to apply {@code
+   * mutation groups} more than once; if the mutation groups are not idempotent, this may lead to a
+   * failure being reported when the mutation group was applied once. For example, an insert may
+   * fail with {@link ErrorCode#ALREADY_EXISTS} even though the row did not exist before this method
+   * was called. For this reason, most users of the library will prefer to use {@link
+   * #write(Iterable)} instead. However, {@code batchWriteAtLeastOnce()} method may be appropriate
+   * for non-atomically committing multiple mutation groups in a single RPC with low latency.
+   *
+   * <p>Example of BatchWriteAtLeastOnce
+   *
+   * <pre>{@code
+   * ServerStream<BatchWriteResponse> responses =
+   *     dbClient.batchWriteAtLeastOnceWithOptions(
+   *         ImmutableList.of(MUTATION_GROUP1, MUTATION_GROUP2));
+   * for (BatchWriteResponse response : responses) {
+   *   // Do something when a response is received.
+   * }
+   * }</pre>
+   *
+   * @return ServerStream\<BatchWriteResponse>
+   */
+  ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<MutationGroup> mutationGroups)
+      throws SpannerException;
+
+  /**
+   * Applies batch of mutation groups in a collection of efficient transactions. The mutation groups
+   * are applied non-atomically in an unspecified order and thus, they must be independent of each
+   * other. Partial failure is possible, i.e., some mutation groups may have been applied
+   * successfully, while some may have failed. The results of individual batches are streamed into
+   * the response as and when the batches are applied.
+   *
+   * <p>Since this method does not feature replay protection, it may attempt to apply {@code
+   * mutation groups} more than once; if the mutation groups are not idempotent, this may lead to a
+   * failure being reported when the mutation group was applied once. For example, an insert may
+   * fail with {@link ErrorCode#ALREADY_EXISTS} even though the row did not exist before this method
+   * was called. For this reason, most users of the library will prefer to use {@link
+   * #write(Iterable)} instead. However, {@code batchWriteAtLeastOnce()} method may be appropriate
+   * for non-atomically committing multiple mutation groups in a single RPC with low latency.
+   *
+   * <p>Example of BatchWriteAtLeastOnceWithOptions
+   *
+   * <pre>{@code
+   * ServerStream<BatchWriteResponse> responses =
+   *     dbClient.batchWriteAtLeastOnceWithOptions(
+   *         ImmutableList.of(MUTATION_GROUP1, MUTATION_GROUP2),
+   *         Options.tag("batch-write-tag"));
+   * for (BatchWriteResponse response : responses) {
+   *   // Do something when a response is received.
+   * }
+   * }</pre>
+   *
+   * Options for a transaction can include:
+   *
+   * <ul>
+   *   <li>{@link Options#priority(com.google.cloud.spanner.Options.RpcPriority)}: The {@link
+   *       RpcPriority} to use for the batch write request.
+   *   <li>{@link Options#tag(String)}: The transaction tag to use for the batch write request.
+   * </ul>
+   *
+   * @return ServerStream\<BatchWriteResponse>
+   */
+  ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
+      Iterable<MutationGroup> mutationGroups, TransactionOption... options) throws SpannerException;
+
   /**
    * Returns a context in which a single read can be performed using {@link TimestampBound#strong()}
    * concurrency. This method will return a {@link ReadContext} that will not return the read

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner;
 
+import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
@@ -24,6 +25,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.spanner.v1.BatchWriteResponse;
 import io.opencensus.common.Scope;
 import io.opencensus.trace.Span;
 import io.opencensus.trace.Tracer;
@@ -106,6 +108,28 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      final Iterable<MutationGroup> mutationGroups) throws SpannerException {
+    return batchWriteAtLeastOnceWithOptions(mutationGroups);
+  }
+
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
+      final Iterable<MutationGroup> mutationGroups, final TransactionOption... options)
+      throws SpannerException {
+    Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
+    try (Scope s = tracer.withSpan(span)) {
+      return runWithSessionRetry(
+          session -> session.batchWriteAtLeastOnceWithOptions(mutationGroups, options));
+    } catch (RuntimeException e) {
+      TraceUtil.setWithFailure(span, e);
+      throw e;
+    } finally {
+      span.end(TraceUtil.END_SPAN_OPTIONS);
+    }
+  }
+
   @Override
   public ReadContext singleUse() {
     Span span = tracer.spanBuilder(READ_ONLY_TRANSACTION).startSpan();

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2023 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.common.base.Preconditions;
+import com.google.spanner.v1.BatchWriteRequest;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Represents a group of Cloud Spanner mutations to be committed together.
+ */
+public class MutationGroup {
+  private final List<Mutation> mutations;
+
+  private MutationGroup(List<Mutation> mutations) {
+    this.mutations = mutations;
+  }
+
+  public static MutationGroup of(Mutation... mutations) {
+    Preconditions.checkArgument(mutations.length > 0, "Should pass in at least one mutation.");
+    return new MutationGroup(Arrays.asList(mutations));
+  }
+
+  public List<Mutation> getMutations() {
+    return mutations;
+  }
+
+  static BatchWriteRequest.MutationGroup toProto(final MutationGroup mutationGroup) {
+    List<com.google.spanner.v1.Mutation> mutationsProto = new ArrayList<>();
+    Mutation.toProto(mutationGroup.getMutations(), mutationsProto);
+    return BatchWriteRequest.MutationGroup.newBuilder().addAllMutations(mutationsProto).build();
+  }
+
+  static List<BatchWriteRequest.MutationGroup> toListProto(
+      final Iterable<MutationGroup> mutationGroups) {
+    List<BatchWriteRequest.MutationGroup> mutationGroupsProto = new ArrayList<>();
+    for (MutationGroup group : mutationGroups) {
+      mutationGroupsProto.add(toProto(group));
+    }
+    return mutationGroupsProto;
+  }
+}

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

@@ -21,6 +21,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction;
 import com.google.cloud.spanner.AbstractReadContext.SingleReadContext;
@@ -35,6 +36,8 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
+import com.google.spanner.v1.BatchWriteRequest;
+import com.google.spanner.v1.BatchWriteResponse;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.RequestOptions;
@@ -195,6 +198,43 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      Iterable<MutationGroup> mutationGroups) throws SpannerException {
+    return batchWriteAtLeastOnceWithOptions(mutationGroups);
+  }
+
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
+      Iterable<MutationGroup> mutationGroups, TransactionOption... transactionOptions)
+      throws SpannerException {
+    setActive(null);
+    Options batchWriteRequestOptions = Options.fromTransactionOptions(transactionOptions);
+    List<BatchWriteRequest.MutationGroup> mutationGroupsProto =
+        MutationGroup.toListProto(mutationGroups);
+    final BatchWriteRequest.Builder requestBuilder =
+        BatchWriteRequest.newBuilder().setSession(name).addAllMutationGroups(mutationGroupsProto);
+    if (batchWriteRequestOptions.hasPriority() || batchWriteRequestOptions.hasTag()) {
+      RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+      if (batchWriteRequestOptions.hasPriority()) {
+        requestOptionsBuilder.setPriority(batchWriteRequestOptions.priority());
+      }
+      if (batchWriteRequestOptions.hasTag()) {
+        requestOptionsBuilder.setTransactionTag(batchWriteRequestOptions.tag());
+      }
+      requestBuilder.setRequestOptions(requestOptionsBuilder.build());
+    }
+    Span span = tracer.spanBuilder(SpannerImpl.BATCH_WRITE).startSpan();
+    try (Scope s = tracer.withSpan(span)) {
+      return spanner.getRpc().batchWriteAtLeastOnce(requestBuilder.build(), this.options);
+    } catch (Throwable e) {
+      TraceUtil.setWithFailure(span, e);
+      throw SpannerExceptionFactory.newSpannerException(e);
+    } finally {
+      span.end(TraceUtil.END_SPAN_OPTIONS);
+    }
+  }
+
   @Override
   public ReadContext singleUse() {
     return singleUse(TimestampBound.strong());

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

@@ -43,6 +43,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.core.ExecutorProvider;
+import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.Timestamp;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
@@ -68,6 +69,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.protobuf.Empty;
+import com.google.spanner.v1.BatchWriteResponse;
 import com.google.spanner.v1.ResultSetStats;
 import io.opencensus.common.Scope;
 import io.opencensus.metrics.DerivedLongCumulative;
@@ -1171,6 +1173,23 @@ public CommitResponse writeAtLeastOnceWithOptions(
       }
     }
 
+    @Override
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+        Iterable<MutationGroup> mutationGroups) throws SpannerException {
+      return batchWriteAtLeastOnceWithOptions(mutationGroups);
+    }
+
+    @Override
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
+        Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+        throws SpannerException {
+      try {
+        return get().batchWriteAtLeastOnceWithOptions(mutationGroups, options);
+      } finally {
+        close();
+      }
+    }
+
     @Override
     public ReadContext singleUse() {
       try {
@@ -1450,6 +1469,24 @@ public CommitResponse writeAtLeastOnceWithOptions(
       }
     }
 
+    @Override
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+        Iterable<MutationGroup> mutationGroups) throws SpannerException {
+      return batchWriteAtLeastOnceWithOptions(mutationGroups);
+    }
+
+    @Override
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
+        Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+        throws SpannerException {
+      try {
+        markUsed();
+        return delegate.batchWriteAtLeastOnceWithOptions(mutationGroups, options);
+      } catch (SpannerException e) {
+        throw lastException = e;
+      }
+    }
+
     @Override
     public long executePartitionedUpdate(Statement stmt, UpdateOption... options)
         throws SpannerException {

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

@@ -66,6 +66,7 @@ class SpannerImpl extends BaseService<SpannerOptions> implements Spanner {
   static final String COMMIT = "CloudSpannerOperation.Commit";
   static final String QUERY = "CloudSpannerOperation.ExecuteStreamingQuery";
   static final String READ = "CloudSpannerOperation.ExecuteStreamingRead";
+  static final String BATCH_WRITE = "CloudSpannerOperation.BatchWrite";
 
   private static final Object CLIENT_ID_LOCK = new Object();
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -157,6 +157,8 @@
 import com.google.spanner.admin.instance.v1.UpdateInstanceMetadata;
 import com.google.spanner.admin.instance.v1.UpdateInstanceRequest;
 import com.google.spanner.v1.BatchCreateSessionsRequest;
+import com.google.spanner.v1.BatchWriteRequest;
+import com.google.spanner.v1.BatchWriteResponse;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.CommitResponse;
@@ -1684,6 +1686,14 @@ public ServerStream<PartialResultSet> executeStreamingPartitionedDml(
     return partitionedDmlStub.executeStreamingSqlCallable().call(request, context);
   }
 
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      BatchWriteRequest request, @Nullable Map<Option, ?> options) {
+    GrpcCallContext context =
+        newCallContext(options, request.getSession(), request, SpannerGrpc.getBatchWriteMethod());
+    return spannerStub.batchWriteCallable().call(request, context);
+  }
+
   @Override
   public StreamingCall executeQuery(
       ExecuteSqlRequest request,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java

@@ -403,6 +403,9 @@ ApiFuture<ResultSet> executeQueryAsync(
   ServerStream<PartialResultSet> executeStreamingPartitionedDml(
       ExecuteSqlRequest request, @Nullable Map<Option, ?> options, Duration timeout);
 
+  ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      BatchWriteRequest request, @Nullable Map<Option, ?> options);
+
   /**
    * Executes a query with streaming result.
    *