feat(NODE-6421): add support for timeoutMS to explain helpers (#4268)

package.json

@@ -97,7 +97,7 @@
     "mocha": "^10.4.0",
     "mocha-sinon": "^2.1.2",
     "mongodb-client-encryption": "^6.1.0",
-    "mongodb-legacy": "^6.1.2",
+    "mongodb-legacy": "^6.1.3",
     "nyc": "^15.1.0",
     "prettier": "^3.3.3",
     "semver": "^7.6.3",

src/cursor/aggregation_cursor.ts

@@ -1,14 +1,19 @@
 import type { Document } from '../bson';
 import { MongoAPIError } from '../error';
-import type { ExplainCommandOptions, ExplainVerbosityLike } from '../explain';
+import {
+  Explain,
+  ExplainableCursor,
+  type ExplainCommandOptions,
+  type ExplainVerbosityLike,
+  validateExplainTimeoutOptions
+} from '../explain';
 import type { MongoClient } from '../mongo_client';
 import { AggregateOperation, type AggregateOptions } from '../operations/aggregate';
 import { executeOperation } from '../operations/execute_operation';
 import type { ClientSession } from '../sessions';
 import type { Sort } from '../sort';
 import { mergeOptions, type MongoDBNamespace } from '../utils';
 import {
-  AbstractCursor,
   type AbstractCursorOptions,
   CursorTimeoutMode,
   type InitialCursorResponse
@@ -24,7 +29,7 @@ export interface AggregationCursorOptions extends AbstractCursorOptions, Aggrega
  * or higher stream
  * @public
  */
-export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
+export class AggregationCursor<TSchema = any> extends ExplainableCursor<TSchema> {
   public readonly pipeline: Document[];
   /** @internal */
   private aggregateOptions: AggregateOptions;
@@ -65,26 +70,47 @@ export class AggregationCursor<TSchema = any> extends AbstractCursor<TSchema> {
 
   /** @internal */
   async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
-    const aggregateOperation = new AggregateOperation(this.namespace, this.pipeline, {
+    const options = {
       ...this.aggregateOptions,
       ...this.cursorOptions,
       session
-    });
+    };
+    try {
+      validateExplainTimeoutOptions(options, Explain.fromOptions(options));
+    } catch {
+      throw new MongoAPIError(
+        'timeoutMS cannot be used with explain when explain is specified in aggregateOptions'
+      );
+    }
+
+    const aggregateOperation = new AggregateOperation(this.namespace, this.pipeline, options);
 
     const response = await executeOperation(this.client, aggregateOperation, this.timeoutContext);
 
     return { server: aggregateOperation.server, session, response };
   }
 
   /** Execute the explain for the cursor */
-  async explain(verbosity?: ExplainVerbosityLike | ExplainCommandOptions): Promise<Document> {
+  async explain(): Promise<Document>;
+  async explain(verbosity: ExplainVerbosityLike | ExplainCommandOptions): Promise<Document>;
+  async explain(options: { timeoutMS?: number }): Promise<Document>;
+  async explain(
+    verbosity: ExplainVerbosityLike | ExplainCommandOptions,
+    options: { timeoutMS?: number }
+  ): Promise<Document>;
+  async explain(
+    verbosity?: ExplainVerbosityLike | ExplainCommandOptions | { timeoutMS?: number },
+    options?: { timeoutMS?: number }
+  ): Promise<Document> {
+    const { explain, timeout } = this.resolveExplainTimeoutOptions(verbosity, options);
     return (
       await executeOperation(
         this.client,
         new AggregateOperation(this.namespace, this.pipeline, {
           ...this.aggregateOptions, // NOTE: order matters here, we may need to refine this
           ...this.cursorOptions,
-          explain: verbosity ?? true
+          ...timeout,
+          explain: explain ?? true
         })
       )
     ).shift(this.deserializationOptions);

src/cursor/find_cursor.ts

@@ -1,7 +1,13 @@
 import { type Document } from '../bson';
 import { CursorResponse } from '../cmap/wire_protocol/responses';
-import { MongoInvalidArgumentError, MongoTailableCursorError } from '../error';
-import { type ExplainCommandOptions, type ExplainVerbosityLike } from '../explain';
+import { MongoAPIError, MongoInvalidArgumentError, MongoTailableCursorError } from '../error';
+import {
+  Explain,
+  ExplainableCursor,
+  type ExplainCommandOptions,
+  type ExplainVerbosityLike,
+  validateExplainTimeoutOptions
+} from '../explain';
 import type { MongoClient } from '../mongo_client';
 import type { CollationOptions } from '../operations/command';
 import { CountOperation, type CountOptions } from '../operations/count';
@@ -11,7 +17,7 @@ import type { Hint } from '../operations/operation';
 import type { ClientSession } from '../sessions';
 import { formatSort, type Sort, type SortDirection } from '../sort';
 import { emitWarningOnce, mergeOptions, type MongoDBNamespace, squashError } from '../utils';
-import { AbstractCursor, type InitialCursorResponse } from './abstract_cursor';
+import { type InitialCursorResponse } from './abstract_cursor';
 
 /** @public Flags allowed for cursor */
 export const FLAGS = [
@@ -24,7 +30,7 @@ export const FLAGS = [
 ] as const;
 
 /** @public */
-export class FindCursor<TSchema = any> extends AbstractCursor<TSchema> {
+export class FindCursor<TSchema = any> extends ExplainableCursor<TSchema> {
   /** @internal */
   private cursorFilter: Document;
   /** @internal */
@@ -63,11 +69,21 @@ export class FindCursor<TSchema = any> extends AbstractCursor<TSchema> {
 
   /** @internal */
   async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
-    const findOperation = new FindOperation(this.namespace, this.cursorFilter, {
+    const options = {
       ...this.findOptions, // NOTE: order matters here, we may need to refine this
       ...this.cursorOptions,
       session
-    });
+    };
+
+    try {
+      validateExplainTimeoutOptions(options, Explain.fromOptions(options));
+    } catch {
+      throw new MongoAPIError(
+        'timeoutMS cannot be used with explain when explain is specified in findOptions'
+      );
+    }
+
+    const findOperation = new FindOperation(this.namespace, this.cursorFilter, options);
 
     const response = await executeOperation(this.client, findOperation, this.timeoutContext);
 
@@ -133,14 +149,27 @@ export class FindCursor<TSchema = any> extends AbstractCursor<TSchema> {
   }
 
   /** Execute the explain for the cursor */
-  async explain(verbosity?: ExplainVerbosityLike | ExplainCommandOptions): Promise<Document> {
+  async explain(): Promise<Document>;
+  async explain(verbosity: ExplainVerbosityLike | ExplainCommandOptions): Promise<Document>;
+  async explain(options: { timeoutMS?: number }): Promise<Document>;
+  async explain(
+    verbosity: ExplainVerbosityLike | ExplainCommandOptions,
+    options: { timeoutMS?: number }
+  ): Promise<Document>;
+  async explain(
+    verbosity?: ExplainVerbosityLike | ExplainCommandOptions | { timeoutMS?: number },
+    options?: { timeoutMS?: number }
+  ): Promise<Document> {
+    const { explain, timeout } = this.resolveExplainTimeoutOptions(verbosity, options);
+
     return (
       await executeOperation(
         this.client,
         new FindOperation(this.namespace, this.cursorFilter, {
           ...this.findOptions, // NOTE: order matters here, we may need to refine this
           ...this.cursorOptions,
-          explain: verbosity ?? true
+          ...timeout,
+          explain: explain ?? true
         })
       )
     ).shift(this.deserializationOptions);

src/explain.ts

@@ -1,3 +1,7 @@
+import { type Document } from './bson';
+import { AbstractCursor } from './cursor/abstract_cursor';
+import { MongoAPIError } from './error';
+
 /** @public */
 export const ExplainVerbosity = Object.freeze({
   queryPlanner: 'queryPlanner',
@@ -86,3 +90,84 @@ export class Explain {
     return new Explain(verbosity, maxTimeMS);
   }
 }
+
+export function validateExplainTimeoutOptions(options: Document, explain?: Explain) {
+  const { maxTimeMS, timeoutMS } = options;
+  if (timeoutMS != null && (maxTimeMS != null || explain?.maxTimeMS != null)) {
+    throw new MongoAPIError('Cannot use maxTimeMS with timeoutMS for explain commands.');
+  }
+}
+
+/**
+ * Applies an explain to a given command.
+ * @internal
+ *
+ * @param command - the command on which to apply the explain
+ * @param options - the options containing the explain verbosity
+ */
+export function decorateWithExplain(
+  command: Document,
+  explain: Explain
+): {
+  explain: Document;
+  verbosity: ExplainVerbosity;
+  maxTimeMS?: number;
+} {
+  type ExplainCommand = ReturnType<typeof decorateWithExplain>;
+  const { verbosity, maxTimeMS } = explain;
+  const baseCommand: ExplainCommand = { explain: command, verbosity };
+
+  if (typeof maxTimeMS === 'number') {
+    baseCommand.maxTimeMS = maxTimeMS;
+  }
+
+  return baseCommand;
+}
+
+/**
+ * @public
+ *
+ * A base class for any cursors that have `explain()` methods.
+ */
+export abstract class ExplainableCursor<TSchema> extends AbstractCursor<TSchema> {
+  /** Execute the explain for the cursor */
+  abstract explain(): Promise<Document>;
+  abstract explain(verbosity: ExplainVerbosityLike | ExplainCommandOptions): Promise<Document>;
+  abstract explain(options: { timeoutMS?: number }): Promise<Document>;
+  abstract explain(
+    verbosity: ExplainVerbosityLike | ExplainCommandOptions,
+    options: { timeoutMS?: number }
+  ): Promise<Document>;
+  abstract explain(
+    verbosity?: ExplainVerbosityLike | ExplainCommandOptions | { timeoutMS?: number },
+    options?: { timeoutMS?: number }
+  ): Promise<Document>;
+
+  protected resolveExplainTimeoutOptions(
+    verbosity?: ExplainVerbosityLike | ExplainCommandOptions | { timeoutMS?: number },
+    options?: { timeoutMS?: number }
+  ): { timeout?: { timeoutMS?: number }; explain?: ExplainVerbosityLike | ExplainCommandOptions } {
+    let explain: ExplainVerbosityLike | ExplainCommandOptions | undefined;
+    let timeout: { timeoutMS?: number } | undefined;
+
+    if (verbosity == null && options == null) {
+      explain = undefined;
+      timeout = undefined;
+    } else if (verbosity != null && options == null) {
+      explain =
+        typeof verbosity !== 'object'
+          ? verbosity
+          : 'verbosity' in verbosity
+            ? verbosity
+            : undefined;
+
+      timeout = typeof verbosity === 'object' && 'timeoutMS' in verbosity ? verbosity : undefined;
+    } else {
+      // @ts-expect-error TS isn't smart enough to determine that if both options are provided, the first is explain options
+      explain = verbosity;
+      timeout = options;
+    }
+
+    return { timeout, explain };
+  }
+}

src/index.ts

@@ -10,6 +10,7 @@ import { ListCollectionsCursor } from './cursor/list_collections_cursor';
 import { ListIndexesCursor } from './cursor/list_indexes_cursor';
 import type { RunCommandCursor } from './cursor/run_command_cursor';
 import { Db } from './db';
+import { ExplainableCursor } from './explain';
 import { GridFSBucket } from './gridfs';
 import { GridFSBucketReadStream } from './gridfs/download';
 import { GridFSBucketWriteStream } from './gridfs/upload';
@@ -91,6 +92,7 @@ export {
   ClientSession,
   Collection,
   Db,
+  ExplainableCursor,
   FindCursor,
   GridFSBucket,
   GridFSBucketReadStream,

src/operations/command.ts

@@ -1,19 +1,19 @@
 import type { BSONSerializeOptions, Document } from '../bson';
 import { type MongoDBResponseConstructor } from '../cmap/wire_protocol/responses';
 import { MongoInvalidArgumentError } from '../error';
-import { Explain, type ExplainOptions } from '../explain';
+import {
+  decorateWithExplain,
+  Explain,
+  type ExplainOptions,
+  validateExplainTimeoutOptions
+} from '../explain';
 import { ReadConcern } from '../read_concern';
 import type { ReadPreference } from '../read_preference';
 import type { Server } from '../sdam/server';
 import { MIN_SECONDARY_WRITE_WIRE_VERSION } from '../sdam/server_selection';
 import type { ClientSession } from '../sessions';
 import { type TimeoutContext } from '../timeout';
-import {
-  commandSupportsReadConcern,
-  decorateWithExplain,
-  maxWireVersion,
-  MongoDBNamespace
-} from '../utils';
+import { commandSupportsReadConcern, maxWireVersion, MongoDBNamespace } from '../utils';
 import { WriteConcern, type WriteConcernOptions } from '../write_concern';
 import type { ReadConcernLike } from './../read_concern';
 import { AbstractOperation, Aspect, type OperationOptions } from './operation';
@@ -97,6 +97,7 @@ export abstract class CommandOperation<T> extends AbstractOperation<T> {
 
     if (this.hasAspect(Aspect.EXPLAINABLE)) {
       this.explain = Explain.fromOptions(options);
+      validateExplainTimeoutOptions(this.options, this.explain);
     } else if (options?.explain != null) {
       throw new MongoInvalidArgumentError(`Option "explain" is not supported on this command`);
     }

src/operations/find.ts

@@ -2,13 +2,17 @@ import type { Document } from '../bson';
 import { CursorResponse, ExplainedCursorResponse } from '../cmap/wire_protocol/responses';
 import { type AbstractCursorOptions, type CursorTimeoutMode } from '../cursor/abstract_cursor';
 import { MongoInvalidArgumentError } from '../error';
-import { type ExplainOptions } from '../explain';
+import {
+  decorateWithExplain,
+  type ExplainOptions,
+  validateExplainTimeoutOptions
+} from '../explain';
 import { ReadConcern } from '../read_concern';
 import type { Server } from '../sdam/server';
 import type { ClientSession } from '../sessions';
 import { formatSort, type Sort } from '../sort';
 import { type TimeoutContext } from '../timeout';
-import { decorateWithExplain, type MongoDBNamespace, normalizeHintField } from '../utils';
+import { type MongoDBNamespace, normalizeHintField } from '../utils';
 import { type CollationOptions, CommandOperation, type CommandOperationOptions } from './command';
 import { Aspect, defineAspects, type Hint } from './operation';
 
@@ -119,6 +123,7 @@ export class FindOperation extends CommandOperation<CursorResponse> {
 
     let findCommand = makeFindCommand(this.ns, this.filter, options);
     if (this.explain) {
+      validateExplainTimeoutOptions(this.options, this.explain);
       findCommand = decorateWithExplain(findCommand, this.explain);
     }