add transaction support

Author: wu-hui

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

@@ -52,12 +52,14 @@
 import com.google.firestore.v1.ExecutePipelineRequest;
 import com.google.firestore.v1.ExecutePipelineResponse;
 import com.google.firestore.v1.StructuredPipeline;
+import com.google.protobuf.ByteString;
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import javax.annotation.Nullable;
 
 /**
  * The Pipeline class provides a flexible and expressive framework for building complex data
@@ -597,29 +599,7 @@ public Pipeline genericStage(String name, List<Object> params) {
    */
   @BetaApi
   public ApiFuture<List<PipelineResult>> execute() {
-    SettableApiFuture<List<PipelineResult>> futureResult = SettableApiFuture.create();
-
-    execute( // Assuming you have this method
-        new PipelineResultObserver() {
-          final List<PipelineResult> results = new ArrayList<>();
-
-          @Override
-          public void onCompleted() {
-            futureResult.set(results);
-          }
-
-          @Override
-          public void onNext(PipelineResult result) {
-            results.add(result);
-          }
-
-          @Override
-          public void onError(Throwable t) {
-            futureResult.setException(t);
-          }
-        });
-
-    return futureResult;
+    return execute(null, null);
   }
 
   /**
@@ -669,14 +649,57 @@ public void onError(Throwable t) {
    */
   @BetaApi
   public void execute(ApiStreamObserver<PipelineResult> observer) {
-    ExecutePipelineRequest request =
+    executeInternal(null, null, observer);
+  }
+
+  ApiFuture<List<PipelineResult>> execute(
+      @Nullable final ByteString transactionId, @Nullable com.google.protobuf.Timestamp readTime) {
+    SettableApiFuture<List<PipelineResult>> futureResult = SettableApiFuture.create();
+
+    executeInternal(
+        transactionId,
+        readTime,
+        new PipelineResultObserver() {
+          final List<PipelineResult> results = new ArrayList<>();
+
+          @Override
+          public void onCompleted() {
+            futureResult.set(results);
+          }
+
+          @Override
+          public void onNext(PipelineResult result) {
+            results.add(result);
+          }
+
+          @Override
+          public void onError(Throwable t) {
+            futureResult.setException(t);
+          }
+        });
+
+    return futureResult;
+  }
+
+  void executeInternal(
+      @Nullable final ByteString transactionId,
+      @Nullable com.google.protobuf.Timestamp readTime,
+      ApiStreamObserver<PipelineResult> observer) {
+    ExecutePipelineRequest.Builder request =
         ExecutePipelineRequest.newBuilder()
             .setDatabase(rpcContext.getDatabaseName())
-            .setStructuredPipeline(StructuredPipeline.newBuilder().setPipeline(toProto()).build())
-            .build();
+            .setStructuredPipeline(StructuredPipeline.newBuilder().setPipeline(toProto()).build());
+
+    if (transactionId != null) {
+      request.setTransaction(transactionId);
+    }
+
+    if (readTime != null) {
+      request.setReadTime(readTime);
+    }
 
     pipelineInternalStream(
-        request,
+        request.build(),
         new PipelineResultObserver() {
           @Override
           public void onCompleted() {

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

@@ -125,6 +125,14 @@ public ApiFuture<AggregateQuerySnapshot> get(@Nonnull AggregateQuery query) {
     }
   }
 
+  @Nonnull
+  @Override
+  public ApiFuture<List<PipelineResult>> execute(@Nonnull Pipeline pipeline) {
+    try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+      return pipeline.execute(null, readTime);
+    }
+  }
+
   @Nonnull
   @Override
   public Transaction create(

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

@@ -260,4 +260,12 @@ public ApiFuture<AggregateQuerySnapshot> get(@Nonnull AggregateQuery query) {
       return query.get(transactionId, null);
     }
   }
+
+  @Nonnull
+  @Override
+  public ApiFuture<List<PipelineResult>> execute(@Nonnull Pipeline pipeline) {
+    try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+      return pipeline.execute(transactionId, null);
+    }
+  }
 }

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

@@ -17,6 +17,7 @@
 package com.google.cloud.firestore;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.core.BetaApi;
 import com.google.api.core.InternalExtensionOnly;
 import com.google.cloud.firestore.telemetry.TraceUtil;
 import com.google.cloud.firestore.telemetry.TraceUtil.Context;
@@ -134,4 +135,9 @@ public abstract ApiFuture<List<DocumentSnapshot>> getAll(
    */
   @Nonnull
   public abstract ApiFuture<AggregateQuerySnapshot> get(@Nonnull AggregateQuery query);
+
+  /** @return The result of the aggregation. */
+  @Nonnull
+  @BetaApi
+  public abstract ApiFuture<List<PipelineResult>> execute(@Nonnull Pipeline pipeline);
 }

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java

@@ -43,6 +43,7 @@
 
 import com.google.cloud.firestore.CollectionReference;
 import com.google.cloud.firestore.LocalFirestoreHelper;
+import com.google.cloud.firestore.Pipeline;
 import com.google.cloud.firestore.PipelineResult;
 import com.google.cloud.firestore.pipeline.expressions.Constant;
 import com.google.cloud.firestore.pipeline.expressions.Field;
@@ -1013,4 +1014,35 @@ public void testNestedFields() throws Exception {
                 map("title", "The Hitchhiker's Guide to the Galaxy", "awards.hugo", true),
                 map("title", "Dune", "awards.hugo", true)));
   }
+
+  @Test
+  public void testPipelineInTransactions() throws Exception {
+    Pipeline pipeline =
+        collection
+            .pipeline()
+            .where(eq("awards.hugo", true))
+            .select("title", "awards.hugo", Field.DOCUMENT_ID);
+
+    firestore
+        .runTransaction(
+            transaction -> {
+              List<PipelineResult> results = transaction.execute(pipeline).get();
+
+              assertThat(data(results))
+                  .isEqualTo(
+                      Lists.newArrayList(
+                          map("title", "The Hitchhiker's Guide to the Galaxy", "awards.hugo", true),
+                          map("title", "Dune", "awards.hugo", true)));
+
+              transaction.update(collection.document("book1"), map("foo", "bar"));
+
+              return "done";
+            })
+        .get();
+
+    List<PipelineResult> result =
+        collection.pipeline().where(eq("foo", "bar")).select("title").execute().get();
+    assertThat(data(result))
+        .isEqualTo(Lists.newArrayList(map("title", "The Hitchhiker's Guide to the Galaxy")));
+  }
 }