1 files changed, 10 insertions, 4 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);
 }
 }