SERVER-35896: Support 'replaceDocuments' mode in $out

2 files changed, 154 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/out/batch_writes.js b/jstests/aggregation/sources/out/batch_writes.js
new file mode 100644
index 00000000000..6bbaac7737a
--- /dev/null
+++ b/jstests/aggregation/sources/out/batch_writes.js
@@ -0,0 +1,61 @@
+// Tests the behavior of an $out stage which encounters an error in the middle of processing. We
+// don't guarantee any particular behavior in this scenario, but this test exists to make sure
+// nothing horrendous happens and to characterize the current behavior.
+// @tags: [assumes_unsharded_collection]
+(function() {
+    "use strict";
+
+    load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+    const coll = db.batch_writes;
+    const outColl = db.batch_writes_out;
+    coll.drop();
+    outColl.drop();
+
+    // Test with 2 very large documents that do not fit into a single batch.
+    const kSize15MB = 15 * 1024 * 1024;
+    const largeArray = new Array(kSize15MB).join("a");
+    assert.commandWorked(coll.insert({_id: 0, a: largeArray}));
+    assert.commandWorked(coll.insert({_id: 1, a: largeArray}));
+
+    // Make sure the $out succeeds without any duplicate keys.
+    ["replaceCollection", "insertDocuments", "replaceDocuments"].forEach(mode => {
+        coll.aggregate([{$out: {to: outColl.getName(), mode: mode}}]);
+        assert.eq(2, outColl.find().itcount());
+        outColl.drop();
+    });
+
+    coll.drop();
+    for (let i = 0; i < 10; i++) {
+        assert.commandWorked(coll.insert({_id: i, a: i}));
+    }
+
+    // Create a unique index on 'a' in the output collection to create a unique key violation when
+    // running the $out. The second document to be written ({_id: 1, a: 1}) will conflict with the
+    // existing document in the output collection. We use a unique index on a field other than _id
+    // because "replaceDocuments" mode will not change _id when one already exists.
+    outColl.drop();
+    assert.commandWorked(outColl.insert({_id: 2, a: 1}));
+    assert.commandWorked(outColl.createIndex({a: 1}, {unique: true}));
+
+    // Test that both batched updates and inserts will successfully write the first document but
+    // fail on the second. We don't guarantee any particular behavior in this case, but this test is
+    // meant to characterize the current behavior.
+    assertErrorCode(coll, [{$out: {to: outColl.getName(), mode: "insertDocuments"}}], 16996);
+    assert.soon(() => {
+        return outColl.find().itcount() == 2;
+    });
+
+    assertErrorCode(coll, [{$out: {to: outColl.getName(), mode: "replaceDocuments"}}], 50904);
+    assert.soon(() => {
+        return outColl.find().itcount() == 2;
+    });
+
+    // Mode "replaceCollection" will drop the contents of the output collection, so there is no
+    // duplicate key error.
+    outColl.drop();
+    assert.commandWorked(outColl.insert({_id: 2, a: 1}));
+    assert.commandWorked(outColl.createIndex({a: 1}, {unique: true}));
+    coll.aggregate([{$out: {to: outColl.getName(), mode: "replaceCollection"}}]);
+    assert.eq(10, outColl.find().itcount());
+}());
diff --git a/jstests/aggregation/sources/out/mode_replace_documents.js b/jstests/aggregation/sources/out/mode_replace_documents.js
new file mode 100644
index 00000000000..8335981a1ff
--- /dev/null
+++ b/jstests/aggregation/sources/out/mode_replace_documents.js
@@ -0,0 +1,93 @@
+// Tests for the $out stage with mode set to "replaceDocuments".
+// @tags: [assumes_unsharded_collection]
+(function() {
+    "use strict";
+
+    load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+    const coll = db.replace_docs;
+    const outColl = db.replace_docs_out;
+    coll.drop();
+    outColl.drop();
+
+    const nDocs = 10;
+    for (let i = 0; i < nDocs; i++) {
+        assert.commandWorked(coll.insert({_id: i, a: i}));
+    }
+
+    // Test that a $out with 'replaceDocuments' mode will default the unique key to "_id".
+    coll.aggregate([{$out: {to: outColl.getName(), mode: "replaceDocuments"}}]);
+    assert.eq(nDocs, outColl.find().itcount());
+
+    // Test that 'replaceDocuments' mode will update existing documents that match the unique key.
+    const nDocsReplaced = 5;
+    coll.aggregate([
+        {$project: {_id: {$mod: ["$_id", nDocsReplaced]}}},
+        {$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1}}}
+    ]);
+    assert.eq(nDocsReplaced, outColl.find({a: {$exists: false}}).itcount());
+
+    // Test 'replaceDocuments' mode with a dotted path unique key.
+    coll.drop();
+    outColl.drop();
+    assert.commandWorked(coll.insert({_id: 0, a: {b: 1}}));
+    assert.commandWorked(coll.insert({_id: 1, a: {b: 1}, c: 1}));
+    coll.aggregate([
+        {$addFields: {_id: 0}},
+        {$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1, "a.b": 1}}}
+    ]);
+    assert.eq([{_id: 0, a: {b: 1}, c: 1}], outColl.find().toArray());
+
+    // TODO SERVER-36100: 'replaceDocuments' mode should allow a missing "_id" unique key.
+    assertErrorCode(coll,
+                    [
+                      {$project: {_id: 0}},
+                      {
+                        $out: {
+                            to: outColl.getName(),
+                            mode: "replaceDocuments",
+                        }
+                      }
+                    ],
+                    50905);
+
+    // Test that 'replaceDocuments' mode with a missing non-id unique key fails.
+    assertErrorCode(
+        coll,
+        [{$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {missing: 1}}}],
+        50905);
+
+    // Test that a replace fails to insert a document if it violates a unique index constraint. In
+    // this example, $out will attempt to insert multiple documents with {a: 0} which is not allowed
+    // with the unique index on {a: 1}.
+    coll.drop();
+    assert.commandWorked(coll.insert([{_id: 0}, {_id: 1}]));
+
+    outColl.drop();
+    assert.commandWorked(outColl.createIndex({a: 1}, {unique: true}));
+    assertErrorCode(
+        coll,
+        [{$addFields: {a: 0}}, {$out: {to: outColl.getName(), mode: "replaceDocuments"}}],
+        50904);
+
+    // Test that $out fails if the unique key contains an array.
+    coll.drop();
+    assert.commandWorked(coll.insert({_id: 0, a: [1, 2]}));
+    assertErrorCode(
+        coll,
+        [
+          {$addFields: {_id: 0}},
+          {$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1, "a.b": 1}}}
+        ],
+        50905);
+
+    coll.drop();
+    assert.commandWorked(coll.insert({_id: 0, a: [{b: 1}]}));
+    assertErrorCode(
+        coll,
+        [
+          {$addFields: {_id: 0}},
+          {$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1, "a.b": 1}}}
+        ],
+        50905);
+}());