From bf52c7f14930e42de59429cfde29b8acf8d86f18 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Tue, 15 Oct 2024 11:33:42 +0200
Subject: [PATCH] add single stage support instead

---
 packages/shell-api/src/database.spec.ts | 19 +++++++++----------
 packages/shell-api/src/database.ts      | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/packages/shell-api/src/database.spec.ts b/packages/shell-api/src/database.spec.ts
index c9684bd0e..a78e97101 100644
--- a/packages/shell-api/src/database.spec.ts
+++ b/packages/shell-api/src/database.spec.ts
@@ -390,19 +390,18 @@ describe('Database', function () {
         serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
       });
 
-      it('throws if the given argument is not an array', async function () {
-        let caughtError: MongoshInvalidInputError | undefined;
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
-        await database.aggregate({} as any).catch((err) => {
-          caughtError = err;
-        });
-        expect(caughtError?.message).contains(
-          'Aggregate pipeline argument must be an array'
+      it('calls serviceProvider.aggregateDb with pipleline and options', async function () {
+        await database.aggregate([{ $piplelineStage: {} }], { options: true });
+
+        expect(serviceProvider.aggregateDb).to.have.been.calledWith(
+          database._name,
+          [{ $piplelineStage: {} }],
+          { options: true }
         );
       });
 
-      it('calls serviceProvider.aggregateDb with pipleline and options', async function () {
-        await database.aggregate([{ $piplelineStage: {} }], { options: true });
+      it('supports a single aggregation stage', async function () {
+        await database.aggregate({ $piplelineStage: {} }, { options: true });
 
         expect(serviceProvider.aggregateDb).to.have.been.calledWith(
           database._name,
diff --git a/packages/shell-api/src/database.ts b/packages/shell-api/src/database.ts
index e038999d7..f129bac10 100644
--- a/packages/shell-api/src/database.ts
+++ b/packages/shell-api/src/database.ts
@@ -412,6 +412,14 @@ export default class Database extends ShellApiWithMongoClass {
     return await this._runAdminCommand(cmd, {});
   }
 
+  async aggregate(
+    singleStage: Document,
+    options?: Document
+  ): Promise<AggregationCursor>;
+  async aggregate(
+    pipeline: Document[],
+    options?: Document
+  ): Promise<AggregationCursor>;
   /**
    * Run an aggregation against the db.
    *
@@ -423,7 +431,7 @@ export default class Database extends ShellApiWithMongoClass {
   @returnType('AggregationCursor')
   @apiVersions([1])
   async aggregate(
-    pipeline: Document[],
+    pipelineOrSingleStage: Document | Document[],
     options?: Document
   ): Promise<AggregationCursor> {
     if ('background' in (options ?? {})) {
@@ -431,6 +439,10 @@ export default class Database extends ShellApiWithMongoClass {
         aggregateBackgroundOptionNotSupportedHelp
       );
     }
+    const pipeline: Document[] = Array.isArray(pipelineOrSingleStage)
+      ? pipelineOrSingleStage
+      : [pipelineOrSingleStage];
+
     assertArgsDefinedType([pipeline], [true], 'Database.aggregate');
 
     if (!Array.isArray(pipeline)) {