SERVER-40238 New stage alias: $set => $addFields

This reverts commit c34bc93a15cbb6fe8b222f12afac5adbab6f0737.

3 files changed, 47 insertions, 5 deletions

diff --git a/src/mongo/db/pipeline/document_source_add_fields.cpp b/src/mongo/db/pipeline/document_source_add_fields.cpp
index e95d420ac4d..319ef9776c6 100644
--- a/src/mongo/db/pipeline/document_source_add_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_add_fields.cpp
@@ -45,27 +45,33 @@ using parsed_aggregation_projection::ParsedAddFields;
 REGISTER_DOCUMENT_SOURCE(addFields,
                          LiteParsedDocumentSourceDefault::parse,
                          DocumentSourceAddFields::createFromBson);
+REGISTER_DOCUMENT_SOURCE(set,
+                         LiteParsedDocumentSourceDefault::parse,
+                         DocumentSourceAddFields::createFromBson);
 
 intrusive_ptr<DocumentSource> DocumentSourceAddFields::create(
-    BSONObj addFieldsSpec, const intrusive_ptr<ExpressionContext>& expCtx) {
+    BSONObj addFieldsSpec, const intrusive_ptr<ExpressionContext>& expCtx, StringData stageName) {
 
     const bool isIndependentOfAnyCollection = false;
     intrusive_ptr<DocumentSourceSingleDocumentTransformation> addFields(
         new DocumentSourceSingleDocumentTransformation(
             expCtx,
             ParsedAddFields::create(expCtx, addFieldsSpec),
-            "$addFields",
+            stageName.toString(),
             isIndependentOfAnyCollection));
     return addFields;
 }
 
 intrusive_ptr<DocumentSource> DocumentSourceAddFields::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
+    const auto specifiedName = elem.fieldNameStringData();
+    invariant(specifiedName == kStageName || specifiedName == kAliasNameSet);
+
     uassert(40272,
-            str::stream() << "$addFields specification stage must be an object, got "
+            str::stream() << specifiedName << " specification stage must be an object, got "
                           << typeName(elem.type()),
             elem.type() == Object);
 
-    return DocumentSourceAddFields::create(elem.Obj(), expCtx);
+    return DocumentSourceAddFields::create(elem.Obj(), expCtx, specifiedName);
 }
 }
diff --git a/src/mongo/db/pipeline/document_source_add_fields.h b/src/mongo/db/pipeline/document_source_add_fields.h
index cf24246676e..5c99a3790bb 100644
--- a/src/mongo/db/pipeline/document_source_add_fields.h
+++ b/src/mongo/db/pipeline/document_source_add_fields.h
@@ -36,14 +36,21 @@ namespace mongo {
 /**
  * $addFields adds or replaces the specified fields to/in the document while preserving the original
  * document. It is modeled on and throws the same errors as $project.
+ *
+ * This stage is also aliased as $set and functions the same way.
  */
 class DocumentSourceAddFields final {
 public:
+    static constexpr StringData kStageName = "$addFields"_sd;
+    static constexpr StringData kAliasNameSet = "$set"_sd;  // An alternate name for this stage.
+
     /**
      * Convenience method for creating a $addFields stage from 'addFieldsSpec'.
      */
     static boost::intrusive_ptr<DocumentSource> create(
-        BSONObj addFieldsSpec, const boost::intrusive_ptr<ExpressionContext>& expCtx);
+        BSONObj addFieldsSpec,
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        StringData stageName = kStageName);
 
     /**
      * Parses a $addFields stage from the user-supplied BSON.
diff --git a/src/mongo/db/pipeline/document_source_add_fields_test.cpp b/src/mongo/db/pipeline/document_source_add_fields_test.cpp
index c0b896d973a..bd8c9483ea7 100644
--- a/src/mongo/db/pipeline/document_source_add_fields_test.cpp
+++ b/src/mongo/db/pipeline/document_source_add_fields_test.cpp
@@ -73,6 +73,35 @@ TEST_F(AddFieldsTest, ShouldKeepUnspecifiedFieldsReplaceExistingFieldsAndAddNewF
     ASSERT_TRUE(addFields->getNext().isEOF());
 }
 
+TEST_F(AddFieldsTest, ShouldSerializeAndParse) {
+    auto addFields = DocumentSourceAddFields::create(BSON("a" << BSON("$const"
+                                                                      << "new")),
+                                                     getExpCtx());
+    ASSERT(addFields->getSourceName() == DocumentSourceAddFields::kStageName);
+    vector<Value> serializedArray;
+    addFields->serializeToArray(serializedArray);
+    auto serializedBson = serializedArray[0].getDocument().toBson();
+    ASSERT_BSONOBJ_EQ(serializedBson, fromjson("{$addFields: {a: {$const: 'new'}}}"));
+    addFields = DocumentSourceAddFields::createFromBson(serializedBson.firstElement(), getExpCtx());
+    ASSERT(addFields != nullptr);
+    ASSERT(addFields->getSourceName() == DocumentSourceAddFields::kStageName);
+}
+
+TEST_F(AddFieldsTest, SetAliasShouldSerializeAndParse) {
+    auto setStage = DocumentSourceAddFields::create(BSON("a" << BSON("$const"
+                                                                     << "new")),
+                                                    getExpCtx(),
+                                                    DocumentSourceAddFields::kAliasNameSet);
+    ASSERT(setStage->getSourceName() == DocumentSourceAddFields::kAliasNameSet);
+    vector<Value> serializedArray;
+    setStage->serializeToArray(serializedArray);
+    auto serializedBson = serializedArray[0].getDocument().toBson();
+    ASSERT_BSONOBJ_EQ(serializedBson, fromjson("{$set: {a: {$const: 'new'}}}"));
+    setStage = DocumentSourceAddFields::createFromBson(serializedBson.firstElement(), getExpCtx());
+    ASSERT(setStage != nullptr);
+    ASSERT(setStage->getSourceName() == DocumentSourceAddFields::kAliasNameSet);
+}
+
 TEST_F(AddFieldsTest, ShouldOptimizeInnerExpressions) {
     auto addFields = DocumentSourceAddFields::create(
         BSON("a" << BSON("$and" << BSON_ARRAY(BSON("$const" << true)))), getExpCtx());