@mongoosejs/studio 0.5.13 → 0.5.14

package/backend/actions/Model/getDocument.js

@@ -4,6 +4,7 @@ const Archetype = require('archetype');
 const getRefFromSchemaType = require('../../helpers/getRefFromSchemaType');
 const removeSpecifiedPaths = require('../../helpers/removeSpecifiedPaths');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetDocumentParams = new Archetype({
   model: {
@@ -19,7 +20,7 @@ const GetDocumentParams = new Archetype({
   }
 }).compile('GetDocumentParams');
 
-module.exports = ({ db }) => async function getDocument(params) {
+module.exports = ({ db, options }) => async function getDocument(params) {
   const { model, documentId, roles } = new GetDocumentParams(params);
 
   await authorize('Model.getDocument', roles);
@@ -31,7 +32,7 @@ module.exports = ({ db }) => async function getDocument(params) {
 
   const doc = await Model.
     findById(documentId).
-    setOptions({ sanitizeFilter: true }).
+    setOptions(omitNullish({ sanitizeFilter: true, maxTimeMS: options?.maxTimeMS })).
     orFail();
   const schemaPaths = {};
   for (const path of Object.keys(Model.schema.paths)) {

package/backend/actions/Model/getDocuments.js

@@ -7,6 +7,7 @@ const getRefFromSchemaType = require('../../helpers/getRefFromSchemaType');
 const getSuggestedProjection = require('../../helpers/getSuggestedProjection');
 const parseProjectionParam = require('../../helpers/parseProjectionParam');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetDocumentsParams = new Archetype({
   model: {
@@ -40,7 +41,7 @@ const GetDocumentsParams = new Archetype({
   }
 }).compile('GetDocumentsParams');
 
-module.exports = ({ db }) => async function getDocuments(params) {
+module.exports = ({ db, options }) => async function getDocuments(params) {
   params = new GetDocumentsParams(params);
   const { roles } = params;
   await authorize('Model.getDocuments', roles);
@@ -75,6 +76,7 @@ module.exports = ({ db }) => async function getDocuments(params) {
   if (projection != null) {
     query = query.select(projection);
   }
+  query.setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS }));
   const cursor = await query.cursor();
   const docs = [];
   for (let doc = await cursor.next(); doc != null; doc = await cursor.next()) {
@@ -106,9 +108,16 @@ module.exports = ({ db }) => async function getDocuments(params) {
   }
   removeSpecifiedPaths(schemaPaths, '.$*');
 
-  const numDocuments = parsedFilter == null ?
-    await Model.estimatedDocumentCount() :
-    await Model.countDocuments(filter);
+  const countQuery = parsedFilter == null ?
+    Model.estimatedDocumentCount() :
+    Model.countDocuments(filter);
+  let numDocuments = null;
+  let numDocsError = null;
+  try {
+    numDocuments = await countQuery.setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS }));
+  } catch (err) {
+    numDocsError = err?.message;
+  }
 
   const suggestedFields = getSuggestedProjection(Model);
 
@@ -116,6 +125,7 @@ module.exports = ({ db }) => async function getDocuments(params) {
     docs: docs.map(doc => doc.toJSON({ virtuals: false, getters: false, transform: false })),
     schemaPaths,
     suggestedFields,
-    numDocs: numDocuments
+    numDocs: numDocuments,
+    numDocsError
   };
 };

package/backend/actions/Model/getDocumentsStream.js

@@ -7,6 +7,7 @@ const getRefFromSchemaType = require('../../helpers/getRefFromSchemaType');
 const getSuggestedProjection = require('../../helpers/getSuggestedProjection');
 const parseProjectionParam = require('../../helpers/parseProjectionParam');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetDocumentsParams = new Archetype({
   model: {
@@ -40,7 +41,7 @@ const GetDocumentsParams = new Archetype({
   }
 }).compile('GetDocumentsParams');
 
-module.exports = ({ db }) => async function* getDocumentsStream(params) {
+module.exports = ({ db, options }) => async function* getDocumentsStream(params) {
   params = new GetDocumentsParams(params);
   const { roles } = params;
   await authorize('Model.getDocumentsStream', roles);
@@ -75,6 +76,7 @@ module.exports = ({ db }) => async function* getDocumentsStream(params) {
   if (projection != null) {
     query = query.select(projection);
   }
+  query.setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS }));
 
   const schemaPaths = {};
   for (const path of Object.keys(Model.schema.paths)) {
@@ -106,32 +108,36 @@ module.exports = ({ db }) => async function* getDocumentsStream(params) {
   yield { schemaPaths, suggestedFields };
 
   // Start counting documents in parallel with streaming documents
-  const numDocsPromise = (parsedFilter == null)
-    ? Model.estimatedDocumentCount().exec()
-    : Model.countDocuments(filter).exec();
-
-  const cursor = await query.cursor();
+  const countQuery = (parsedFilter == null)
+    ? Model.estimatedDocumentCount()
+    : Model.countDocuments(filter);
+  const numDocsPromise = countQuery.
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS })).
+    exec();
 
   let numDocsYielded = false;
   let numDocumentsPromiseResolved = false;
   let numDocumentsValue;
   let numDocumentsError;
 
-  try {
-    // Start listening for numDocsPromise resolution
-    numDocsPromise.then(num => {
-      numDocumentsPromiseResolved = true;
-      numDocumentsValue = num;
-    }).catch(err => {
-      numDocumentsPromiseResolved = true;
-      numDocumentsError = err;
-    });
+  // Attach both handlers immediately so a fast count failure cannot become
+  // an unhandled rejection while the document cursor is still opening.
+  numDocsPromise.then(num => {
+    numDocumentsPromiseResolved = true;
+    numDocumentsValue = num;
+  }).catch(err => {
+    numDocumentsPromiseResolved = true;
+    numDocumentsError = err;
+  });
 
+  const cursor = await query.cursor();
+
+  try {
     for (let doc = await cursor.next(); doc != null; doc = await cursor.next()) {
       // If numDocsPromise has resolved and not yet yielded, yield it first
       if (numDocumentsPromiseResolved && !numDocsYielded) {
         if (numDocumentsError) {
-          yield { error: numDocumentsError };
+          yield { numDocsError: numDocumentsError.message };
         } else {
           yield { numDocs: numDocumentsValue };
         }
@@ -142,8 +148,12 @@ module.exports = ({ db }) => async function* getDocumentsStream(params) {
 
     // If numDocsPromise hasn't resolved yet, wait for it and yield
     if (!numDocsYielded) {
-      const numDocuments = await numDocsPromise;
-      yield { numDocs: numDocuments };
+      try {
+        const numDocuments = await numDocsPromise;
+        yield { numDocs: numDocuments };
+      } catch (err) {
+        yield { numDocsError: err.message };
+      }
     }
   } finally {
     await cursor.close();

package/backend/actions/Model/getEstimatedDocumentCounts.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetEstimatedDocumentCountsParams = new Archetype({
   roles: {
@@ -9,7 +10,7 @@ const GetEstimatedDocumentCountsParams = new Archetype({
   }
 }).compile('GetEstimatedDocumentCountsParams');
 
-module.exports = ({ db }) => async function getEstimatedDocumentCounts(params) {
+module.exports = ({ db, options }) => async function getEstimatedDocumentCounts(params) {
   const { roles } = new GetEstimatedDocumentCountsParams(params);
   await authorize('Model.getEstimatedDocumentCounts', roles);
 
@@ -20,7 +21,9 @@ module.exports = ({ db }) => async function getEstimatedDocumentCounts(params) {
   const results = await Promise.allSettled(
     modelNames.map(name => {
       const Model = db.models[name];
-      return Model.estimatedDocumentCount().exec();
+      return Model.estimatedDocumentCount().
+        setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS })).
+        exec();
     })
   );
 

package/backend/actions/Model/updateDocument.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const UpdateDocumentsParams = new Archetype({
   model: {
@@ -24,7 +25,7 @@ const UpdateDocumentsParams = new Archetype({
   }
 }).compile('UpdateDocumentsParams');
 
-module.exports = ({ db }) => async function updateDocument(params) {
+module.exports = ({ db, options }) => async function updateDocument(params) {
   const { model, _id, update, unset, roles } = new UpdateDocumentsParams(params);
 
   await authorize('Model.updateDocument', roles);
@@ -58,6 +59,7 @@ module.exports = ({ db }) => async function updateDocument(params) {
   }
 
   const doc = await Model.
-    findByIdAndUpdate(_id, updateOperation, { sanitizeFilter: true, returnDocument: 'after', overwriteImmutable: true, runValidators: false });
+    findByIdAndUpdate(_id, updateOperation, { sanitizeFilter: true, returnDocument: 'after', overwriteImmutable: true, runValidators: false }).
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS }));
   return { doc };
 };

package/backend/actions/Model/updateDocuments.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 
 const UpdateDocumentsParams = new Archetype({
   model: {
@@ -21,7 +22,7 @@ const UpdateDocumentsParams = new Archetype({
   }
 }).compile('UpdateDocumentsParams');
 
-module.exports = ({ db }) => async function updateDocuments(params) {
+module.exports = ({ db, options }) => async function updateDocuments(params) {
   const { model, _id, update, roles } = new UpdateDocumentsParams(params);
 
   await authorize('Model.updateDocuments', roles);
@@ -60,7 +61,8 @@ module.exports = ({ db }) => async function updateDocuments(params) {
   }
 
   const result = await Model.
-    updateMany({ _id: { $in: _id } }, updateOperation, { overwriteImmutable: true, runValidators: false });
+    updateMany({ _id: { $in: _id } }, updateOperation, { overwriteImmutable: true, runValidators: false }).
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS }));
 
   return { result };
 };

package/backend/actions/Model/validateDocument.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const authorize = require('../../authorize');
+const omitNullish = require('../../helpers/omitNullish');
 const validateDocumentWithTimeout = require('../../helpers/validateDocumentWithTimeout');
 
 const ValidateDocumentParams = new Archetype({
@@ -18,7 +19,7 @@ const ValidateDocumentParams = new Archetype({
   }
 }).compile('ValidateDocumentParams');
 
-module.exports = ({ db }) => async function validateDocument(params) {
+module.exports = ({ db, options }) => async function validateDocument(params) {
   const { model, documentId, roles } = new ValidateDocumentParams(params);
 
   await authorize('Model.validateDocument', roles);
@@ -28,7 +29,9 @@ module.exports = ({ db }) => async function validateDocument(params) {
     throw new Error(`Model ${model} not found`);
   }
 
-  const doc = await Model.findById(documentId).setOptions({ sanitizeFilter: true }).orFail();
+  const doc = await Model.findById(documentId).
+    setOptions(omitNullish({ sanitizeFilter: true, maxTimeMS: options?.maxTimeMS })).
+    orFail();
   return {
     result: await validateDocumentWithTimeout(doc)
   };

package/backend/actions/Task/cancelTask.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const mongoose = require('mongoose');
+const omitNullish = require('../../helpers/omitNullish');
 
 const CancelTaskParams = new Archetype({
   taskId: {
@@ -10,15 +11,17 @@ const CancelTaskParams = new Archetype({
   }
 }).compile('CancelTaskParams');
 
-module.exports = ({ db }) => async function cancelTask(params) {
+module.exports = ({ db, options }) => async function cancelTask(params) {
   params = new CancelTaskParams(params);
   const { taskId } = params;
   const { Task } = db.models;
 
-  const task = await Task.findOne({ _id: taskId }).orFail();
+  const task = await Task.findOne({ _id: taskId }).
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS })).
+    orFail();
 
   const cancelledTask = await Task.cancelTask({ _id: taskId });
   return {
     task: cancelledTask
   };
-};
+};

package/backend/actions/Task/getTaskOverview.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const escape = require('regexp.escape');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetTaskOverviewParams = new Archetype({
   start: { $type: Date },
@@ -34,7 +35,7 @@ function buildMatch(params) {
   return match;
 }
 
-module.exports = ({ db }) => async function getTaskOverview(params) {
+module.exports = ({ db, options }) => async function getTaskOverview(params) {
   params = new GetTaskOverviewParams(params);
   if (typeof params.status === 'string') params.status = params.status.trim();
   if (typeof params.name === 'string') params.name = params.name.trim();
@@ -93,7 +94,7 @@ module.exports = ({ db }) => async function getTaskOverview(params) {
     }
   ];
 
-  const [result] = await Task.aggregate(pipeline);
+  const [result] = await Task.aggregate(pipeline).option(omitNullish({ maxTimeMS: options?.maxTimeMS }));
 
   return {
     statusCounts: result.statusCounts?.[0] ?? {},

package/backend/actions/Task/getTasks.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const escape = require('regexp.escape');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetTasksParams = new Archetype({
   start: { $type: Date },
@@ -55,7 +56,7 @@ const TASK_PROJECT_STAGE = {
   parameters: '$payload'
 };
 
-module.exports = ({ db }) => async function getTasks(params) {
+module.exports = ({ db, options }) => async function getTasks(params) {
   params = new GetTasksParams(params);
   if (typeof params.status === 'string') params.status = params.status.trim();
   if (typeof params.name === 'string') params.name = params.name.trim();
@@ -96,7 +97,7 @@ module.exports = ({ db }) => async function getTasks(params) {
     }
   ];
 
-  const [result] = await Task.aggregate(pipeline);
+  const [result] = await Task.aggregate(pipeline).option(omitNullish({ maxTimeMS: options?.maxTimeMS }));
   const tasks = result.tasks || [];
   const numDocs = (result.count && result.count[0] && result.count[0].total) || 0;
   const statusCounts = result.statusCounts?.[0] ?? {};

package/backend/actions/Task/getTasksOverTime.js

@@ -1,6 +1,7 @@
 'use strict';
 
 const Archetype = require('archetype');
+const omitNullish = require('../../helpers/omitNullish');
 
 const GetTasksOverTimeParams = new Archetype({
   start: { $type: Date },
@@ -10,7 +11,7 @@ const GetTasksOverTimeParams = new Archetype({
 
 const TRACKED_STATUSES = ['succeeded', 'failed', 'cancelled'];
 
-module.exports = ({ db }) => async function getTasksOverTime(params) {
+module.exports = ({ db, options }) => async function getTasksOverTime(params) {
   params = new GetTasksOverTimeParams(params);
   const { Task } = db.models;
   const { start, end, bucketSizeMs } = params;
@@ -54,7 +55,7 @@ module.exports = ({ db }) => async function getTasksOverTime(params) {
     { $sort: { _id: 1 } }
   ];
 
-  const results = await Task.aggregate(pipeline);
+  const results = await Task.aggregate(pipeline).option(omitNullish({ maxTimeMS: options?.maxTimeMS }));
 
   return results.map(r => {
     const bucket = { timestamp: r._id, succeeded: 0, failed: 0, cancelled: 0 };

package/backend/actions/Task/rescheduleTask.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const mongoose = require('mongoose');
+const omitNullish = require('../../helpers/omitNullish');
 
 const RescheduleTaskParams = new Archetype({
   taskId: {
@@ -14,12 +15,14 @@ const RescheduleTaskParams = new Archetype({
   }
 }).compile('RescheduleTaskParams');
 
-module.exports = ({ db }) => async function rescheduleTask(params) {
+module.exports = ({ db, options }) => async function rescheduleTask(params) {
   params = new RescheduleTaskParams(params);
   const { taskId, scheduledAt } = params;
   const { Task } = db.models;
 
-  const task = await Task.findOne({ _id: taskId }).orFail();
+  const task = await Task.findOne({ _id: taskId }).
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS })).
+    orFail();
 
   if (scheduledAt < Date.now()) {
     throw new Error('Cannot reschedule a task for the past');
@@ -36,4 +39,4 @@ module.exports = ({ db }) => async function rescheduleTask(params) {
   return {
     task
   };
-};
+};

package/backend/actions/Task/runTask.js

@@ -2,6 +2,7 @@
 
 const Archetype = require('archetype');
 const mongoose = require('mongoose');
+const omitNullish = require('../../helpers/omitNullish');
 
 const RunTaskParams = new Archetype({
   taskId: {
@@ -10,16 +11,18 @@ const RunTaskParams = new Archetype({
   }
 }).compile('RunTaskParams');
 
-module.exports = ({ db }) => async function runTask(params) {
+module.exports = ({ db, options }) => async function runTask(params) {
   params = new RunTaskParams(params);
   const { taskId } = params;
   const { Task } = db.models;
 
-  const task = await Task.findOne({ _id: taskId }).orFail();
+  const task = await Task.findOne({ _id: taskId }).
+    setOptions(omitNullish({ maxTimeMS: options?.maxTimeMS })).
+    orFail();
 
   const executedTask = await Task.execute(task);
  
   return {
     task: executedTask
   };
-};
+};

package/backend/helpers/omitNullish.js

@@ -0,0 +1,5 @@
+'use strict';
+
+module.exports = function omitNullish(obj) {
+  return Object.fromEntries(Object.entries(obj).filter(([, value]) => value != null));
+};

package/backend/index.js

@@ -16,10 +16,10 @@ module.exports = function backend(db, studioConnection, options) {
   }
 
   studioConnection = studioConnection ?? db;
-  const Dashboard = studioConnection.model('__Studio_Dashboard', dashboardSchema, 'studio__dashboards');
-  const DashboardResult = studioConnection.model('__Studio_DashboardResult', dashboardResultSchema, 'studio__dashboardResults');
-  const ChatMessage = studioConnection.model('__Studio_ChatMessage', chatMessageSchema, 'studio__chatMessages');
-  const ChatThread = studioConnection.model('__Studio_ChatThread', chatThreadSchema, 'studio__chatThreads');
+  studioConnection.model('__Studio_Dashboard', dashboardSchema, 'studio__dashboards');
+  studioConnection.model('__Studio_DashboardResult', dashboardResultSchema, 'studio__dashboardResults');
+  studioConnection.model('__Studio_ChatMessage', chatMessageSchema, 'studio__chatMessages');
+  studioConnection.model('__Studio_ChatThread', chatThreadSchema, 'studio__chatThreads');
 
   let changeStream = null;
   if (options?.changeStream) {

package/backend/sandbox/createSandbox.js

@@ -5,7 +5,7 @@ const mongoose = require('mongoose');
 const vm = require('vm');
 const { createScriptDb } = require('./createScriptDb');
 
-module.exports = function createSandbox({ db }) {
+module.exports = function createSandbox({ db, maxTimeMS }) {
   const logs = [];
 
   // db must be a connection, not a Mongoose instance
@@ -13,7 +13,7 @@ module.exports = function createSandbox({ db }) {
     db = db.connection;
   }
 
-  const scriptDb = createScriptDb(db);
+  const scriptDb = createScriptDb(db, { maxTimeMS });
   if (!scriptDb.db.Types) {
     scriptDb.db.Types = mongoose.Types;
   }

package/backend/sandbox/createScriptDb.js

@@ -1,8 +1,11 @@
 'use strict';
 
+const omitNullish = require('../helpers/omitNullish');
+
 const wrappedCollections = new WeakSet();
 
 const collectionMethodOptionsIndex = new Map([
+  ['aggregate', 1],
   ['bulkWrite', 1],
   ['countDocuments', 1],
   ['createIndex', 1],
@@ -21,26 +24,34 @@ const collectionMethodOptionsIndex = new Map([
   ['findOneAndUpdate', 2],
   ['insertMany', 1],
   ['insertOne', 1],
+  ['listIndexes', 0],
   ['rename', 1],
   ['replaceOne', 2],
   ['updateMany', 2],
   ['updateOne', 2]
 ]);
 
-function createScriptDb(db) {
+function createScriptDb(db, options = {}) {
   const sourceConnection = db;
 
   const scriptConnection = sourceConnection.useDb(sourceConnection.name, { useCache: false });
+  // Mongoose useDb() shares the source connection's options object. Copy it
+  // before setting sandbox-only options so the application connection remains untouched.
+  scriptConnection.options = { ...scriptConnection.options };
   scriptConnection.config = {
     ...scriptConnection.config,
     autoCreate: false,
     autoIndex: false
   };
+  if (options.maxTimeMS != null) {
+    scriptConnection.set('maxTimeMS', options.maxTimeMS);
+  }
 
   let dryRunSession = null;
   const getSession = () => dryRunSession;
-  cloneModels(sourceConnection, scriptConnection, getSession);
-  wrapCollectionAccessors(scriptConnection, getSession);
+  const getMaxTimeMS = () => options.maxTimeMS;
+  cloneModels(sourceConnection, scriptConnection, getSession, getMaxTimeMS);
+  wrapCollectionAccessors(scriptConnection, getSession, getMaxTimeMS);
 
   const originalDebug = sourceConnection.options?.debug;
   scriptConnection.set('debug', function() {
@@ -68,14 +79,14 @@ function createScriptDb(db) {
   };
 }
 
-function cloneModels(sourceConnection, scriptConnection, getSession) {
+function cloneModels(sourceConnection, scriptConnection, getSession, getMaxTimeMS) {
   for (const Model of Object.values(sourceConnection.models ?? {})) {
     if (Model?.schema == null || typeof scriptConnection.model !== 'function') {
       continue;
     }
 
     const ClonedModel = scriptConnection.model(Model.modelName, Model.schema, getCollectionName(Model));
-    wrapModelCollections(ClonedModel, getSession);
+    wrapModelCollections(ClonedModel, getSession, getMaxTimeMS);
   }
 }
 
@@ -83,13 +94,13 @@ function getCollectionName(Model) {
   return Model.collection?.collectionName ?? Model.collection?.name ?? Model.$__collection?.collectionName;
 }
 
-function wrapModelCollections(Model, getSession) {
+function wrapModelCollections(Model, getSession, getMaxTimeMS) {
   const collection = Model?.collection ?? Model?.$__collection;
   if (collection == null) {
     return;
   }
 
-  const wrappedCollection = wrapCollection(collection, getSession);
+  const wrappedCollection = wrapCollection(collection, getSession, getMaxTimeMS);
   Model.collection = wrappedCollection;
   Model.$__collection = wrappedCollection;
 
@@ -108,7 +119,7 @@ function setModelCollectionSymbols(target, collection) {
   }
 }
 
-function wrapCollectionAccessors(scriptConnection, getSession) {
+function wrapCollectionAccessors(scriptConnection, getSession, getMaxTimeMS) {
   const wrappedCollectionCache = new WeakMap();
 
   const nativeDb = scriptConnection.db;
@@ -125,19 +136,19 @@ function wrapCollectionAccessors(scriptConnection, getSession) {
           collection.collection = originalDbCollection.call(nativeDb, collection.name);
         }
       }
-      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession);
+      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession, getMaxTimeMS);
     };
   }
 
   if (originalDbCollection != null) {
     nativeDb.collection = function() {
       const collection = originalDbCollection.apply(this, arguments);
-      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession);
+      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession, getMaxTimeMS);
     };
   }
 }
 
-function getOrCreateWrappedCollection(cache, collection, getSession) {
+function getOrCreateWrappedCollection(cache, collection, getSession, getMaxTimeMS) {
   if (collection == null || (typeof collection !== 'object' && typeof collection !== 'function')) {
     return collection;
   }
@@ -147,12 +158,12 @@ function getOrCreateWrappedCollection(cache, collection, getSession) {
   if (cache.has(collection)) {
     return cache.get(collection);
   }
-  const wrapped = wrapCollection(collection, getSession);
+  const wrapped = wrapCollection(collection, getSession, getMaxTimeMS);
   cache.set(collection, wrapped);
   return wrapped;
 }
 
-function wrapCollection(collection, getSession) {
+function wrapCollection(collection, getSession, getMaxTimeMS) {
   if (wrappedCollections.has(collection)) {
     return collection;
   }
@@ -165,15 +176,20 @@ function wrapCollection(collection, getSession) {
     }
 
     wrapped[methodName] = function() {
-      const args = addSessionOption(Array.from(arguments), optionsIndex, getSession());
+      const args = addOperationOptions(
+        Array.from(arguments),
+        optionsIndex,
+        getSession(),
+        getMaxTimeMS()
+      );
       return method.apply(collection, args);
     };
   }
   return wrapped;
 }
 
-function addSessionOption(args, optionsIndex, session) {
-  if (session == null) {
+function addOperationOptions(args, optionsIndex, session, maxTimeMS) {
+  if (session == null && maxTimeMS == null) {
     return args;
   }
 
@@ -184,14 +200,21 @@ function addSessionOption(args, optionsIndex, session) {
   const options = args[optionsIndex];
 
   if (options != null && typeof options !== 'object') {
-    throw new Error('Cannot run dry run on script where options arg is a non-object');
+    if (session != null) {
+      throw new Error('Cannot run dry run on script where options arg is a non-object');
+    }
+    throw new Error('Cannot apply maxTimeMS to script where options arg is a non-object');
   }
 
-  if (options?.session != null) {
+  if (session != null && options?.session != null) {
     throw new Error('Cannot run dry run on script that uses sessions');
   }
 
-  args[optionsIndex] = { ...(options ?? {}), session };
+  args[optionsIndex] = omitNullish({
+    ...(options ?? {}),
+    maxTimeMS: maxTimeMS ?? options?.maxTimeMS,
+    session: session ?? options?.session
+  });
   return args;
 }