Add $out and $merge validation

Author: W-A-James

src/cursor/aggregation_cursor.ts

@@ -1,4 +1,5 @@
 import type { Document } from '../bson';
+import { MongoAPIError } from '../error';
 import type { ExplainVerbosityLike } from '../explain';
 import type { MongoClient } from '../mongo_client';
 import { AggregateOperation, type AggregateOptions } from '../operations/aggregate';
@@ -8,7 +9,7 @@ import type { Sort } from '../sort';
 import type { MongoDBNamespace } from '../utils';
 import { mergeOptions } from '../utils';
 import type { AbstractCursorOptions, InitialCursorResponse } from './abstract_cursor';
-import { AbstractCursor } from './abstract_cursor';
+import { AbstractCursor, CursorTimeoutMode } from './abstract_cursor';
 
 /** @public */
 export interface AggregationCursorOptions extends AbstractCursorOptions, AggregateOptions {}
@@ -36,6 +37,15 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
 
     this.pipeline = pipeline;
     this.aggregateOptions = options;
+
+    if (
+      this.cursorOptions.timeoutMS != null &&
+      this.cursorOptions.timeoutMode === CursorTimeoutMode.ITERATION &&
+      this.pipeline.filter(stage => {
+        Object.hasOwn(stage, '$out') || Object.hasOwn(stage, '$merge');
+      }).length > 0
+    )
+      throw new MongoAPIError('Cannot use $out or $merge stage with ITERATION timeoutMode');
   }
 
   clone(): AggregationCursor<TSchema> {
@@ -93,6 +103,13 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
   addStage<T = Document>(stage: Document): AggregationCursor<T>;
   addStage<T = Document>(stage: Document): AggregationCursor<T> {
     this.throwIfInitialized();
+    if (
+      this.cursorOptions.timeoutMS != null &&
+      this.cursorOptions.timeoutMode === CursorTimeoutMode.ITERATION &&
+      (Object.hasOwn(stage, '$out') || Object.hasOwn(stage, '$merge'))
+    ) {
+      throw new MongoAPIError('Cannot use $out or $merge stage with ITERATION timeoutMode');
+    }
     this.pipeline.push(stage);
     return this as unknown as AggregationCursor<T>;
   }

test/unit/cursor/aggregation_cursor.test.ts

@@ -1,6 +1,11 @@
 import { expect } from 'chai';
 
-import { type AggregationCursor, MongoClient } from '../../mongodb';
+import {
+  type AggregationCursor,
+  CursorTimeoutMode,
+  MongoAPIError,
+  MongoClient
+} from '../../mongodb';
 
 describe('class AggregationCursor', () => {
   let client: MongoClient;
@@ -126,6 +131,38 @@ describe('class AggregationCursor', () => {
   });
 
   context('when addStage, bespoke stage methods, or array is used to construct pipeline', () => {
+    context('when CSOT is enabled', () => {
+      let aggregationCursor: AggregationCursor;
+      before(function () {
+        aggregationCursor = client
+          .db('test')
+          .collection('test')
+          .aggregate([], { timeoutMS: 100, timeoutMode: CursorTimeoutMode.ITERATION });
+      });
+
+      context('when a $out stage is add with .addStage()', () => {
+        it('throws a MongoAPIError', function () {
+          expect(() => {
+            aggregationCursor.addStage({ $out: 'test' });
+          }).to.throw(MongoAPIError);
+        });
+      });
+      context('when a $merge stage is add with .addStage()', () => {
+        it('throws a MongoAPIError', function () {
+          expect(() => {
+            aggregationCursor.addStage({ $merge: {} });
+          }).to.throw(MongoAPIError);
+        });
+      });
+      context('when a $out stage is add with .out()', () => {
+        it('throws a MongoAPIError', function () {
+          expect(() => {
+            aggregationCursor.out('test');
+          }).to.throw(MongoAPIError);
+        });
+      });
+    });
+
     it('sets deeply identical aggregations pipelines', () => {
       const collection = client.db().collection('test');
 
@@ -157,4 +194,29 @@ describe('class AggregationCursor', () => {
       expect(builderGenericStageCursor.pipeline).to.deep.equal(expectedPipeline);
     });
   });
+
+  describe('constructor()', () => {
+    context('when CSOT is enabled', () => {
+      let client: MongoClient;
+      before(function () {
+        client = new MongoClient('mongodb://iLoveJavascript', { timeoutMS: 100 });
+      });
+      context('when timeoutMode=ITERATION and a $out stage is provided', function () {
+        expect(() => {
+          client
+            .db('test')
+            .collection('test')
+            .aggregate([{ $out: 'test' }], { timeoutMode: CursorTimeoutMode.ITERATION });
+        }).to.throw(MongoAPIError);
+      });
+      context('when timeoutMode=ITERATION and a $merge stage is provided', function () {
+        expect(() => {
+          client
+            .db('test')
+            .collection('test')
+            .aggregate([{ $merge: 'test' }], { timeoutMode: CursorTimeoutMode.ITERATION });
+        }).to.throw(MongoAPIError);
+      });
+    });
+  });
 });