Merge pull request #11900 from Automattic/vkarpov15/driver-fixes

feat(mongoose): add `setDriver()` function to allow overwriting `driver` in a more consistent way

Author: vkarpov15

lib/connection.js

@@ -1068,7 +1068,7 @@ Connection.prototype.collection = function(name, options) {
   };
   options = Object.assign({}, defaultOptions, options ? utils.clone(options) : {});
   options.$wasForceClosed = this.$wasForceClosed;
-  const Collection = driver.get().Collection;
+  const Collection = this.base && this.base.__driver && this.base.__driver.Collection || driver.get().Collection;
   if (!(name in this.collections)) {
     this.collections[name] = new Collection(name, this, options);
   }

lib/index.js

@@ -35,6 +35,7 @@ const shardingPlugin = require('./plugins/sharding');
 const trusted = require('./helpers/query/trusted').trusted;
 const sanitizeFilter = require('./helpers/query/sanitizeFilter');
 const isBsonType = require('./helpers/isBsonType');
+const MongooseError = require('./error/mongooseError');
 
 const defaultMongooseSymbol = Symbol.for('mongoose:default');
 
@@ -62,6 +63,7 @@ function Mongoose(options) {
   this.connections = [];
   this.models = {};
   this.events = new EventEmitter();
+  this.__driver = driver.get();
   // default global options
   this.options = Object.assign({
     pluralization: true,
@@ -135,13 +137,44 @@ Mongoose.prototype.ConnectionStates = STATES;
  * uses to communicate with the database. A driver is a Mongoose-specific interface that defines functions
  * like `find()`.
  *
+ * @deprecated
  * @memberOf Mongoose
  * @property driver
  * @api public
  */
 
 Mongoose.prototype.driver = driver;
 
+/**
+ * Overwrites the current driver used by this Mongoose instance. A driver is a
+ * Mongoose-specific interface that defines functions like `find()`.
+ *
+ * @memberOf Mongoose
+ * @method setDriver
+ * @api public
+ */
+
+Mongoose.prototype.setDriver = function setDriver(driver) {
+  const _mongoose = this instanceof Mongoose ? this : mongoose;
+
+  if (_mongoose.__driver === driver) {
+    return _mongoose;
+  }
+
+  const openConnection = _mongoose.connections && _mongoose.connections.find(conn => conn.readyState !== STATES.disconnected);
+  if (openConnection) {
+    const msg = 'Cannot modify Mongoose driver if a connection is already open. ' +
+      'Call `mongoose.disconnect()` before modifying the driver';
+    throw new MongooseError(msg);
+  }
+  _mongoose.__driver = driver;
+
+  const Connection = driver.getConnection();
+  _mongoose.connections = [new Connection(_mongoose)];
+
+  return _mongoose;
+};
+
 /**
  * Sets mongoose options
  *
@@ -278,7 +311,7 @@ Mongoose.prototype.get = Mongoose.prototype.set;
 Mongoose.prototype.createConnection = function(uri, options, callback) {
   const _mongoose = this instanceof Mongoose ? this : mongoose;
 
-  const Connection = driver.get().getConnection();
+  const Connection = _mongoose.__driver.getConnection();
   const conn = new Connection(_mongoose);
   if (typeof options === 'function') {
     callback = options;
@@ -474,7 +507,6 @@ Mongoose.prototype.pluralize = function(fn) {
  */
 
 Mongoose.prototype.model = function(name, schema, collection, options) {
-
   const _mongoose = this instanceof Mongoose ? this : mongoose;
 
   if (typeof schema === 'string') {
@@ -691,7 +723,7 @@ Mongoose.prototype.__defineGetter__('connection', function() {
 });
 
 Mongoose.prototype.__defineSetter__('connection', function(v) {
-  if (v instanceof Connection) {
+  if (v instanceof this.__driver.getConnection()) {
     this.connections[0] = v;
     this.models = v.models;
   }
@@ -720,18 +752,6 @@ Mongoose.prototype.__defineSetter__('connection', function(v) {
 
 Mongoose.prototype.connections;
 
-/*!
- * Connection
- */
-
-const Connection = driver.get().getConnection();
-
-/*!
- * Collection
- */
-
-const Collection = driver.get().Collection;
-
 /**
  * The Mongoose Aggregate constructor
  *
@@ -748,7 +768,14 @@ Mongoose.prototype.Aggregate = Aggregate;
  * @api public
  */
 
-Mongoose.prototype.Collection = Collection;
+Object.defineProperty(Mongoose.prototype, 'Collection', {
+  get: function() {
+    return this.__driver.Collection;
+  },
+  set: function(Collection) {
+    this.__driver.Collection = Collection;
+  }
+});
 
 /**
  * The Mongoose [Connection](#connection_Connection) constructor
@@ -759,7 +786,18 @@ Mongoose.prototype.Collection = Collection;
  * @api public
  */
 
-Mongoose.prototype.Connection = Connection;
+Object.defineProperty(Mongoose.prototype, 'Connection', {
+  get: function() {
+    return this.__driver.getConnection();
+  },
+  set: function(Connection) {
+    if (Connection === this.__driver.getConnection()) {
+      return;
+    }
+
+    this.__driver.getConnection = () => Connection;
+  }
+});
 
 /**
  * The Mongoose version

lib/query.js

@@ -2197,7 +2197,6 @@ function _castArrayFilters(query) {
  * @api private
  */
 Query.prototype._find = wrapThunk(function(callback) {
-
   this._castConditions();
 
   if (this.error() != null) {

lib/utils.js

@@ -935,6 +935,9 @@ exports.getOption = function(name) {
   const sources = Array.prototype.slice.call(arguments, 1);
 
   for (const source of sources) {
+    if (source == null) {
+      continue;
+    }
     if (source[name] != null) {
       return source[name];
     }

test/index.test.js

@@ -7,6 +7,7 @@ const start = require('./common');
 const assert = require('assert');
 const random = require('./util').random;
 const stream = require('stream');
+const { EventEmitter } = require('events');
 
 const collection = 'blogposts_' + random();
 
@@ -1044,4 +1045,45 @@ describe('mongoose module:', function() {
       });
     });
   });
+
+  describe('custom drivers', function() {
+    it('can set custom driver (gh-11900)', async function() {
+      const m = new mongoose.Mongoose();
+
+      class Collection {
+        findOne(filter, options, cb) {
+          cb(null, { answer: 42 });
+        }
+      }
+      class Connection extends EventEmitter {
+        constructor(base) {
+          super();
+          this.base = base;
+          this.models = {};
+        }
+
+        collection() {
+          return new Collection();
+        }
+
+        openUri(uri, opts, callback) {
+          this.readyState = mongoose.ConnectionStates.connected;
+          callback();
+        }
+      }
+      const driver = {
+        Collection,
+        getConnection: () => Connection
+      };
+
+      m.setDriver(driver);
+
+      await m.connect();
+
+      const Test = m.model('Test', m.Schema({ answer: Number }));
+
+      const res = await Test.findOne();
+      assert.deepEqual(res.toObject(), { answer: 42 });
+    });
+  });
 });

test/types/populate.test.ts

@@ -247,6 +247,31 @@ async function _11532() {
   expectError(leanResult?.__v);
 }
 
+async function gh11710() {
+
+  // `Parent` represents the object as it is stored in MongoDB
+  interface Parent {
+    child?: Types.ObjectId,
+    name?: string
+  }
+  interface Child {
+    name: string;
+  }
+  interface PopulatedParent {
+    child: Child | null;
+  }
+  const ParentModel = model<Parent>('Parent', new Schema({
+    child: { type: Schema.Types.ObjectId, ref: 'Child' },
+    name: String
+  }));
+  const childSchema: Schema = new Schema({ name: String });
+  const ChildModel = model<Child>('Child', childSchema);
+
+  // Populate with `Paths` generic `{ child: Child }` to override `child` path
+  const doc = await ParentModel.findOne({}).populate<Pick<PopulatedParent, 'child'>>('child').orFail();
+  expectType<Child | null>(doc.child);
+}
+
 function gh11758() {
   interface NestedChild {
     name: string

types/index.d.ts

@@ -81,6 +81,12 @@ declare module 'mongoose' {
   /** Returns an array of model names created on this instance of Mongoose. */
   export function modelNames(): Array<string>;
 
+  /**
+   * Overwrites the current driver used by this Mongoose instance. A driver is a
+   * Mongoose-specific interface that defines functions like `find()`.
+   */
+  export function setDriver(driver: any): Mongoose;
+
   /** The node-mongodb-native driver Mongoose uses. */
   export const mongo: typeof mongodb;