@mongoosejs/studio 0.4.1 → 0.5.1

package/backend/actions/ChatMessage/executeScript.js

@@ -2,9 +2,8 @@
 
 const Archetype = require('archetype');
 const authorize = require('../../authorize');
-const MongooseStudioChartColors = require('../../constants/mongooseStudioChartColors');
+const createSandbox = require('../../sandbox/createSandbox');
 const mongoose = require('mongoose');
-const vm = require('vm');
 
 const ExecuteScriptParams = new Archetype({
   initiatedById: {
@@ -17,13 +16,16 @@ const ExecuteScriptParams = new Archetype({
   script: {
     $type: 'string'
   },
+  dryRun: {
+    $type: 'boolean'
+  },
   roles: {
     $type: ['string']
   }
 }).compile('ExecuteScriptParams');
 
 module.exports = ({ db, studioConnection }) => async function executeScript(params) {
-  const { initiatedById, chatMessageId, script, roles } = new ExecuteScriptParams(params);
+  const { initiatedById, chatMessageId, script, dryRun, roles } = new ExecuteScriptParams(params);
   const ChatThread = studioConnection.model('__Studio_ChatThread');
   const ChatMessage = studioConnection.model('__Studio_ChatMessage');
 
@@ -39,32 +41,19 @@ module.exports = ({ db, studioConnection }) => async function executeScript(para
     throw new Error('Unauthorized');
   }
 
-  // Create a sandbox with the db object
-  const logs = [];
-  if (!db.Types) {
-    db.Types = mongoose.Types;
-  }
-  const sandbox = { db, mongoose, console: {}, ObjectId: mongoose.Types.ObjectId, MongooseStudioChartColors };
-
-  // Capture console logs
-  sandbox.console.log = function() {
-    const args = Array.from(arguments);
-    logs.push(args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg).join(' '));
-  };
-
-  const context = vm.createContext(sandbox);
+  const sandbox = createSandbox({ db });
 
   let output;
   let error;
 
   try {
     // Execute the script in the sandbox
-    output = await vm.runInContext(wrappedScript(script), context);
+    output = await sandbox.runScript({ script, dryRun });
 
     const updatedContent = updateContentWithScript(chatMessage.content, chatMessage.script, script);
     chatMessage.script = script;
     chatMessage.content = updatedContent;
-    chatMessage.executionResult = { output, logs: logs.join('\n'), error: null };
+    chatMessage.executionResult = { output, logs: sandbox.getLogs(), error: null, dryRun: !!dryRun };
     await chatMessage.save();
 
     return { chatMessage };
@@ -80,20 +69,23 @@ module.exports = ({ db, studioConnection }) => async function executeScript(para
         content: updatedContent,
         executionResult: {
           output: null,
-          logs: logs.join('\n'),
-          error
+          logs: sandbox.getLogs(),
+          error,
+          dryRun: !!dryRun
         }
       }
     );
 
     throw new Error(`Script execution failed: ${error}`);
+  } finally {
+    try {
+      await sandbox.close();
+    } catch (_) {
+      // Ignore sandbox cleanup errors so they do not mask the primary result.
+    }
   }
 };
 
-const wrappedScript = script => `(async () => {
-  ${script}
-})()`;
-
 function updateContentWithScript(content, previousScript, newScript) {
   if (typeof content !== 'string') {
     return content;

package/backend/actions/Model/exportQueryResults.js

@@ -62,9 +62,7 @@ module.exports = ({ db }) => async function exportQueryResults(params, req, res)
   }
   const csv = stringify(rows);
 
-  res.set({
-    'Content-Type': 'text/csv',
-    'Content-Disposition': `attachment; filename="${model.toLowerCase()}-export.csv"`
-  });
+  res.setHeader('Content-Type', 'text/csv');
+  res.setHeader('Content-Disposition', `attachment; filename="${model.toLowerCase()}-export.csv"`);
   return csv;
 };

package/backend/actions/getCapabilities.js

@@ -0,0 +1,42 @@
+'use strict';
+
+module.exports = ({ db, options }) => async function getCapabilities() {
+  const connection = db?.connection ?? db;
+  const admin = typeof connection?.db?.admin === 'function' ? connection.db.admin() : null;
+  const supportsAI = hasAIProviderKey(options);
+
+  if (admin == null || typeof admin.command !== 'function') {
+    return {
+      supportsChangeStreams: false,
+      supportsTransactions: false,
+      supportsAI
+    };
+  }
+
+  let hello;
+  try {
+    hello = await admin.command({ hello: 1 });
+  } catch (err) {
+    hello = await admin.command({ isMaster: 1 });
+  }
+
+  const maxWireVersion = hello?.maxWireVersion ?? 0;
+  const supportsModernMongoFeatures = maxWireVersion >= 7;
+  const isReplicaSet = !!hello?.setName;
+  const isMongos = hello?.msg === 'isdbgrid';
+  const hasSessions = hello?.logicalSessionTimeoutMinutes != null;
+
+  return {
+    supportsChangeStreams: supportsModernMongoFeatures && (isReplicaSet || isMongos),
+    supportsTransactions: supportsModernMongoFeatures && hasSessions && (isReplicaSet || isMongos),
+    supportsAI
+  };
+};
+
+function hasAIProviderKey(options) {
+  return [
+    options?.openAIAPIKey,
+    options?.anthropicAPIKey,
+    options?.googleGeminiAPIKey
+  ].some(value => typeof value === 'string' ? value.trim().length > 0 : !!value);
+}

package/backend/actions/index.js

@@ -3,7 +3,7 @@
 exports.ChatMessage = require('./ChatMessage');
 exports.ChatThread = require('./ChatThread');
 exports.Dashboard = require('./Dashboard');
+exports.getCapabilities = require('./getCapabilities');
 exports.Model = require('./Model');
-exports.Script = require('./Script');
 exports.status = require('./status');
 exports.Task = require('./Task');

package/backend/db/chatMessageSchema.js

@@ -28,7 +28,8 @@ const chatMessageSchema = new mongoose.Schema({
   executionResult: {
     output: mongoose.Schema.Types.Mixed,
     logs: String,
-    error: String
+    error: String,
+    dryRun: Boolean
   }
 }, { timestamps: true });
 

package/backend/next.js

@@ -81,8 +81,12 @@ module.exports = function next(conn, options) {
       throw new Error(`Action ${actionName} not found`);
     }
 
-    return actionFn(params)
+    return actionFn(params, req, res)
       .then(result => {
+        if (typeof result === 'string') {
+          res.status(200).send(result);
+          return result;
+        }
         res.status(200).json(result);
         return result;
       })

package/backend/sandbox/createSandbox.js

@@ -0,0 +1,83 @@
+'use strict';
+
+const MongooseStudioChartColors = require('../constants/mongooseStudioChartColors');
+const mongoose = require('mongoose');
+const vm = require('vm');
+const { createScriptDb } = require('./createScriptDb');
+
+module.exports = function createSandbox({ db }) {
+  const logs = [];
+  const scriptDb = createScriptDb(db);
+  if (!scriptDb.db.Types) {
+    scriptDb.db.Types = mongoose.Types;
+  }
+
+  const sandbox = {
+    db: scriptDb.db,
+    mongoose,
+    console: {},
+    ObjectId: mongoose.Types.ObjectId,
+    MongooseStudioChartColors
+  };
+
+  sandbox.console.log = function() {
+    const args = Array.from(arguments);
+    logs.push(args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg).join(' '));
+  };
+
+  return {
+    context: vm.createContext(sandbox),
+    db: scriptDb.db,
+    getLogs() {
+      return logs.join('\n');
+    },
+    setDryRunSession: scriptDb.setDryRunSession,
+    async runScript({ script, dryRun }) {
+      const wrappedScript = `(async () => {
+        ${script}
+      })()`;
+
+      if (!dryRun) {
+        return await vm.runInContext(wrappedScript, this.context);
+      }
+
+      return await runDryRunScript({
+        db: this.db,
+        context: this.context,
+        setDryRunSession: this.setDryRunSession,
+        wrappedScript
+      });
+    },
+    close: scriptDb.close
+  };
+};
+
+async function runDryRunScript({ db, context, setDryRunSession, wrappedScript }) {
+  if (typeof db?.startSession !== 'function') {
+    throw new Error('Dry run mode requires MongoDB sessions support');
+  }
+
+  const session = await db.startSession();
+  setDryRunSession(session);
+
+  try {
+    session.startTransaction();
+    const result = await vm.runInContext(wrappedScript, context);
+    await abortTransaction(session);
+    return result;
+  } catch (err) {
+    await abortTransaction(session);
+    throw err;
+  } finally {
+    setDryRunSession(null);
+    await session.endSession().catch(() => {});
+  }
+}
+
+async function abortTransaction(session) {
+  if (typeof session?.inTransaction === 'function' && !session.inTransaction()) {
+    return;
+  }
+
+  await session.abortTransaction().catch(() => {});
+}

