Merge pull request #14505 from Automattic/vkarpov15/set-driver-model-db-2

fix(mongoose): make setDriver() update mongoose.model() connections and collections

Author: vkarpov15

lib/model.js

@@ -4895,6 +4895,34 @@ Model.compile = function compile(name, schema, collectionName, connection, base)
   return model;
 };
 
+/**
+ * Update this model to use the new connection, including updating all internal
+ * references and creating a new `COllection` instance using the new connection.
+ * Not for external use, only used by `setDriver()` to ensure that you can still
+ * call `setDriver()` after creating a model using `mongoose.model()`.
+ *
+ * @param {Connection} newConnection the new connection to use
+ * @api private
+ */
+
+Model.$__updateConnection = function $__updateConnection(newConnection) {
+  this.db = newConnection;
+  this.prototype.db = newConnection;
+  this.prototype[modelDbSymbol] = newConnection;
+
+  const collection = newConnection.collection(
+    this.collection.collectionName,
+    this.collection.opts
+  );
+
+  this.prototype.collection = collection;
+  this.prototype.$collection = collection;
+  this.prototype[modelCollectionSymbol] = collection;
+
+  this.collection = collection;
+  this.$__collection = collection;
+};
+
 /**
  * Register custom query methods for this model
  *

lib/mongoose.js

@@ -166,9 +166,19 @@ Mongoose.prototype.setDriver = function setDriver(driver) {
   _mongoose.__driver = driver;
 
   const Connection = driver.Connection;
+  const oldDefaultConnection = _mongoose.connections[0];
   _mongoose.connections = [new Connection(_mongoose)];
   _mongoose.connections[0].models = _mongoose.models;
 
+  // Update all models that pointed to the old default connection to
+  // the new default connection, including collections
+  for (const model of Object.values(_mongoose.models)) {
+    if (model.db !== oldDefaultConnection) {
+      continue;
+    }
+    model.$__updateConnection(_mongoose.connections[0]);
+  }
+
   return _mongoose;
 };