feat: add success and error callbacks to BulkWriter

google-cloud-firestore/clirr-ignored-differences.xml

@@ -223,6 +223,30 @@
     <to>java.util.List</to>
   </difference>
 
+  <!--
+  UpdateBuilder
+  -->
+  <difference>
+    <differenceType>6001</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>pendingOperations</field>
+  </difference>
+  <difference>
+    <differenceType>6001</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>state</field>
+  </difference>
+  <difference>
+    <differenceType>6010</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>writes</field>
+  </difference>
+  <difference>
+    <differenceType>7009</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <method>int getMutationsSize()</method>
+  </difference>
+
   <!--
   FakeCredentials Refactor
   com.google.cloud.firestore.FirestoreOptions$Builder$FakeCredentials -> com.google.cloud.firestore.FirestoreOptions$EmulatorCredentials

google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java

@@ -16,32 +16,177 @@
 
 package com.google.cloud.firestore;
 
+import com.google.api.core.ApiAsyncFunction;
+import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
+import com.google.cloud.Timestamp;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.firestore.v1.BatchWriteRequest;
+import com.google.firestore.v1.BatchWriteResponse;
+import io.grpc.Status;
+import io.opencensus.trace.AttributeValue;
+import io.opencensus.trace.Tracing;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArraySet;
+import javax.annotation.Nullable;
 
 /** Used to represent a batch on the BatchQueue. */
 class BulkCommitBatch extends UpdateBuilder<ApiFuture<WriteResult>> {
+  /**
+   * Used to represent the state of batch.
+   *
+   * <p>Writes can only be added while the batch is OPEN. For a batch to be sent, the batch must be
+   * READY_TO_SEND. After a batch is sent, it is marked as SENT.
+   */
+  enum BatchState {
+    OPEN,
+    READY_TO_SEND,
+    SENT,
+  }
+
+  private BatchState state = BatchState.OPEN;
+
+  private final List<SettableApiFuture<BatchWriteResult>> pendingOperations = new ArrayList<>();
+  private final Set<DocumentReference> documents = new CopyOnWriteArraySet<>();
+  private final int maxBatchSize;
 
   BulkCommitBatch(FirestoreImpl firestore, int maxBatchSize) {
-    super(firestore, maxBatchSize);
+    super(firestore);
+    this.maxBatchSize = maxBatchSize;
   }
 
-  BulkCommitBatch(FirestoreImpl firestore, BulkCommitBatch retryBatch) {
-    super(firestore);
+  ApiFuture<WriteResult> wrapResult(DocumentReference documentReference) {
+    return processLastOperation(documentReference);
+  }
+
+  /**
+   * Commits all pending operations to the database and verifies all preconditions.
+   *
+   * <p>The writes in the batch are not applied atomically and can be applied out of order.
+   */
+  ApiFuture<List<BatchWriteResult>> bulkCommit() {
+    Tracing.getTracer()
+        .getCurrentSpan()
+        .addAnnotation(
+            TraceUtil.SPAN_NAME_BATCHWRITE,
+            ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(getWrites().size())));
+
+    Preconditions.checkState(
+        isReadyToSend(), "The batch should be marked as READY_TO_SEND before committing");
+    state = BatchState.SENT;
 
-    // Create a new BulkCommitBatch containing only the indexes from the provided indexes to retry.
-    for (int index : retryBatch.getPendingIndexes()) {
-      this.writes.add(retryBatch.writes.get(index));
+    final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder();
+    request.setDatabase(firestore.getDatabaseName());
+
+    for (WriteOperation writeOperation : getWrites()) {
+      request.addWrites(writeOperation.write);
     }
 
+    ApiFuture<BatchWriteResponse> response =
+        firestore.sendRequest(request.build(), firestore.getClient().batchWriteCallable());
+
+    committed = true;
+
+    return ApiFutures.transform(
+        response,
+        new ApiFunction<BatchWriteResponse, List<BatchWriteResult>>() {
+          @Override
+          public List<BatchWriteResult> apply(BatchWriteResponse batchWriteResponse) {
+            List<com.google.firestore.v1.WriteResult> writeResults =
+                batchWriteResponse.getWriteResultsList();
+
+            List<com.google.rpc.Status> statuses = batchWriteResponse.getStatusList();
+
+            List<BatchWriteResult> result = new ArrayList<>();
+
+            for (int i = 0; i < writeResults.size(); ++i) {
+              com.google.firestore.v1.WriteResult writeResult = writeResults.get(i);
+              com.google.rpc.Status status = statuses.get(i);
+              Status code = Status.fromCodeValue(status.getCode());
+              @Nullable Timestamp updateTime = null;
+              @Nullable Exception exception = null;
+              if (code == Status.OK) {
+                updateTime = Timestamp.fromProto(writeResult.getUpdateTime());
+              } else {
+                exception = FirestoreException.serverRejected(code, status.getMessage());
+              }
+              result.add(new BatchWriteResult(updateTime, exception));
+            }
+
+            return result;
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  int getPendingOperationCount() {
+    return pendingOperations.size();
+  }
+
+  ApiFuture<WriteResult> processLastOperation(DocumentReference documentReference) {
     Preconditions.checkState(
-        retryBatch.state == BatchState.SENT,
-        "Batch should be SENT when creating a new BulkCommitBatch for retry");
-    this.state = retryBatch.state;
-    this.pendingOperations = retryBatch.pendingOperations;
+        !documents.contains(documentReference),
+        "Batch should not contain writes to the same document");
+    documents.add(documentReference);
+    Preconditions.checkState(state == BatchState.OPEN, "Batch should be OPEN when adding writes");
+    SettableApiFuture<BatchWriteResult> resultFuture = SettableApiFuture.create();
+    pendingOperations.add(resultFuture);
+
+    if (getPendingOperationCount() == maxBatchSize) {
+      state = BatchState.READY_TO_SEND;
+    }
+
+    return ApiFutures.transformAsync(
+        resultFuture,
+        new ApiAsyncFunction<BatchWriteResult, WriteResult>() {
+          public ApiFuture<WriteResult> apply(BatchWriteResult batchWriteResult) throws Exception {
+            if (batchWriteResult.getException() == null) {
+              return ApiFutures.immediateFuture(new WriteResult(batchWriteResult.getWriteTime()));
+            } else {
+              throw batchWriteResult.getException();
+            }
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  /**
+   * Resolves the individual operations in the batch with the results and removes the entry from the
+   * pendingOperations map if the result is not retryable.
+   */
+  void processResults(List<BatchWriteResult> results) {
+    for (int i = 0; i < results.size(); i++) {
+      SettableApiFuture<BatchWriteResult> resultFuture = pendingOperations.get(i);
+      BatchWriteResult result = results.get(i);
+      if (result.getException() == null) {
+        resultFuture.set(result);
+      } else {
+        resultFuture.setException(result.getException());
+      }
+    }
+  }
+
+  void markReadyToSend() {
+    if (state == BatchState.OPEN) {
+      state = BatchState.READY_TO_SEND;
+    }
+  }
+
+  boolean isOpen() {
+    return state == BatchState.OPEN;
+  }
+
+  boolean isReadyToSend() {
+    return state == BatchState.READY_TO_SEND;
   }
 
-  ApiFuture<WriteResult> wrapResult(ApiFuture<WriteResult> result) {
-    return result;
+  boolean has(DocumentReference documentReference) {
+    return documents.contains(documentReference);
   }
 }