Simplify Pipeline

Author: tom-andersen

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

@@ -19,11 +19,7 @@
 import com.google.cloud.firestore.pipeline.expressions.Selectable;
 import com.google.cloud.firestore.pipeline.stages.AddFields;
 import com.google.cloud.firestore.pipeline.stages.Aggregate;
-import com.google.cloud.firestore.pipeline.stages.Collection;
-import com.google.cloud.firestore.pipeline.stages.CollectionGroup;
-import com.google.cloud.firestore.pipeline.stages.Database;
 import com.google.cloud.firestore.pipeline.stages.Distinct;
-import com.google.cloud.firestore.pipeline.stages.Documents;
 import com.google.cloud.firestore.pipeline.stages.FindNearest;
 import com.google.cloud.firestore.pipeline.stages.GenericStage;
 import com.google.cloud.firestore.pipeline.stages.Limit;
@@ -33,9 +29,8 @@
 import com.google.cloud.firestore.pipeline.stages.Stage;
 import com.google.cloud.firestore.pipeline.stages.StageUtils;
 import com.google.cloud.firestore.pipeline.stages.Where;
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.ExecutePipelineRequest;
 import com.google.firestore.v1.ExecutePipelineResponse;
@@ -48,7 +43,6 @@
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
 
 /**
  * The Pipeline class provides a flexible and expressive framework for building complex data
@@ -103,32 +97,21 @@
 @BetaApi
 public final class Pipeline {
   private static Logger logger = Logger.getLogger(Pipeline.class.getName());
-  private final ImmutableList<Stage> stages;
+  private final FluentIterable<Stage> stages;
   private final Firestore db;
 
-  private Pipeline(Firestore db, List<Stage> stages) {
+  private Pipeline(Firestore db, FluentIterable<Stage> stages) {
     this.db = db;
-    this.stages = ImmutableList.copyOf(stages);
+    this.stages = stages;
   }
 
   @InternalApi
-  Pipeline(Firestore db, Collection collection) {
-    this(db, Lists.newArrayList(collection));
+  Pipeline(Firestore db, Stage stage) {
+    this(db, FluentIterable.of(stage));
   }
 
-  @InternalApi
-  Pipeline(Firestore db, CollectionGroup group) {
-    this(db, Lists.newArrayList(group));
-  }
-
-  @InternalApi
-  Pipeline(Firestore firestore, Database db) {
-    this(firestore, Lists.newArrayList(db));
-  }
-
-  @InternalApi
-  Pipeline(Firestore db, Documents docs) {
-    this(db, Lists.newArrayList(docs));
+  private Pipeline append(Stage stage) {
+    return new Pipeline(this.db, stages.append(stage));
   }
 
   /**
@@ -162,12 +145,7 @@ private Pipeline(Firestore db, List<Stage> stages) {
    */
   @BetaApi
   public Pipeline addFields(Selectable... fields) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new AddFields(PipelineUtils.selectablesToMap(fields)))
-            .build());
+    return append(new AddFields(PipelineUtils.selectablesToMap(fields)));
   }
 
   /**
@@ -201,12 +179,7 @@ public Pipeline addFields(Selectable... fields) {
    */
   @BetaApi
   public Pipeline select(Selectable... selections) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new Select(PipelineUtils.selectablesToMap(selections)))
-            .build());
+    return append(new Select(PipelineUtils.selectablesToMap(selections)));
   }
 
   /**
@@ -232,12 +205,7 @@ public Pipeline select(Selectable... selections) {
    */
   @BetaApi
   public Pipeline select(String... fields) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new Select(PipelineUtils.fieldNamesToMap(fields)))
