CSHARP-2624: Add $replaceWith stage.

Author: DmitryLukyanov

src/MongoDB.Driver/AggregateFluent.cs

@@ -187,6 +187,11 @@ public override IAggregateFluent<TNewResult> ReplaceRoot<TNewResult>(AggregateEx
             return WithPipeline(_pipeline.ReplaceRoot(newRoot));
         }
 
+        public override IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateExpressionDefinition<TResult, TNewResult> newRoot)
+        {
+            return WithPipeline(_pipeline.ReplaceWith(newRoot));
+        }
+
         public override IAggregateFluent<TResult> Skip(int skip)
         {
             return WithPipeline(_pipeline.Skip(skip));

src/MongoDB.Driver/AggregateFluentBase.cs

@@ -167,6 +167,12 @@ public virtual IAggregateFluent<TNewResult> ReplaceRoot<TNewResult>(AggregateExp
             throw new NotImplementedException();
         }
 
+        /// <inheritdoc />
+        public virtual IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateExpressionDefinition<TResult, TNewResult> newRoot)
+        {
+            throw new NotImplementedException();
+        }
+
         /// <inheritdoc />
         public abstract IAggregateFluent<TResult> Skip(int skip);
 

src/MongoDB.Driver/IAggregateFluent.cs

@@ -276,6 +276,14 @@ IAggregateFluent<TNewResult> Lookup<TForeignDocument, TAsElement, TAs, TNewResul
         /// <returns>The fluent aggregate interface.</returns>
         IAggregateFluent<TNewResult> ReplaceRoot<TNewResult>(AggregateExpressionDefinition<TResult, TNewResult> newRoot);
 
+        /// <summary>
+        /// Appends a $replaceWith stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TNewResult">The type of the new result.</typeparam>
+        /// <param name="newRoot">The new root.</param>
+        /// <returns>The fluent aggregate interface.</returns>
+        IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateExpressionDefinition<TResult, TNewResult> newRoot);
+
         /// <summary>
         /// Appends a skip stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/IAggregateFluentExtensions.cs

@@ -493,6 +493,24 @@ public static IAggregateFluent<TNewResult> ReplaceRoot<TResult, TNewResult>(
             return aggregate.AppendStage(PipelineStageDefinitionBuilder.ReplaceRoot(newRoot));
         }
 
+        /// <summary>
+        /// Appends a $replaceWith stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <typeparam name="TNewResult">The type of the new result.</typeparam>
+        /// <param name="aggregate">The aggregate.</param>
+        /// <param name="newRoot">The new root.</param>
+        /// <returns>
+        /// The fluent aggregate interface.
+        /// </returns>
+        public static IAggregateFluent<TNewResult> ReplaceWith<TResult, TNewResult>(
+            this IAggregateFluent<TResult> aggregate,
+            Expression<Func<TResult, TNewResult>> newRoot)
+        {
+            Ensure.IsNotNull(aggregate, nameof(aggregate));
+            return aggregate.AppendStage(PipelineStageDefinitionBuilder.ReplaceWith(newRoot));
+        }
+
         /// <summary>
         /// Appends an ascending sort stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/PipelineDefinitionBuilder.cs

@@ -946,6 +946,46 @@ public static PipelineDefinition<TInput, TOutput> ReplaceRoot<TInput, TIntermedi
             return pipeline.AppendStage(PipelineStageDefinitionBuilder.ReplaceRoot(newRoot, translationOptions));
         }
 
+        /// <summary>
+        /// Appends a $replaceWith stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TIntermediate">The type of the intermediate documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="newRoot">The new root.</param>
+        /// <returns>
+        /// A new pipeline with an additional stage.
+        /// </returns>
+        public static PipelineDefinition<TInput, TOutput> ReplaceWith<TInput, TIntermediate, TOutput>(
+            this PipelineDefinition<TInput, TIntermediate> pipeline,
+            AggregateExpressionDefinition<TIntermediate, TOutput> newRoot)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(PipelineStageDefinitionBuilder.ReplaceWith(newRoot));
+        }
+
+        /// <summary>
+        /// Appends a $replaceWith stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TIntermediate">The type of the intermediate documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="newRoot">The new root.</param>
+        /// <param name="translationOptions">The translation options.</param>
+        /// <returns>
+        /// The fluent aggregate interface.
+        /// </returns>
+        public static PipelineDefinition<TInput, TOutput> ReplaceWith<TInput, TIntermediate, TOutput>(
+            this PipelineDefinition<TInput, TIntermediate> pipeline,
+            Expression<Func<TIntermediate, TOutput>> newRoot,
+            ExpressionTranslationOptions translationOptions = null)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(PipelineStageDefinitionBuilder.ReplaceWith(newRoot, translationOptions));
+        }
+
         /// <summary>
         /// Appends a $skip stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/PipelineStageDefinitionBuilder.cs

