changes based on updated batch write API

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

@@ -194,75 +194,58 @@ CommitResponse writeAtLeastOnceWithOptions(
       Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException;
 
   /**
-   * Batches the supplied mutations in a collection of efficient transactions. The mutations are
-   * applied non-atomically in an unspecified order and thus, they must be independent of each
-   * other. Partial failure is possible, i.e., some mutations 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.
+   * 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
-   * mutations} more than once; if the mutations are not idempotent, this may lead to a failure
-   * being reported when the mutation 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 mutations in a single RPC with low latency.
+   * 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
+   * <p>Example of BatchWriteAtLeastOnce
    *
    * <pre>{@code
-   * long singerId = my_singer_id;
-   * Mutation mutation = Mutation.newInsertBuilder("Singers")
-   *         .set("SingerId")
-   *         .to(singerId)
-   *         .set("FirstName")
-   *         .to("Billy")
-   *         .set("LastName")
-   *         .to("Joel")
-   *         .build();
    * ServerStream<BatchWriteResponse> responses =
-   *     dbClient.batchWriteAtLeastOnce(Collections.singletonList(mutation));
+   *     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<Mutation> mutations)
+  ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<MutationGroup> mutationGroups)
       throws SpannerException;
 
   /**
-   * Batches the supplied mutations in a collection of efficient transactions. The mutations are
-   * applied non-atomically in an unspecified order and thus, they must be independent of each
-   * other. Partial failure is possible, i.e., some mutations 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.
+   * 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
-   * mutations} more than once; if the mutations are not idempotent, this may lead to a failure
-   * being reported when the mutation 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 mutations in a single RPC with low latency.
+   * 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
+   * <p>Example of BatchWriteAtLeastOnceWithOptions
    *
    * <pre>{@code
-   * long singerId = my_singer_id;
-   * Mutation mutation = Mutation.newInsertBuilder("Singers")
-   *         .set("SingerId")
-   *         .to(singerId)
-   *         .set("FirstName")
-   *         .to("Billy")
-   *         .set("LastName")
-   *         .to("Joel")
-   *         .build();
    * ServerStream<BatchWriteResponse> responses =
-   *     dbClient.batchWriteAtLeastOnce(
-   *             Collections.singletonList(mutation),
-   *             Options.priority(RpcPriority.LOW));
+   *     dbClient.batchWriteAtLeastOnceWithOptions(
+   *         ImmutableList.of(MUTATION_GROUP1, MUTATION_GROUP2),
+   *         Options.tag("batch-write-tag"));
    * for (BatchWriteResponse response : responses) {
    *   // Do something when a response is received.
    * }
@@ -279,7 +262,7 @@ ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<Mutation> mutati
    * @return ServerStream\<BatchWriteResponse>
    */
   ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
-      Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException;
+      Iterable<MutationGroup> mutationGroups, TransactionOption... options) throws SpannerException;
 
   /**
    * Returns a context in which a single read can be performed using {@link TimestampBound#strong()}

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

@@ -109,19 +109,19 @@ public CommitResponse writeAtLeastOnceWithOptions(
   }
 
   @Override
-  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(final Iterable<Mutation> mutations)
-      throws SpannerException {
-    return batchWriteAtLeastOnceWithOptions(mutations);
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      final Iterable<MutationGroup> mutationGroups) throws SpannerException {
+    return batchWriteAtLeastOnceWithOptions(mutationGroups);
   }
 
   @Override
   public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
-      final Iterable<Mutation> mutations, final TransactionOption... options)
+      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(mutations, options));
+          session -> session.batchWriteAtLeastOnceWithOptions(mutationGroups, options));
     } catch (RuntimeException e) {
       TraceUtil.setWithFailure(span, e);
       throw e;

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

@@ -0,0 +1,39 @@
+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;
+
+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

@@ -199,21 +199,21 @@ public CommitResponse writeAtLeastOnceWithOptions(
   }
 
   @Override
-  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<Mutation> mutations)
-      throws SpannerException {
-    return batchWriteAtLeastOnceWithOptions(mutations);
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      Iterable<MutationGroup> mutationGroups) throws SpannerException {
+    return batchWriteAtLeastOnceWithOptions(mutationGroups);
   }
 
   @Override
   public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
-      Iterable<Mutation> mutations, TransactionOption... transactionOptions)
+      Iterable<MutationGroup> mutationGroups, TransactionOption... transactionOptions)
       throws SpannerException {
     setActive(null);
     Options batchWriteRequestOptions = Options.fromTransactionOptions(transactionOptions);
-    List<com.google.spanner.v1.Mutation> mutationsProto = new ArrayList<>();
-    Mutation.toProto(mutations, mutationsProto);
+    List<BatchWriteRequest.MutationGroup> mutationGroupsProto =
+        MutationGroup.toListProto(mutationGroups);
     final BatchWriteRequest.Builder requestBuilder =
-        BatchWriteRequest.newBuilder().setSession(name).addAllMutations(mutationsProto);
+        BatchWriteRequest.newBuilder().setSession(name).addAllMutationGroups(mutationGroupsProto);
     if (batchWriteRequestOptions.hasPriority() || batchWriteRequestOptions.hasTag()) {
       RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
       if (batchWriteRequestOptions.hasPriority()) {

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

@@ -1173,16 +1173,17 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
 
     @Override
-    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<Mutation> mutations)
-        throws SpannerException {
-      return batchWriteAtLeastOnceWithOptions(mutations);
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+        Iterable<MutationGroup> mutationGroups) throws SpannerException {
+      return batchWriteAtLeastOnceWithOptions(mutationGroups);
     }
 
     @Override
     public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
-        Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException {
+        Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+        throws SpannerException {
       try {
-        return get().batchWriteAtLeastOnceWithOptions(mutations, options);
+        return get().batchWriteAtLeastOnceWithOptions(mutationGroups, options);
       } finally {
         close();
       }
@@ -1436,17 +1437,18 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
 
     @Override
-    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(Iterable<Mutation> mutations)
-        throws SpannerException {
-      return batchWriteAtLeastOnceWithOptions(mutations);
+    public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+        Iterable<MutationGroup> mutationGroups) throws SpannerException {
+      return batchWriteAtLeastOnceWithOptions(mutationGroups);
     }
 
     @Override
     public ServerStream<BatchWriteResponse> batchWriteAtLeastOnceWithOptions(
-        Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException {
+        Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+        throws SpannerException {
       try {
         markUsed();
-        return delegate.batchWriteAtLeastOnceWithOptions(mutations, options);
+        return delegate.batchWriteAtLeastOnceWithOptions(mutationGroups, options);
       } catch (SpannerException e) {
         throw lastException = e;
       }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -133,12 +133,19 @@ public class DatabaseClientImplTest {
   private static final long UPDATE_COUNT = 1L;
   private static final com.google.rpc.Status STATUS_OK =
       com.google.rpc.Status.newBuilder().setCode(com.google.rpc.Code.OK_VALUE).build();
-  private static final Iterable<Mutation> MUTATIONS =
+  private static final Iterable<MutationGroup> MUTATION_GROUPS =
       ImmutableList.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(),
-          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("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()));
   private static final Iterable<BatchWriteResponse> BATCH_WRITE_RESPONSES =
       ImmutableList.of(
           BatchWriteResponse.newBuilder()
@@ -339,7 +346,7 @@ public void testBatchWriteAtLeastOnce() {
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
 
-    ServerStream<BatchWriteResponse> responseStream = client.batchWriteAtLeastOnce(MUTATIONS);
+    ServerStream<BatchWriteResponse> responseStream = client.batchWriteAtLeastOnce(MUTATION_GROUPS);
     int idx = 0;
     for (BatchWriteResponse response : responseStream) {
       assertEquals(
@@ -353,7 +360,7 @@ public void testBatchWriteAtLeastOnce() {
     List<BatchWriteRequest> requests = mockSpanner.getRequestsOfType(BatchWriteRequest.class);
     assertEquals(requests.size(), 1);
     BatchWriteRequest request = requests.get(0);
-    assertEquals(request.getMutationsCount(), 4);
+    assertEquals(request.getMutationGroupsCount(), 4);
     assertEquals(request.getRequestOptions().getPriority(), Priority.PRIORITY_UNSPECIFIED);
   }
 
@@ -362,14 +369,14 @@ public void testBatchWriteAtLeastOnceWithOptions() {
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     ServerStream<BatchWriteResponse> responseStream =
-        client.batchWriteAtLeastOnceWithOptions(MUTATIONS, Options.priority(RpcPriority.LOW));
+        client.batchWriteAtLeastOnceWithOptions(MUTATION_GROUPS, Options.priority(RpcPriority.LOW));
     for (BatchWriteResponse response : responseStream) {}
 
     assertNotNull(responseStream);
     List<BatchWriteRequest> requests = mockSpanner.getRequestsOfType(BatchWriteRequest.class);
     assertEquals(requests.size(), 1);
     BatchWriteRequest request = requests.get(0);
-    assertEquals(request.getMutationsCount(), 4);
+    assertEquals(request.getMutationGroupsCount(), 4);
     assertEquals(request.getRequestOptions().getPriority(), Priority.PRIORITY_LOW);
   }
 
@@ -378,14 +385,15 @@ public void testBatchWriteAtLeastOnceWithTagOptions() {
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     ServerStream<BatchWriteResponse> responseStream =
-        client.batchWriteAtLeastOnceWithOptions(MUTATIONS, Options.tag("app=spanner,env=test"));
+        client.batchWriteAtLeastOnceWithOptions(
+            MUTATION_GROUPS, Options.tag("app=spanner,env=test"));
     for (BatchWriteResponse response : responseStream) {}
 
     assertNotNull(responseStream);
     List<BatchWriteRequest> requests = mockSpanner.getRequestsOfType(BatchWriteRequest.class);
     assertEquals(requests.size(), 1);
     BatchWriteRequest request = requests.get(0);
-    assertEquals(request.getMutationsCount(), 4);
+    assertEquals(request.getMutationGroupsCount(), 4);
     assertEquals(request.getRequestOptions().getTransactionTag(), "app=spanner,env=test");
     assertThat(request.getRequestOptions().getRequestTag()).isEmpty();
   }