-            .build());
+    return append(new Select(PipelineUtils.fieldNamesToMap(fields)));
   }
 
   /**
@@ -273,8 +241,7 @@ public Pipeline select(String... fields) {
    */
   @BetaApi
   public Pipeline where(FilterCondition condition) {
-    return new Pipeline(
-        this.db, ImmutableList.<Stage>builder().addAll(stages).add(new Where(condition)).build());
+    return append(new Where(condition));
   }
 
   /**
@@ -299,8 +266,7 @@ public Pipeline where(FilterCondition condition) {
    */
   @BetaApi
   public Pipeline offset(int offset) {
-    return new Pipeline(
-        this.db, ImmutableList.<Stage>builder().addAll(stages).add(new Offset(offset)).build());
+    return append(new Offset(offset));
   }
 
   /**
@@ -330,8 +296,7 @@ public Pipeline offset(int offset) {
    */
   @BetaApi
   public Pipeline limit(int limit) {
-    return new Pipeline(
-        this.db, ImmutableList.<Stage>builder().addAll(stages).add(new Limit(limit)).build());
+    return append(new Limit(limit));
   }
 
   /**
@@ -358,12 +323,7 @@ public Pipeline limit(int limit) {
    */
   @BetaApi
   public Pipeline aggregate(ExprWithAlias<Accumulator>... accumulators) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(Aggregate.withAccumulators(accumulators))
-            .build());
+    return append(Aggregate.withAccumulators(accumulators));
   }
 
   /**
@@ -400,8 +360,7 @@ public Pipeline aggregate(ExprWithAlias<Accumulator>... accumulators) {
    */
   @BetaApi
   public Pipeline aggregate(Aggregate aggregate) {
-    return new Pipeline(
-        this.db, ImmutableList.<Stage>builder().addAll(stages).add(aggregate).build());
+    return append(aggregate);
   }
 
   /**
@@ -423,12 +382,7 @@ public Pipeline aggregate(Aggregate aggregate) {
    */
   @BetaApi
   public Pipeline distinct(String... fields) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new Distinct(PipelineUtils.fieldNamesToMap(fields)))
-            .build());
+    return append(new Distinct(PipelineUtils.fieldNamesToMap(fields)));
   }
 
   /**
@@ -460,12 +414,7 @@ public Pipeline distinct(String... fields) {
    */
   @BetaApi
   public Pipeline distinct(Selectable... selectables) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new Distinct(PipelineUtils.selectablesToMap(selectables)))
-            .build());
+    return append(new Distinct(PipelineUtils.selectablesToMap(selectables)));
   }
 
   /**
@@ -528,14 +477,8 @@ public Pipeline findNearest(
   public Pipeline findNearest(
       Expr property, double[] vector, FindNearest.FindNearestOptions options) {
     // Implementation for findNearest (add the FindNearest stage if needed)
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(
-                new FindNearest(
-                    property, vector, options)) // Assuming FindNearest takes these arguments
-            .build());
+    return append(
+        new FindNearest(property, vector, options)); // Assuming FindNearest takes these arguments
   }
 
   /**
@@ -571,12 +514,7 @@ public Pipeline findNearest(
    */
   @BetaApi
   public Pipeline sort(List<Ordering> orders, Sort.Density density, Sort.Truncation truncation) {
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new Sort(orders, density, truncation))
-            .build());
+    return append(new Sort(orders, density, truncation));
   }
 
   /**
@@ -633,12 +571,7 @@ public Pipeline sort(Ordering... orders) {
   @BetaApi
   public Pipeline genericStage(String name, List<Object> params) {
     // Implementation for genericStage (add the GenericStage if needed)
-    return new Pipeline(
-        this.db,
-        ImmutableList.<Stage>builder()
-            .addAll(stages)
-            .add(new GenericStage(name, params)) // Assuming GenericStage takes a list of params
-            .build());
+    return append(new GenericStage(name, params)); // Assuming GenericStage takes a list of params
   }
 
   /**
@@ -805,11 +738,7 @@ private Value toProto() {
     return Value.newBuilder()
         .setPipelineValue(
             com.google.firestore.v1.Pipeline.newBuilder()
-                .addAllStages(
-                    stages.stream()
-                        .map(StageUtils::toStageProto)
-                        .collect(Collectors.toList())) // Use the static method
-            )
+                .addAllStages(stages.transform(StageUtils::toStageProto)))
         .build();
   }