Refactor terminal stages

Author: bwaldvogel

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java

@@ -265,9 +265,9 @@ public void validate(Document query) {
         }
     }
 
-    private boolean isLastStage(TerminalStage aggregationStage) {
+    private boolean isLastStage(TerminalStage stage) {
         Assert.notEmpty(stages);
-        return stages.indexOf(aggregationStage) == stages.size() - 1;
+        return stages.indexOf(stage) == stages.size() - 1;
     }
 
     public boolean isModifying() {

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/MergeStage.java

@@ -26,7 +26,7 @@
 import de.bwaldvogel.mongo.exception.MongoServerError;
 import de.bwaldvogel.mongo.exception.TypeMismatchException;
 
-public class MergeStage implements TerminalStage {
+public class MergeStage extends TerminalStage {
 
     private static final Set<String> KNOWN_KEYS = new HashSet<>(Arrays.asList("into", "on", "whenMatched", "whenNotMatched"));
 
@@ -166,7 +166,7 @@ public String name() {
     }
 
     @Override
-    public Stream<Document> apply(Stream<Document> stream) {
+    public void applyLast(Stream<Document> stream) {
         MongoCollection<?> collection = targetCollectionSupplier.get();
 
         stream.forEach(document -> {
@@ -178,11 +178,7 @@ public Stream<Document> apply(Stream<Document> stream) {
                     case merge:
                         Document mergedDocument = existingDocument.clone();
                         mergedDocument.merge(document);
-                        if (!mergedDocument.get("_id").equals(existingDocument.get("_id"))) {
-                            throw new ImmutableFieldException("$merge failed to update the matching document, did you attempt to modify the _id or the shard key?" +
-                                " :: caused by :: " +
-                                "Performing an update on the path '_id' would modify the immutable field '_id'");
-                        }
+                        assertIdHasNotChanged(existingDocument, mergedDocument);
                         replaceDocument(collection, existingDocument, mergedDocument);
                         break;
                     case replace:
@@ -211,8 +207,14 @@ public Stream<Document> apply(Stream<Document> stream) {
                 }
             }
         });
+    }
 
-        return Stream.empty();
+    private static void assertIdHasNotChanged(Document one, Document other) {
+        if (!one.get("_id").equals(other.get("_id"))) {
+            throw new ImmutableFieldException("$merge failed to update the matching document, did you attempt to modify the _id or the shard key?" +
+                " :: caused by :: " +
+                "Performing an update on the path '_id' would modify the immutable field '_id'");
+        }
     }
 
     private Document getJoinQuery(Document document) {

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/OutStage.java

@@ -1,6 +1,5 @@
 package de.bwaldvogel.mongo.backend.aggregation.stage;
 
-import java.util.Collections;
 import java.util.stream.Stream;
 
 import org.slf4j.Logger;
@@ -12,7 +11,7 @@
 import de.bwaldvogel.mongo.exception.IllegalOperationException;
 import de.bwaldvogel.mongo.oplog.NoopOplog;
 
-public class OutStage implements TerminalStage {
+public class OutStage extends TerminalStage {
 
     private static final Logger log = LoggerFactory.getLogger(OutStage.class);
 
@@ -30,20 +29,19 @@ public String name() {
     }
 
     @Override
-    public Stream<Document> apply(Stream<Document> stream) {
+    public void applyLast(Stream<Document> stream) {
         if (this.collectionName.contains("$")) {
             throw new IllegalOperationException("error with target namespace: Invalid collection name: " + this.collectionName);
         }
         String tempCollectionName = "_tmp" + System.currentTimeMillis() + "_" + this.collectionName;
         MongoCollection<?> tempCollection = database.createCollectionOrThrowIfExists(tempCollectionName);
-        stream.forEach(document -> tempCollection.insertDocuments(Collections.singletonList(document)));
+        stream.forEach(tempCollection::addDocument);
         MongoCollection<?> existingCollection = database.resolveCollection(this.collectionName, false);
         if (existingCollection != null) {
             log.info("Dropping existing collection {}", existingCollection);
             database.dropCollection(this.collectionName, NoopOplog.get());
         }
         database.moveCollection(database, tempCollection, this.collectionName);
-        return Stream.empty();
     }
 
     @Override

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/TerminalStage.java

@@ -1,7 +1,18 @@
 package de.bwaldvogel.mongo.backend.aggregation.stage;
 
-public interface TerminalStage extends AggregationStage {
+import java.util.stream.Stream;
 
-    String name();
+import de.bwaldvogel.mongo.bson.Document;
 
+public abstract class TerminalStage implements AggregationStage {
+
+    public abstract String name();
+
+    abstract void applyLast(Stream<Document> stream);
+
+    @Override
+    public final Stream<Document> apply(Stream<Document> stream) {
+        applyLast(stream);
+        return Stream.empty();
+    }
 }