feat: add support for BatchWriteAtLeastOnce (#2520)

* feat: add support for BatchWriteAtleastOnce * test: add batchwrite() support to MockSpannerServiceImpl * test: add commit timestamp to proto * test: add commit timestamp to proto * test: add commit timestamp to proto * consume the stream in tests * refactor tests * refactor tests * test if mutations are correctly applied * null check * skip for emulator * add method documentation * add method documentation * add method documentation * remove autogenerated code * remove autogenerated tests * batchWriteAtleastOnce -> batchWriteAtLeastOnce * batchWriteAtleastOnceWithOptions -> batchWriteAtLeastOnceWithOptions * changes based on updated batch write API * add copyright and doc * address review comments * address review comments * add more documentation --------- Co-authored-by: Arpan Mishra <akmish3@gmail.com>

Author: rajatbhatta

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,56 @@ 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>One BatchWriteResponse can contain the results for multiple MutationGroups. Inspect the
+ * indexes field to determine the MutationGroups that the BatchWriteResponse is for.
+ *
+ * <p>The mutation groups may be applied more than once. This can lead to failures if the mutation
+ * groups are non-idempotent. For example, an insert that is replayed can return an {@link
+ * ErrorCode#ALREADY_EXISTS} error. For this reason, users of the library may 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
+ * Iterable<MutationGroup> mutationGroups =
+ * ImmutableList.of(
+ * MutationGroup.of(
+ * Mutation.newInsertBuilder("FOO1").set("ID").to(1L).set("NAME").to("Bar1").build(),
+ * Mutation.newInsertBuilder("FOO2").set("ID").to(2L).set("NAME").to("Bar2").build()),
+ * MutationGroup.of(
+ * Mutation.newInsertBuilder("FOO3").set("ID").to(3L).set("NAME").to("Bar3").build(),
+ * Mutation.newInsertBuilder("FOO4").set("ID").to(4L).set("NAME").to("Bar4").build()),
+ * MutationGroup.of(
+ * Mutation.newInsertBuilder("FOO4").set("ID").to(4L).set("NAME").to("Bar4").build(),
+ * Mutation.newInsertBuilder("FOO5").set("ID").to(5L).set("NAME").to("Bar5").build()),
+ * MutationGroup.of(
+ * Mutation.newInsertBuilder("FOO6").set("ID").to(6L).set("NAME").to("Bar6").build()));
+ * ServerStream<BatchWriteResponse> responses =
+ * dbClient.batchWriteAtLeastOnce(mutationGroups, 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>
+ */
+ ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+ 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,21 @@ public CommitResponse writeAtLeastOnceWithOptions(
  }
  }
 
+ @Override
+ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+ 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.batchWriteAtLeastOnce(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,63 @@
+/*
+ * 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.common.collect.ImmutableList;
+import com.google.spanner.v1.BatchWriteRequest;
+import java.util.ArrayList;
+import java.util.List;
+
+/** Represents a group of Cloud Spanner mutations to be committed together. */
+public class MutationGroup {
+ private final ImmutableList<Mutation> mutations;
+
+ private MutationGroup(ImmutableList<Mutation> mutations) {
+ this.mutations = mutations;
+ }
+
+ /** Creates a {@code MutationGroup} given a vararg of mutations. */
+ public static MutationGroup of(Mutation... mutations) {
+ Preconditions.checkArgument(mutations.length > 0, "Should pass in at least one mutation.");
+ return new MutationGroup(ImmutableList.copyOf(mutations));
+ }
+
+ /** Creates a {@code MutationGroup} given an iterable of mutations. */
+ public static MutationGroup of(Iterable<Mutation> mutations) {
+ return new MutationGroup(ImmutableList.copyOf(mutations));
+ }
+
+ /** Returns corresponding mutations for this MutationGroup. */
+ public ImmutableList<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;
@@ -160,7 +163,6 @@ public CommitResponse writeAtLeastOnceWithOptions(
  Iterable<Mutation> mutations, TransactionOption... transactionOptions)
  throws SpannerException {
  setActive(null);
- Options commitRequestOptions = Options.fromTransactionOptions(transactionOptions);
  List<com.google.spanner.v1.Mutation> mutationsProto = new ArrayList<>();
  Mutation.toProto(mutations, mutationsProto);
  final CommitRequest.Builder requestBuilder =
@@ -172,15 +174,9 @@ public CommitResponse writeAtLeastOnceWithOptions(
  .setSingleUseTransaction(
  TransactionOptions.newBuilder()
  .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()));
- if (commitRequestOptions.hasPriority() || commitRequestOptions.hasTag()) {
- RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
- if (commitRequestOptions.hasPriority()) {
- requestOptionsBuilder.setPriority(commitRequestOptions.priority());
- }
- if (commitRequestOptions.hasTag()) {
- requestOptionsBuilder.setTransactionTag(commitRequestOptions.tag());
- }
- requestBuilder.setRequestOptions(requestOptionsBuilder.build());
+ RequestOptions commitRequestOptions = getRequestOptions(transactionOptions);
+ if (commitRequestOptions != null) {
+ requestBuilder.setRequestOptions(commitRequestOptions);
  }
  CommitRequest request = requestBuilder.build();
  Span span = tracer.spanBuilder(SpannerImpl.COMMIT).startSpan();
@@ -195,6 +191,45 @@ public CommitResponse writeAtLeastOnceWithOptions(
  }
  }
 
+ private RequestOptions getRequestOptions(TransactionOption... transactionOptions) {
+ Options requestOptions = Options.fromTransactionOptions(transactionOptions);
+ if (requestOptions.hasPriority() || requestOptions.hasTag()) {
+ RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+ if (requestOptions.hasPriority()) {
+ requestOptionsBuilder.setPriority(requestOptions.priority());
+ }
+ if (requestOptions.hasTag()) {
+ requestOptionsBuilder.setTransactionTag(requestOptions.tag());
+ }
+ return requestOptionsBuilder.build();
+ }
+ return null;
+ }
+
+ @Override
+ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+ Iterable<MutationGroup> mutationGroups, TransactionOption... transactionOptions)
+ throws SpannerException {
+ setActive(null);
+ List<BatchWriteRequest.MutationGroup> mutationGroupsProto =
+ MutationGroup.toListProto(mutationGroups);
+ final BatchWriteRequest.Builder requestBuilder =
+ BatchWriteRequest.newBuilder().setSession(name).addAllMutationGroups(mutationGroupsProto);
+ RequestOptions batchWriteRequestOptions = getRequestOptions(transactionOptions);
+ if (batchWriteRequestOptions != null) {
+ requestBuilder.setRequestOptions(batchWriteRequestOptions);
+ }
+ 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;
@@ -69,6 +70,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;
@@ -1172,6 +1174,17 @@ public CommitResponse writeAtLeastOnceWithOptions(
  }
  }
 
+ @Override
+ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+ Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+ throws SpannerException {
+ try {
+ return get().batchWriteAtLeastOnce(mutationGroups, options);
+ } finally {
+ close();
+ }
+ }
+
  @Override
  public ReadContext singleUse() {
  try {
@@ -1465,6 +1478,18 @@ public CommitResponse writeAtLeastOnceWithOptions(
  }
  }
 
+ @Override
+ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+ Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+ throws SpannerException {
+ try {
+ markUsed();
+ return delegate.batchWriteAtLeastOnce(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.
  *