From 06c38bcc2d8b50054bc7c9d844bf4359288cdc98 Mon Sep 17 00:00:00 2001
From: Tim Griesser <tgriesser10@gmail.com>
Date: Tue, 10 Dec 2013 14:54:51 -0500
Subject: [PATCH] Allow for returning in update queries, #132

---
 clients/server/postgres/grammar.js  | 17 +++++++++++++++--
 lib/builder.js                      |  3 ++-
 test/integration/builder/updates.js | 10 ++++++++++
 test/integration/output/Updates.js  | 27 +++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/clients/server/postgres/grammar.js b/clients/server/postgres/grammar.js
index aa64ce72..86e0890b 100644
--- a/clients/server/postgres/grammar.js
+++ b/clients/server/postgres/grammar.js
@@ -38,7 +38,20 @@ exports.grammar = _.defaults({
       }
       sql += "(" + this.columnize(columns) + ") values " + paramBlocks.join(', ');
     }
+    sql += this.compileReturning(qb);
+    return sql;
+  },
 
+  // Compiles an `update` query, allowing for a return value.
+  compileUpdate: function(qb) {
+    var sql = baseGrammar.compileUpdate.apply(this, arguments);
+    sql += this.compileReturning(qb);
+    return sql;
+  },
+
+  // Adds the returning value to the statement.
+  compileReturning: function(qb) {
+    var sql = '';
     if (qb.flags.returning) {
       if (_.isArray(qb.flags.returning)) {
         sql += ' returning ' + this.wrapArray(qb.flags.returning);
@@ -51,9 +64,9 @@ exports.grammar = _.defaults({
 
   // Ensures the response is returned in the same format as other clients.
   handleResponse: function(builder, response) {
+    var returning = builder.flags.returning;
     if (response.command === 'SELECT') return response.rows;
-    if (response.command === 'INSERT') {
-      var returning = builder.flags.returning;
+    if (response.command === 'INSERT' || (response.command === 'UPDATE' && returning)) {
       return _.map(response.rows, function(row) {
         if (returning === '*' || _.isArray(returning)) return row;
         return row[returning];
diff --git a/lib/builder.js b/lib/builder.js
index 57f4a302..267595bb 100644
--- a/lib/builder.js
+++ b/lib/builder.js
@@ -404,7 +404,8 @@ _.extend(Builder.prototype, Common, {
   },
 
   // Sets the values for an `update` query.
-  update: function(values) {
+  update: function(values, returning) {
+    if (returning) this.returning(returning);
     var obj = Helpers.sortObject(values);
     var bindings = [];
     for (var i = 0, l = obj.length; i < l; i++) {
diff --git a/test/integration/builder/updates.js b/test/integration/builder/updates.js
index 4c97f406..d931abf5 100644
--- a/test/integration/builder/updates.js
+++ b/test/integration/builder/updates.js
@@ -50,6 +50,16 @@ module.exports = function(knex) {
 
     });
 
+    it('should allow returning for updates in postgresql', function() {
+
+      return knex('accounts').logMe().where('id', 1).update({
+        first_name: 'UpdatedUser',
+        last_name: 'UpdatedTest',
+        email:'test100@example.com'
+      }, '*');
+
+    });
+
   });
 
 };
\ No newline at end of file
diff --git a/test/integration/output/Updates.js b/test/integration/output/Updates.js
index 5c07d5a3..b57c7bcb 100644
--- a/test/integration/output/Updates.js
+++ b/test/integration/output/Updates.js
@@ -15,5 +15,32 @@ module.exports = {
       sql: 'update "accounts" set "email" = ?, "first_name" = ?, "last_name" = ? where "id" = ?',
       result: 1
     }
+  },
+  'should allow returning for updates in postgresql': {
+    mysql: {
+      bindings: ['test100@example.com','UpdatedUser','UpdatedTest',1],
+      sql: 'update `accounts` set `email` = ?, `first_name` = ?, `last_name` = ? where `id` = ?',
+      result: 1
+    },
+    postgresql: {
+      bindings: ['test100@example.com','UpdatedUser','UpdatedTest',1],
+      sql: 'update "accounts" set "email" = ?, "first_name" = ?, "last_name" = ? where "id" = ? returning *',
+      result: [{
+        id: '1',
+        first_name: 'UpdatedUser',
+        last_name: 'UpdatedTest',
+        email: 'test100@example.com',
+        logins: 1,
+        about: 'Lorem ipsum Dolore labore incididunt enim.',
+        created_at: new Date(),
+        updated_at: new Date(),
+        phone: null
+      }]
+    },
+    sqlite3: {
+      bindings: ['test100@example.com','UpdatedUser','UpdatedTest',1],
+      sql: 'update "accounts" set "email" = ?, "first_name" = ?, "last_name" = ? where "id" = ?',
+      result: 1
+    }
   }
 };
\ No newline at end of file