@@ -1077,6 +1077,47 @@ public static PipelineStageDefinition<TInput, TOutput> ReplaceRoot<TInput, TOutp
             return ReplaceRoot(new ExpressionAggregateExpressionDefinition<TInput, TOutput>(newRoot, translationOptions));
         }
 
+        /// <summary>
+        /// Creates a $replaceWith stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="newRoot">The new root.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TOutput> ReplaceWith<TInput, TOutput>(
+            AggregateExpressionDefinition<TInput, TOutput> newRoot)
+        {
+            Ensure.IsNotNull(newRoot, nameof(newRoot));
+
+            const string operatorName = "$replaceWith";
+            var stage = new DelegatedPipelineStageDefinition<TInput, TOutput>(
+                operatorName,
+                (s, sr) =>
+                {
+                    var document = new BsonDocument(operatorName, newRoot.Render(s, sr));
+                    var outputSerializer = sr.GetSerializer<TOutput>();
+                    return new RenderedPipelineStageDefinition<TOutput>(operatorName, document, outputSerializer);
+                });
+
+            return stage;
+        }
+
+        /// <summary>
+        /// Creates a $replaceWith stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="newRoot">The new root.</param>
+        /// <param name="translationOptions">The translation options.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TOutput> ReplaceWith<TInput, TOutput>(
+            Expression<Func<TInput, TOutput>> newRoot,
+            ExpressionTranslationOptions translationOptions = null)
+        {
+            Ensure.IsNotNull(newRoot, nameof(newRoot));
+            return ReplaceWith(new ExpressionAggregateExpressionDefinition<TInput, TOutput>(newRoot, translationOptions));
+        }
+
         /// <summary>
         /// Creates a $skip stage.
         /// </summary>

src/MongoDB.Driver/UpdateDefinitionBuilder.cs

@@ -1052,6 +1052,16 @@ public UpdateDefinition<TDocument> Mul<TField>(Expression<Func<TDocument, TField
             return Mul(new ExpressionFieldDefinition<TDocument, TField>(field), value);
         }
 
+        /// <summary>
+        /// Creates an update pipeline.
+        /// </summary>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <returns>An update pipeline.</returns>
+        public UpdateDefinition<TDocument> Pipeline(PipelineDefinition<TDocument, TDocument> pipeline)
+        {
+            return new PipelineUpdateDefinition<TDocument>(pipeline);
+        }
+
         /// <summary>
         /// Creates a pop operator.
         /// </summary>

tests/MongoDB.Driver.Tests/AggregateFluentTests.cs

@@ -628,6 +628,47 @@ public void ReplaceRoot_should_add_the_expected_stage(
             }
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public void ReplaceWith_should_add_the_expected_stage([Values(false, true)] bool async)
+        {
+            var subject = CreateSubject();
+
+            var result = subject.ReplaceWith<BsonDocument>("$X");
+
+            Predicate<PipelineDefinition<C, BsonDocument>> isExpectedPipeline = pipeline =>
+            {
+                var renderedPipeline = RenderPipeline(pipeline);
+                return
+                    renderedPipeline.Documents.Count == 1 &&
+                    renderedPipeline.Documents[0] == BsonDocument.Parse("{ $replaceWith : '$X' }") &&
+                    renderedPipeline.OutputSerializer.ValueType == typeof(BsonDocument);
+            };
+
+            if (async)
+            {
+                result.ToCursorAsync().GetAwaiter().GetResult();
+
+                _mockCollection.Verify(
+                    c => c.AggregateAsync<BsonDocument>(
+                        It.Is<PipelineDefinition<C, BsonDocument>>(pipeline => isExpectedPipeline(pipeline)),
+                        It.IsAny<AggregateOptions>(),
+                        CancellationToken.None),
+                    Times.Once);
+            }
+            else
+            {
+                result.ToCursor();
+
+                _mockCollection.Verify(
+                    c => c.Aggregate<BsonDocument>(
+                        It.Is<PipelineDefinition<C, BsonDocument>>(pipeline => isExpectedPipeline(pipeline)),
+                        It.IsAny<AggregateOptions>(),
+                        CancellationToken.None),
+                    Times.Once);
+            }
+        }
+
         [Theory]
         [ParameterAttributeData]
         public void SortByCount_should_add_the_expected_stage(

tests/MongoDB.Driver.Tests/IAggregateFluentExtensionsTests.cs

@@ -288,6 +288,28 @@ public void ReplaceRoot_should_generate_the_correct_stage_with_anonymous_class()
             AssertLast(subject, expectedStage);
         }
 
