From 8fca1aaaeb611e26745c837868b7e7f47b287ad5 Mon Sep 17 00:00:00 2001
From: Alex <alexandru.andries@isa.utm.md>
Date: Thu, 18 Apr 2024 18:12:16 +0300
Subject: [PATCH] feat(NODE-3639): add a general stage to the aggregation
 pipeline builder (#4079)

---
 src/cursor/aggregation_cursor.ts          | 64 +++++++++++----------
 test/integration/crud/aggregation.test.ts | 68 ++++++++++++++++++++++-
 2 files changed, 97 insertions(+), 35 deletions(-)

diff --git a/src/cursor/aggregation_cursor.ts b/src/cursor/aggregation_cursor.ts
index 7b16887f5a..cba77e9b52 100644
--- a/src/cursor/aggregation_cursor.ts
+++ b/src/cursor/aggregation_cursor.ts
@@ -86,33 +86,45 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
     );
   }
 
+  /** Add a stage to the aggregation pipeline
+   * @example
+   * ```
+   * const documents = await users.aggregate().addStage({ $match: { name: /Mike/ } }).toArray();
+   * ```
+   * @example
+   * ```
+   * const documents = await users.aggregate()
+   *   .addStage<{ name: string }>({ $project: { name: true } })
+   *   .toArray(); // type of documents is { name: string }[]
+   * ```
+   */
+  addStage(stage: Document): this;
+  addStage<T = Document>(stage: Document): AggregationCursor<T>;
+  addStage<T = Document>(stage: Document): AggregationCursor<T> {
+    assertUninitialized(this);
+    this[kPipeline].push(stage);
+    return this as unknown as AggregationCursor<T>;
+  }
+
   /** Add a group stage to the aggregation pipeline */
   group<T = TSchema>($group: Document): AggregationCursor<T>;
   group($group: Document): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $group });
-    return this;
+    return this.addStage({ $group });
   }
 
   /** Add a limit stage to the aggregation pipeline */
   limit($limit: number): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $limit });
-    return this;
+    return this.addStage({ $limit });
   }
 
   /** Add a match stage to the aggregation pipeline */
   match($match: Document): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $match });
-    return this;
+    return this.addStage({ $match });
   }
 
   /** Add an out stage to the aggregation pipeline */
   out($out: { db: string; coll: string } | string): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $out });
-    return this;
+    return this.addStage({ $out });
   }
 
   /**
@@ -157,50 +169,36 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
    * ```
    */
   project<T extends Document = Document>($project: Document): AggregationCursor<T> {
-    assertUninitialized(this);
-    this[kPipeline].push({ $project });
-    return this as unknown as AggregationCursor<T>;
+    return this.addStage<T>({ $project });
   }
 
   /** Add a lookup stage to the aggregation pipeline */
   lookup($lookup: Document): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $lookup });
-    return this;
+    return this.addStage({ $lookup });
   }
 
   /** Add a redact stage to the aggregation pipeline */
   redact($redact: Document): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $redact });
-    return this;
+    return this.addStage({ $redact });
   }
 
   /** Add a skip stage to the aggregation pipeline */
   skip($skip: number): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $skip });
-    return this;
+    return this.addStage({ $skip });
   }
 
   /** Add a sort stage to the aggregation pipeline */
   sort($sort: Sort): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $sort });
-    return this;
+    return this.addStage({ $sort });
   }
 
   /** Add a unwind stage to the aggregation pipeline */
   unwind($unwind: Document | string): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $unwind });
-    return this;
+    return this.addStage({ $unwind });
   }
 
   /** Add a geoNear stage to the aggregation pipeline */
   geoNear($geoNear: Document): this {
-    assertUninitialized(this);
-    this[kPipeline].push({ $geoNear });
-    return this;
+    return this.addStage({ $geoNear });
   }
 }
diff --git a/test/integration/crud/aggregation.test.ts b/test/integration/crud/aggregation.test.ts
index 5e91614941..a91571d127 100644
--- a/test/integration/crud/aggregation.test.ts
+++ b/test/integration/crud/aggregation.test.ts
@@ -1,10 +1,10 @@
 import { expect } from 'chai';
 
-import { MongoInvalidArgumentError } from '../../mongodb';
+import { type MongoClient, MongoInvalidArgumentError } from '../../mongodb';
 import { filterForCommands } from '../shared';
 
 describe('Aggregation', function () {
-  let client;
+  let client: MongoClient;
 
   beforeEach(async function () {
     client = this.configuration.newClient();
@@ -939,4 +939,68 @@ describe('Aggregation', function () {
         .finally(() => client.close());
     }
   });
+
+  it('should return identical results for array aggregations and builder aggregations', async function () {
+    const databaseName = this.configuration.db;
+    const db = client.db(databaseName);
+    const collection = db.collection(
+      'shouldReturnIdenticalResultsForArrayAggregationsAndBuilderAggregations'
+    );
+
+    const docs = [
+      {
+        title: 'this is my title',
+        author: 'bob',
+        posted: new Date(),
+        pageViews: 5,
+        tags: ['fun', 'good', 'fun'],
+        other: { foo: 5 },
+        comments: [
+          { author: 'joe', text: 'this is cool' },
+          { author: 'sam', text: 'this is bad' }
+        ]
+      }
+    ];
+
+    await collection.insertMany(docs, { writeConcern: { w: 1 } });
+    const arrayPipelineCursor = collection.aggregate([
+      {
+        $project: {
+          author: 1,
+          tags: 1
+        }
+      },
+      { $unwind: '$tags' },
+      {
+        $group: {
+          _id: { tags: '$tags' },
+          authors: { $addToSet: '$author' }
+        }
+      },
+      { $sort: { _id: -1 } }
+    ]);
+
+    const builderPipelineCursor = collection
+      .aggregate()
+      .project({ author: 1, tags: 1 })
+      .unwind('$tags')
+      .group({ _id: { tags: '$tags' }, authors: { $addToSet: '$author' } })
+      .sort({ _id: -1 });
+
+    const builderGenericStageCursor = collection
+      .aggregate()
+      .addStage({ $project: { author: 1, tags: 1 } })
+      .addStage({ $unwind: '$tags' })
+      .addStage({ $group: { _id: { tags: '$tags' }, authors: { $addToSet: '$author' } } })
+      .addStage({ $sort: { _id: -1 } });
+
+    const expectedResults = [
+      { _id: { tags: 'good' }, authors: ['bob'] },
+      { _id: { tags: 'fun' }, authors: ['bob'] }
+    ];
+
+    expect(await arrayPipelineCursor.toArray()).to.deep.equal(expectedResults);
+    expect(await builderPipelineCursor.toArray()).to.deep.equal(expectedResults);
+    expect(await builderGenericStageCursor.toArray()).to.deep.equal(expectedResults);
+  });
 });