fix: backwards compat for connect() #3790 (Fix #3847)

Automattic · Feb 3, 2016 · dc9612a · dc9612a
1 parent 559fa6c
commit dc9612a
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 7 deletions.
diff --git a/lib/index.js b/lib/index.js
@@ -230,12 +230,13 @@ Mongoose.prototype.createConnection = function(uri, options) {
 
 Mongoose.prototype.connect = function() {
   var conn = this.connection;
-
   if (rgxReplSet.test(arguments[0]) || checkReplicaSetInUri(arguments[0])) {
-    return conn.openSet.apply(conn, arguments);
+    this.$opPromise = conn.openSet.apply(conn, arguments);
+    return this;
   }
 
-  return conn.open.apply(conn, arguments);
+  this.$opPromise = conn.open.apply(conn, arguments);
+  return this;
 };
 
 /**
@@ -248,7 +249,6 @@ Mongoose.prototype.connect = function() {
 
 Mongoose.prototype.disconnect = function(fn) {
   var error;
-
   this.connections.forEach(function(conn) {
     conn.close(function(err) {
       if (error) {
@@ -261,14 +261,50 @@ Mongoose.prototype.disconnect = function(fn) {
   });
 
   var Promise = PromiseProvider.get();
-  return new Promise.ES6(function(resolve, reject) {
+  this.$opPromise = new Promise.ES6(function(resolve, reject) {
     fn && fn(error);
     if (error) {
       reject(error);
       return;
     }
     resolve();
   });
+  return this;
+};
+
+/**
+ * Ability to use mongoose object as a pseudo-promise so `.connect().then()`
+ * and `.disconnect().then()` are viable.
+ *
+ * @param {Function} onFulfilled
+ * @param {Function} onRejected
+ * @return {Promise}
+ * @api public
+ */
+
+Mongoose.prototype.then = function(onFulfilled, onRejected) {
+  var Promise = PromiseProvider.get();
+  if (!this.$opPromise) {
+    return new Promise.ES6(function(resolve, reject) {
+      reject(new Error('Can only call `.then()` if connect() or disconnect() ' +
+        'has been called'));
+    }).then(onFulfilled, onRejected);
+  }
+  return this.$opPromise.then(onFulfilled, onRejected);
+};
+
+/**
+ * Ability to use mongoose object as a pseudo-promise so `.connect().then()`
+ * and `.disconnect().then()` are viable.
+ *
+ * @param {Function} onFulfilled
+ * @param {Function} onRejected
+ * @return {Promise}
+ * @api public
+ */
+
+Mongoose.prototype.catch = function(onRejected) {
+  return this.then(null, onRejected);
 };
 
 /**

diff --git a/test/index.test.js b/test/index.test.js
@@ -37,6 +37,16 @@ describe('mongoose module:', function() {
         });
       });
     });
+
+    it('with promise (gh-3790)', function(done) {
+      var goose = new Mongoose;
+      var db = goose.connection,
+          uri = 'mongodb://localhost/mongoose_test';
+
+      goose.connect(process.env.MONGOOSE_TEST_URI || uri).then(function() {
+        db.close(done);
+      });
+    });
   });
 
   it('{g,s}etting options', function(done) {
@@ -128,8 +138,8 @@ describe('mongoose module:', function() {
           cb(new Error('bam'));
         };
 
-        mong.disconnect()
-          .on('error', function(error) {
+        mong.disconnect().connection.
+          on('error', function(error) {
             assert.equal('bam', error.message);
           });
 
@@ -149,6 +159,17 @@ describe('mongoose module:', function() {
         });
       });
     });
+
+    it('with promise (gh-3790)', function(done) {
+      var mong = new Mongoose();
+      var uri = 'mongodb://localhost/mongoose_test';
+
+      mong.connect(process.env.MONGOOSE_TEST_URI || uri);
+
+      mong.connection.on('open', function() {
+        mong.disconnect().then(function() { done(); });
+      });
+    });
   });
 
   describe('model()', function() {