+        [Fact]
+        public void ReplaceWith_should_generate_the_correct_stage()
+        {
+            var subject = CreateSubject()
+                .ReplaceWith(x => x.PhoneNumber);
+
+            var expectedStage = BsonDocument.Parse("{ $replaceWith : '$PhoneNumber' }");
+
+            AssertLast(subject, expectedStage);
+        }
+
+        [Fact]
+        public void ReplaceWith_should_generate_the_correct_stage_with_anonymous_class()
+        {
+            var subject = CreateSubject()
+                .ReplaceWith(x => new { Name = x.FirstName + " " + x.LastName });
+
+            var expectedStage = BsonDocument.Parse("{ $replaceWith : { Name : { $concat : [ '$FirstName', ' ', '$LastName' ] } } }");
+
+            AssertLast(subject, expectedStage);
+        }
+
         [Theory]
         [ParameterAttributeData]
         public void Single_should_add_limit_and_call_ToCursor(

tests/MongoDB.Driver.Tests/PipelineUpdateDefinitionTests.cs

@@ -40,9 +40,19 @@ public void PipelineUpdateDefinition_should_work_with_pipeline_builder()
             result.Should().Be("[{ \"$addFields\" : { \"x\" : 2 } }]");
         }
 
+        [Fact]
+        public void PipelineUpdateDefinition_should_work_with_pipeline_builder_and_replaceWith()
+        {
+            var pipeline = new EmptyPipelineDefinition<BsonDocument>()
+                .ReplaceWith((AggregateExpressionDefinition<BsonDocument, BsonDocument>)"{ _id : \"$_id\", s : { $sum : [\"$X\", \"$Y\"] } }");
+            var subject = CreateSubject(pipeline);
+            var result = subject.ToString();
+            result.Should().Be("[{ \"$replaceWith\" : { \"_id\" : \"$_id\", \"s\" : { \"$sum\" : [\"$X\", \"$Y\"] } } }]");
+        }
+
         private PipelineUpdateDefinition<BsonDocument> CreateSubject(params string[] stages)
         {
-            return CreateSubject(PipelineDefinition<BsonDocument,BsonDocument>.Create(stages));
+            return CreateSubject(PipelineDefinition<BsonDocument, BsonDocument>.Create(stages));
         }
 
         private PipelineUpdateDefinition<BsonDocument> CreateSubject(PipelineDefinition<BsonDocument, BsonDocument> pipeline)

tests/MongoDB.Driver.Tests/UpdateDefinitionBuilderTests.cs

@@ -329,6 +329,16 @@ public void Mul_Typed()
             Assert(subject.Mul("Age", 2), "{$mul: {age: 2}}");
         }
 
+        [Fact]
+        public void Pipeline()
+        {
+            var subject = CreateSubject<BsonDocument>();
+            var pipeline = new EmptyPipelineDefinition<BsonDocument>()
+                .AppendStage<BsonDocument, BsonDocument, BsonDocument>("{ $addFields : { x : 2 } }");
+
+            Assert(subject.Pipeline(pipeline), new[] { "{ \"$addFields\" : { \"x\" : 2 } }" });
+        }
+
         [Fact]
         public void PopFirst()
         {
@@ -639,11 +649,19 @@ public void Unset_Typed()
 
         private void Assert<TDocument>(UpdateDefinition<TDocument> update, BsonDocument expected)
         {
-            var renderedUpdate = Render(update);
+            var renderedUpdate = Render(update).AsBsonDocument;
 
             renderedUpdate.Should().Be(expected);
         }
 
+        private void Assert<TDocument>(UpdateDefinition<TDocument> update, string[] expectedArrayItems)
+        {
+            var renderedUpdate = Render(update).AsBsonArray;
+
+            var bsonArray = new BsonArray(expectedArrayItems.Select(BsonDocument.Parse));
+            renderedUpdate.Should().Be(bsonArray);
+        }
+
         private void Assert<TDocument>(UpdateDefinition<TDocument> update, string expected)
         {
             Assert(update, BsonDocument.Parse(expected));
@@ -661,10 +679,10 @@ private UpdateDefinitionBuilder<TDocument> CreateSubject<TDocument>()
             return new UpdateDefinitionBuilder<TDocument>();
         }
 
-        private BsonDocument Render<TDocument>(UpdateDefinition<TDocument> update)
+        private BsonValue Render<TDocument>(UpdateDefinition<TDocument> update)
         {
             var documentSerializer = BsonSerializer.SerializerRegistry.GetSerializer<TDocument>();
-            return update.Render(documentSerializer, BsonSerializer.SerializerRegistry).AsBsonDocument;
+            return update.Render(documentSerializer, BsonSerializer.SerializerRegistry);
         }
 
         private class Person