Refactor and fix pipelines (#1922)

* Refactor pipelines

* Pretty

* Simplify

Author: tom-andersen

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/AbstractStage.java

@@ -16,12 +16,15 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.InternalApi;
 import com.google.cloud.firestore.pipeline.expressions.Expr;
+import com.google.firestore.v1.Pipeline;
 import java.util.Map;
 
 @InternalApi
-public final class AddFields implements Stage {
+public final class AddFields extends Stage {
 
   private static final String name = "add_fields";
   private final Map<String, Expr> fields;
@@ -31,13 +34,8 @@ public AddFields(Map<String, Expr> fields) {
     this.fields = fields;
   }
 
-  @InternalApi
-  public Map<String, Expr> getFields() {
-    return fields;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder().setName(name).addArgs(encodeValue(fields)).build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/AddFields.java

@@ -16,20 +16,22 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.BetaApi;
-import com.google.api.core.InternalApi;
 import com.google.cloud.firestore.PipelineUtils;
 import com.google.cloud.firestore.pipeline.expressions.Accumulator;
 import com.google.cloud.firestore.pipeline.expressions.Expr;
 import com.google.cloud.firestore.pipeline.expressions.ExprWithAlias;
 import com.google.cloud.firestore.pipeline.expressions.Selectable;
+import com.google.firestore.v1.Pipeline;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Map;
 import java.util.stream.Collectors;
 
 @BetaApi
-public final class Aggregate implements Stage {
+public final class Aggregate extends Stage {
 
   private static final String name = "aggregate";
   private final Map<String, Expr> groups;
@@ -63,18 +65,12 @@ private Aggregate(Map<String, Expr> groups, Map<String, Accumulator> accumulator
     this.accumulators = accumulators;
   }
 
-  @InternalApi
-  Map<String, Expr> getGroups() {
-    return groups;
-  }
-
-  @InternalApi
-  Map<String, Accumulator> getAccumulators() {
-    return accumulators;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder()
+        .setName(name)
+        .addArgs(encodeValue(accumulators))
+        .addArgs(encodeValue(groups))
+        .build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Aggregate.java

@@ -17,10 +17,12 @@
 package com.google.cloud.firestore.pipeline.stages;
 
 import com.google.api.core.InternalApi;
+import com.google.firestore.v1.Pipeline;
+import com.google.firestore.v1.Value;
 import javax.annotation.Nonnull;
 
 @InternalApi
-public final class Collection implements Stage {
+public final class Collection extends Stage {
 
   private static final String name = "collection";
   @Nonnull private final String path;
@@ -34,13 +36,11 @@ public Collection(@Nonnull String path) {
     }
   }
 
-  @InternalApi
-  public String getPath() {
-    return path;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder()
+        .setName(name)
+        .addArgs(Value.newBuilder().setReferenceValue(path).build())
+        .build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Collection.java

@@ -16,10 +16,14 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.InternalApi;
+import com.google.firestore.v1.Pipeline;
+import com.google.firestore.v1.Value;
 
 @InternalApi
-public final class CollectionGroup implements Stage {
+public final class CollectionGroup extends Stage {
 
   private static final String name = "collection_group";
   private final String collectionId;
@@ -29,13 +33,12 @@ public CollectionGroup(String collectionId) {
     this.collectionId = collectionId;
   }
 
-  @InternalApi
-  public String getCollectionId() {
-    return collectionId;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder()
+        .setName(name)
+        .addArgs(Value.newBuilder().setReferenceValue("").build())
+        .addArgs(encodeValue(collectionId))
+        .build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/CollectionGroup.java

@@ -17,16 +17,17 @@
 package com.google.cloud.firestore.pipeline.stages;
 
 import com.google.api.core.InternalApi;
+import com.google.firestore.v1.Pipeline;
 
 @InternalApi
-public final class Database implements Stage {
+public final class Database extends Stage {
   private static final String name = "database";
 
   @InternalApi
   public Database() {}
 
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder().setName(name).build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Database.java

@@ -16,13 +16,16 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.cloud.firestore.pipeline.expressions.Expr;
+import com.google.firestore.v1.Pipeline;
 import java.util.Map;
 
 @BetaApi
-public final class Distinct implements Stage {
+public final class Distinct extends Stage {
 
   private static final String name = "distinct";
   private final Map<String, Expr> groups;
@@ -32,13 +35,8 @@ public Distinct(Map<String, Expr> groups) {
     this.groups = groups;
   }
 
-  @InternalApi
-  Map<String, Expr> getGroups() {
-    return groups;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder().setName(name).addArgs(encodeValue(groups)).build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Distinct.java

@@ -18,12 +18,13 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.firestore.DocumentReference;
+import com.google.firestore.v1.Pipeline;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 
 @InternalApi
-public final class Documents implements Stage {
+public final class Documents extends Stage {
 
   private static final String name = "documents";
   private List<String> documents;
@@ -39,13 +40,12 @@ public static Documents of(DocumentReference... documents) {
         Arrays.stream(documents).map(doc -> "/" + doc.getPath()).collect(Collectors.toList()));
   }
 
-  @InternalApi
-  public List<String> getDocuments() {
-    return documents;
-  }
-
   @Override
-  public String getName() {
-    return name;
+  Pipeline.Stage toStageProto() {
+    Pipeline.Stage.Builder builder = Pipeline.Stage.newBuilder().setName(name);
+    for (String document : documents) {
+      builder.addArgsBuilder().setStringValue(document);
+    }
+    return builder.build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Documents.java

@@ -16,12 +16,15 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.cloud.firestore.pipeline.expressions.Expr;
+import com.google.firestore.v1.Pipeline;
 
 @BetaApi
-public final class FindNearest implements Stage {
+public final class FindNearest extends Stage {
 
   public interface DistanceMeasure {
 
@@ -109,28 +112,14 @@ public FindNearest(
   }
 
   @Override
-  @InternalApi
-  public String getName() {
-    return name;
-  }
-
-  @InternalApi
-  public Expr getProperty() {
-    return property;
-  }
-
-  @InternalApi
-  public double[] getVector() {
-    return vector;
-  }
-
-  @InternalApi
-  public DistanceMeasure getDistanceMeasure() {
-    return distanceMeasure;
-  }
-
-  @InternalApi
-  public FindNearestOptions getOptions() {
-    return options;
+  Pipeline.Stage toStageProto() {
+    return Pipeline.Stage.newBuilder()
+        .setName(name)
+        .addArgs(encodeValue(property))
+        .addArgs(encodeValue(vector))
+        .addArgs(encodeValue(distanceMeasure.toProtoString()))
+        .putOptions("limit", encodeValue(options.getLimit()))
+        .putOptions("distance_field", encodeValue(options.getDistanceField()))
+        .build();
   }
 }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/FindNearest.java

@@ -16,11 +16,14 @@
 
 package com.google.cloud.firestore.pipeline.stages;
 
+import static com.google.cloud.firestore.PipelineUtils.encodeValue;
+
 import com.google.api.core.InternalApi;
+import com.google.firestore.v1.Pipeline;
 import java.util.List;
 
 @InternalApi
-public final class GenericStage implements Stage {
+public final class GenericStage extends Stage {
 
   private final String name;
   private List<Object> params;
@@ -32,13 +35,11 @@ public GenericStage(String name, List<Object> params) {
   }
 
   @Override
-  @InternalApi
-  public String getName() {
-    return name;
-  }
-
-  @InternalApi
-  public List<Object> getParams() {
-    return params;
+  Pipeline.Stage toStageProto() {
+    Pipeline.Stage.Builder builder = Pipeline.Stage.newBuilder().setName(name);
+    for (Object param : params) {
+      builder.addArgs(encodeValue(param));
+    }
+    return builder.build();
   }
 }