Add option to return affected rows

brabeji · Jun 25, 2014 · 8198aa0 · 8198aa0
1 parent cc99c74
commit 8198aa0
Show file tree

Hide file tree

Showing 8 changed files with 136 additions and 107 deletions.
diff --git a/lib/dialects/abstract/query-generator.js b/lib/dialects/abstract/query-generator.js
@@ -143,7 +143,9 @@ module.exports = (function() {
     /*
       Returns an insert into command. Parameters: table name + hash of attribute-value-pairs.
     */
-    insertQuery: function(table, valueHash, modelAttributes) {
+    insertQuery: function(table, valueHash, modelAttributes, options) {
+      options = options || {};
+
       var query
         , valueQuery = 'INSERT INTO <%= table %> (<%= attributes %>) VALUES (<%= values %>)'
         , emptyQuery = 'INSERT INTO <%= table %>'
@@ -168,7 +170,7 @@ module.exports = (function() {
         emptyQuery += ' VALUES ()';
       }
 
-      if (this._dialect.supports['RETURNING']) {
+      if (this._dialect.supports['RETURNING'] && options.returning) {
         valueQuery += ' RETURNING *';
         emptyQuery += ' RETURNING *';
       }
@@ -235,7 +237,7 @@ module.exports = (function() {
         query += ' LIMIT ' + this.escape(options.limit) + ' ';
       }
 
-      if (this._dialect.supports['RETURNING'] && (options.returning || options.returning === undefined)) {
+      if (this._dialect.supports['RETURNING'] && options.returning) {
         query += ' RETURNING *';
       }
 

diff --git a/lib/dialects/postgres/query-generator.js b/lib/dialects/postgres/query-generator.js
@@ -253,11 +253,17 @@ module.exports = (function() {
     },
 
     bulkInsertQuery: function(tableName, attrValueHashes, options, modelAttributes) {
-      var query = 'INSERT INTO <%= table %> (<%= attributes %>) VALUES <%= tuples %> RETURNING *;'
+      options = options || {};
+
+      var query = 'INSERT INTO <%= table %> (<%= attributes %>) VALUES <%= tuples %>'
         , tuples = []
         , serials = []
         , allAttributes = [];
 
+      if (this._dialect.supports['RETURNING'] && options.returning) {
+        query += ' RETURNING *';
+      }
+
       Utils._.forEach(attrValueHashes, function(attrValueHash) {
         Utils._.forOwn(attrValueHash, function(value, key) {
           if (allAttributes.indexOf(key) === -1) {
@@ -289,6 +295,8 @@ module.exports = (function() {
       , tuples: tuples.join(',')
       };
 
+      query = query + ';';
+
       return Utils._.template(query)(replacements);
     },
 

diff --git a/lib/dialects/postgres/query.js b/lib/dialects/postgres/query.js
@@ -146,6 +146,10 @@ module.exports = (function() {
       } else if (self.send('isShowOrDescribeQuery')) {
         return results;
       } else if (QueryTypes.BULKUPDATE === self.options.type) {
+        if (!self.options.returning) {
+          return result.rowCount;
+        }
+
         if (!!self.callee && !!self.callee._hasHstoreAttributes) {
           rows.forEach(function(row) {
             parseHstoreFields(self.callee, row);

diff --git a/lib/model.js b/lib/model.js
@@ -1341,6 +1341,7 @@ module.exports = (function() {
    * @param  {Boolean}  [options.validate=true]         Should each row be subject to validation before it is inserted. The whole insert will fail if one row fails validation
    * @param  {Boolean}  [options.hooks=true]            Run before / after bulk update hooks?
    * @param  {Boolean}  [options.individualHooks=false] Run before / after update hooks?
+   * @param  {Boolean}  [options.returning=false]       Return the affected rows (only for postgres)
    * @param  {Number}   [options.limit]                 How many rows to update (only for mysql and mariadb)
    * @deprecated The syntax is due for change, in order to make `where` more consistent with the rest of the API
    *
@@ -1353,6 +1354,7 @@ module.exports = (function() {
       validate: true,
       hooks: true,
       individualHooks: false,
+      returning: false,
       force: false
     }, options || {});
 
@@ -1458,6 +1460,10 @@ module.exports = (function() {
 
       // Run query to update all rows
       return self.QueryInterface.bulkUpdate(self.getTableName(), attrValueHashUse, where, options, self.tableAttributes).then(function(affectedRows) {
+        if (options.returning) {
+          return [affectedRows.length, affectedRows];
+        }
+
         return [affectedRows];
       });
     }).tap(function(result) {

diff --git a/lib/query-interface.js b/lib/query-interface.js
@@ -414,7 +414,7 @@ module.exports = (function() {
   };
 
   QueryInterface.prototype.insert = function(dao, tableName, values, options) {
-    var sql = this.QueryGenerator.insertQuery(tableName, values, dao.Model.rawAttributes);
+    var sql = this.QueryGenerator.insertQuery(tableName, values, dao.Model.rawAttributes, options);
 
     return this.sequelize.query(sql, dao, options).then(function(result) {
       result.isNewRecord = false;

diff --git a/test/dao-factory.test.js b/test/dao-factory.test.js
@@ -257,8 +257,8 @@ describe(Support.getTestDialectTeaser("DAOFactory"), function () {
       var titleSetter = sinon.spy()
         , Task = this.sequelize.define('TaskBuild', {
           title:  {
-            type: Sequelize.STRING(50), 
-            allowNull: false, 
+            type: Sequelize.STRING(50),
+            allowNull: false,
             defaultValue: ''
           }
         }, {
@@ -783,47 +783,47 @@ describe(Support.getTestDialectTeaser("DAOFactory"), function () {
       })
     })
 
-    if (dialect === "postgres") {
-      it('returns the affected rows', function(_done) {
-       var self = this
-          , data = [{ username: 'Peter', secretValue: '42' },
-                    { username: 'Paul',  secretValue: '42' },
-                    { username: 'Bob',   secretValue: '43' }]
-          , done = _.after(2, _done)
+    it('returns the number of affected rows', function(_done) {
+     var self = this
+        , data = [{ username: 'Peter', secretValue: '42' },
+                  { username: 'Paul',  secretValue: '42' },
+                  { username: 'Bob',   secretValue: '43' }]
+        , done = _.after(2, _done)
 
-        this.User.bulkCreate(data).success(function() {
-          self.User.update({username: 'Bill'}, {secretValue: '42'}).spread(function(affectedRows) {
-            expect(affectedRows).to.have.length(2)
+      this.User.bulkCreate(data).success(function() {
+        self.User.update({username: 'Bill'}, {secretValue: '42'}).spread(function(affectedRows) {
+          expect(affectedRows).to.equal(2)
 
-            done()
-          })
+          done()
+        })
 
-          self.User.update({username: 'Bill'}, {secretValue: '44'}).spread(function(affectedRows) {
-            expect(affectedRows).to.have.length(0)
+        self.User.update({username: 'Bill'}, {secretValue: '44'}).spread(function(affectedRows) {
+          expect(affectedRows).to.equal(0)
 
-            done()
-          })
+          done()
         })
       })
-    }
+    })
 
-    if (dialect !== "postgres") {
-      it('returns the number of affected rows', function(_done) {
+    if (dialect === "postgres") {
+      it('returns the affected rows if `options.returning` is true', function(_done) {
        var self = this
           , data = [{ username: 'Peter', secretValue: '42' },
                     { username: 'Paul',  secretValue: '42' },
                     { username: 'Bob',   secretValue: '43' }]
           , done = _.after(2, _done)
 
         this.User.bulkCreate(data).success(function() {
-          self.User.update({username: 'Bill'}, {secretValue: '42'}).spread(function(affectedRows) {
-            expect(affectedRows).to.equal(2)
+          self.User.update({ username: 'Bill' }, { secretValue: '42' }, { returning: true }).spread(function(count, rows) {
+            expect(count).to.equal(2)
+            expect(rows).to.have.length(2)
 
             done()
           })
 
-          self.User.update({username: 'Bill'}, {secretValue: '44'}).spread(function(affectedRows) {
-            expect(affectedRows).to.equal(0)
+          self.User.update({ username: 'Bill'}, { secretValue: '44' }, { returning: true }).spread(function(count, rows) {
+            expect(count).to.equal(0)
+            expect(rows).to.have.length(0)
 
             done()
           })

diff --git a/test/postgres/dao.test.js b/test/postgres/dao.test.js
@@ -295,16 +295,16 @@ if (dialect.match(/^postgres/)) {
           .error(console.log)
       })
 
-      it("should update hstore correctly and return affected rows", function(done) {
+      it("should update hstore correctly and return the affected rows", function(done) {
         var self = this
 
         this.User
           .create({ username: 'user', email: ['[email protected]'], settings: { created: { test: '"value"' }}})
           .success(function(oldUser) {
             // Update the user and check that the returned object's fields have been parsed by the hstore library
-            self.User.update({settings: {should: 'update', to: 'this', first: 'place'}}, oldUser.identifiers).spread(function(newUsers) {
-              expect(newUsers).to.have.length(1);
-              expect(newUsers[0].settings).to.deep.equal({should: 'update', to: 'this', first: 'place'})
+            self.User.update({settings: {should: 'update', to: 'this', first: 'place'}}, oldUser.identifiers, { returning: true }).spread(function(count, users) {
+              expect(count).to.equal(1);
+              expect(users[0].settings).to.deep.equal({should: 'update', to: 'this', first: 'place'})
               done()
             })
           })