package/backend/sandbox/createScriptDb.js

@@ -0,0 +1,201 @@
+'use strict';
+
+const wrappedCollections = new WeakSet();
+
+const collectionMethodOptionsIndex = new Map([
+  ['bulkWrite', 1],
+  ['countDocuments', 1],
+  ['createIndex', 1],
+  ['createIndexes', 1],
+  ['deleteMany', 1],
+  ['deleteOne', 1],
+  ['distinct', 2],
+  ['drop', 0],
+  ['dropIndex', 1],
+  ['dropIndexes', 0],
+  ['estimatedDocumentCount', 0],
+  ['find', 1],
+  ['findOne', 1],
+  ['findOneAndDelete', 1],
+  ['findOneAndReplace', 2],
+  ['findOneAndUpdate', 2],
+  ['insertMany', 1],
+  ['insertOne', 1],
+  ['rename', 1],
+  ['replaceOne', 2],
+  ['updateMany', 2],
+  ['updateOne', 2]
+]);
+
+function createScriptDb(db) {
+  const sourceConnection = db;
+
+  const scriptConnection = sourceConnection.useDb(sourceConnection.name, { useCache: false });
+  scriptConnection.config = {
+    ...scriptConnection.config,
+    autoCreate: false,
+    autoIndex: false
+  };
+
+  let dryRunSession = null;
+  const getSession = () => dryRunSession;
+  cloneModels(sourceConnection, scriptConnection, getSession);
+  wrapCollectionAccessors(scriptConnection, getSession);
+
+  const originalDebug = sourceConnection.options?.debug;
+  scriptConnection.set('debug', function() {
+    if (typeof originalDebug === 'function') {
+      return originalDebug.apply(this, arguments);
+    }
+  });
+
+  return {
+    db: scriptConnection,
+    setDryRunSession(session) {
+      dryRunSession = session;
+    },
+    close() {
+      if (Array.isArray(sourceConnection.otherDbs)) {
+        sourceConnection.otherDbs = sourceConnection.otherDbs.filter(db => db !== scriptConnection);
+      }
+      if (Array.isArray(scriptConnection.otherDbs)) {
+        scriptConnection.otherDbs = [];
+      }
+      if (sourceConnection.relatedDbs?.[scriptConnection.name] === scriptConnection) {
+        delete sourceConnection.relatedDbs[scriptConnection.name];
+      }
+    }
+  };
+}
+
+function cloneModels(sourceConnection, scriptConnection, getSession) {
+  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);
+  }
+}
+
+function getCollectionName(Model) {
+  return Model.collection?.collectionName ?? Model.collection?.name ?? Model.$__collection?.collectionName;
+}
+
+function wrapModelCollections(Model, getSession) {
+  const collection = Model?.collection ?? Model?.$__collection;
+  if (collection == null) {
+    return;
+  }
+
+  const wrappedCollection = wrapCollection(collection, getSession);
+  Model.collection = wrappedCollection;
+  Model.$__collection = wrappedCollection;
+
+  if (Model.prototype != null) {
+    Model.prototype.collection = wrappedCollection;
+    Model.prototype.$collection = wrappedCollection;
+    setModelCollectionSymbols(Model.prototype, wrappedCollection);
+  }
+}
+
+function setModelCollectionSymbols(target, collection) {
+  for (const symbol of Object.getOwnPropertySymbols(target)) {
+    if (symbol.description === 'mongoose#Model#collection') {
+      target[symbol] = collection;
+    }
+  }
+}
+
+function wrapCollectionAccessors(scriptConnection, getSession) {
+  const wrappedCollectionCache = new WeakMap();
+
+  const nativeDb = scriptConnection.db;
+  const originalDbCollection = nativeDb != null && typeof nativeDb.collection === 'function'
+    ? nativeDb.collection
+    : null;
+
+  if (typeof scriptConnection.collection === 'function') {
+    const originalConnectionCollection = scriptConnection.collection;
+    scriptConnection.collection = function() {
+      const collection = originalConnectionCollection.apply(this, arguments);
+      if (collection != null && originalDbCollection != null && collection.name != null) {
+        if (collection.collection == null || wrappedCollections.has(collection.collection)) {
+          collection.collection = originalDbCollection.call(nativeDb, collection.name);
+        }
+      }
+      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession);
+    };
+  }
+
+  if (originalDbCollection != null) {
+    nativeDb.collection = function() {
+      const collection = originalDbCollection.apply(this, arguments);
+      return getOrCreateWrappedCollection(wrappedCollectionCache, collection, getSession);
+    };
+  }
+}
+
+function getOrCreateWrappedCollection(cache, collection, getSession) {
+  if (collection == null || (typeof collection !== 'object' && typeof collection !== 'function')) {
+    return collection;
+  }
+  if (wrappedCollections.has(collection)) {
+    return collection;
+  }
+  if (cache.has(collection)) {
+    return cache.get(collection);
+  }
+  const wrapped = wrapCollection(collection, getSession);
+  cache.set(collection, wrapped);
+  return wrapped;
+}
+
+function wrapCollection(collection, getSession) {
+  if (wrappedCollections.has(collection)) {
+    return collection;
+  }
+  const wrapped = Object.create(collection);
+  wrappedCollections.add(wrapped);
+  for (const [methodName, optionsIndex] of collectionMethodOptionsIndex) {
+    const method = collection[methodName];
+    if (typeof method !== 'function') {
+      continue;
+    }
+
+    wrapped[methodName] = function() {
+      const args = addSessionOption(Array.from(arguments), optionsIndex, getSession());
+      return method.apply(collection, args);
+    };
+  }
+  return wrapped;
+}
+
+function addSessionOption(args, optionsIndex, session) {
+  if (session == null) {
+    return args;
+  }
+
+  while (args.length < optionsIndex) {
+    args.push(undefined);
+  }
+
+  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 (options?.session != null) {
+    throw new Error('Cannot run dry run on script that uses sessions');
+  }
+
+  args[optionsIndex] = { ...(options ?? {}), session };
+  return args;
+}
+
+module.exports = {
+  createScriptDb,
+  collectionMethodOptionsIndex
+};