feat(upsert): return upserted record with options.returning=true (#8924)

Sushant · GitHub
Commit a4a0f883 authored Jan 22, 2018 by Sushant Committed by GitHub Jan 22, 2018
Showing with 191 additions and 95 deletions
lib/dialects/mssql/query-generator.js
lib/dialects/mssql/query.js
lib/dialects/postgres/query-generator.js
lib/dialects/postgres/query.js
lib/dialects/sqlite/query-interface.js
lib/model.js
lib/query-interface.js
test/integration/model/upsert.test.js
test/unit/model/upsert.test.js
--- a/lib/dialects/mssql/query-generator.js
+++ b/lib/dialects/mssql/query-generator.js
@@ -335,7 +335,6 @@ const QueryGenerator = {
    const tableNameQuoted = this.quoteTable(tableName);
    let needIdentityInsertWrapper = false;
    //Obtain primaryKeys, uniquekeys and identity attrs from rawAttributes as model is not passed
    for (const key in model.rawAttributes) {
      if (model.rawAttributes[key].primaryKey) {

--- a/lib/dialects/mssql/query.js
+++ b/lib/dialects/mssql/query.js
@@ -213,12 +213,7 @@ class Query extends AbstractQuery {
    } else if (this.isSelectQuery()) {
      result = this.handleSelectQuery(data);
    } else if (this.isUpsertQuery()) {
-      //Use the same return value as that of MySQL & Postgres
+      result = data[0];
-      if (data[0].$action === 'INSERT') {
-        result = 1;
-      } else {
-        result = 2;
-      }
    } else if (this.isCallQuery()) {
      result = data[0];
    } else if (this.isBulkUpdateQuery()) {

--- a/lib/dialects/postgres/query-generator.js
+++ b/lib/dialects/postgres/query-generator.js
@@ -345,34 +345,38 @@ const QueryGenerator = {
    return `ALTER TABLE ${this.quoteTable(tableName)} RENAME COLUMN ${attrString.join(', ')};`;
  },
-  fn(fnName, tableName, body, returns, language) {
+  fn(fnName, tableName, parameters, body, returns, language) {
    fnName = fnName || 'testfunc';
    language = language || 'plpgsql';
-    returns = returns || 'SETOF ' + this.quoteTable(tableName);
+    returns = returns ? `RETURNS ${returns}` : '';
+    parameters = parameters || '';
-    return `CREATE OR REPLACE FUNCTION pg_temp.${fnName}() RETURNS ${returns} AS $func$ BEGIN ${body} END; $func$ LANGUAGE ${language}; SELECT * FROM pg_temp.${fnName}();`;
+    return `CREATE OR REPLACE FUNCTION pg_temp.${fnName}(${parameters}) ${returns} AS $func$ BEGIN ${body} END; $func$ LANGUAGE ${language}; SELECT * FROM pg_temp.${fnName}();`;
  },
-  exceptionFn(fnName, tableName, main, then, when, returns, language) {
+  exceptionFn(fnName, tableName, parameters, main, then, when, returns, language) {
    when = when || 'unique_violation';
    const body = `${main} EXCEPTION WHEN ${when} THEN ${then};`;
-    return this.fn(fnName, tableName, body, returns, language);
+    return this.fn(fnName, tableName, parameters, body, returns, language);
  },
  upsertQuery(tableName, insertValues, updateValues, where, model, options) {
-    const insert = this.insertQuery(tableName, insertValues, model.rawAttributes, options);
+    const primaryField = this.quoteIdentifier(model.primaryKeyField);
-    const update = this.updateQuery(tableName, updateValues, where, options, model.rawAttributes);
+    let insert = this.insertQuery(tableName, insertValues, model.rawAttributes, options);
+    let update = this.updateQuery(tableName, updateValues, where, options, model.rawAttributes);
+    insert = insert.replace('RETURNING *', `RETURNING ${primaryField} INTO primary_key`);
+    update = update.replace('RETURNING *', `RETURNING ${primaryField} INTO primary_key`);
-    // The numbers here are selected to match the number of affected rows returned by MySQL
    return this.exceptionFn(
      'sequelize_upsert',
      tableName,
-      insert + ' RETURN 1;',
+      'OUT created boolean, OUT primary_key text',
-      update + '; RETURN 2',
+      `${insert} created := true;`,
-      'unique_violation',
+      `${update}; created := false`
-      'integer'
    );
  },

--- a/lib/dialects/postgres/query.js
+++ b/lib/dialects/postgres/query.js
@@ -236,12 +236,11 @@ class Query extends AbstractQuery {
          if (!this.options.returning) {
            return parseInt(rowCount, 10);
          }
          return this.handleSelectQuery(rows);
        } else if (QueryTypes.BULKDELETE === this.options.type) {
          return parseInt(rowCount, 10);
        } else if (this.isUpsertQuery()) {
-          return rows[0].sequelize_upsert;
+          return rows[0];
        } else if (this.isInsertQuery() || this.isUpdateQuery()) {
          if (this.instance && this.instance.dataValues) {
            for (const key in rows[0]) {

--- a/lib/dialects/sqlite/query-interface.js
+++ b/lib/dialects/sqlite/query-interface.js
@@ -171,6 +171,8 @@ exports.addConstraint = addConstraint;
 *
 * @param {String} tableName
 * @param {Object} options  Query Options
+ *
+ * @private
 * @returns {Promise}
 */
 function getForeignKeyReferencesForTable(tableName, options) {

--- a/lib/model.js
+++ b/lib/model.js
@@ -2215,22 +2215,25 @@ class Model {
   * @param  {Boolean}      [options.validate=true] Run validations before the row is inserted
   * @param  {Array}        [options.fields=Object.keys(this.attributes)] The fields to insert / update. Defaults to all changed fields
   * @param  {Boolean}      [options.hooks=true]  Run before / after upsert hooks?
+   * @param  {Boolean}      [options.returning=false] Append RETURNING * to get back auto generated values (Postgres only)
   * @param  {Transaction}  [options.transaction] Transaction to run query under
   * @param  {Function}     [options.logging=false] A function that gets executed while running the query to log the sql.
   * @param  {Boolean}      [options.benchmark=false] Pass query execution time in milliseconds as second argument to logging function (options.logging).
   * @param  {String}       [options.searchPath=DEFAULT] An optional parameter to specify the schema search_path (Postgres only)
   *
-   * @return {Promise<created>} Returns a boolean indicating whether the row was created or updated.
+   * @return {Promise<created>} Returns a boolean indicating whether the row was created or updated. For Postgres/MSSQL with (options.returning=true), it returns record and created boolean with signature `<Model, created>`.
   */
  static upsert(values, options) {
-    options = _.extend({
+    options = Object.assign({
-      hooks: true
+      hooks: true,
+      returning: false
    }, Utils.cloneDeep(options || {}));
    options.model = this;
    const createdAtAttr = this._timestampAttributes.createdAt;
    const updatedAtAttr = this._timestampAttributes.updatedAt;
-    const hadPrimary = this.primaryKeyField in values || this.primaryKeyAttribute in values;
+    const hasPrimary = this.primaryKeyField in values || this.primaryKeyAttribute in values;
    const instance = this.build(values);
    if (!options.fields) {
@@ -2256,7 +2259,7 @@ class Model {
      // Build adds a null value for the primary key, if none was given by the user.
      // We need to remove that because of some Postgres technicalities.
-      if (!hadPrimary && this.primaryKeyAttribute && !this.rawAttributes[this.primaryKeyAttribute].defaultValue) {
+      if (!hasPrimary && this.primaryKeyAttribute && !this.rawAttributes[this.primaryKeyAttribute].defaultValue) {
        delete insertValues[this.primaryKeyField];
        delete updateValues[this.primaryKeyField];
      }
@@ -2267,6 +2270,12 @@ class Model {
        }
      }).then(() => {
        return this.QueryInterface.upsert(this.getTableName(options), insertValues, updateValues, instance.where(), this, options);
+      }).spread((created, primaryKey) => {
+        if (options.returning === true && primaryKey) {
+          return this.findById(primaryKey, options).then(record => [record, created]);
+        }
+        return created;
      }).tap(result => {
        if (options.hooks) {
          return this.runHooks('afterUpsert', result, options);

--- a/lib/query-interface.js
+++ b/lib/query-interface.js
@@ -904,9 +904,21 @@ class QueryInterface {
    });
  }
-  upsert(tableName, valuesByField, updateValues, where, model, options) {
+  /**
+   * Upsert
+   *
+   * @param {String} tableName
+   * @param {Object} insertValues values to be inserted, mapped to field name
+   * @param {Object} updateValues values to be updated, mapped to field name
+   * @param {Object} where        various conditions
+   * @param {Model}  model
+   * @param {Object} options
+   *
+   * @returns {Promise<created, primaryKey>}
+   */
+  upsert(tableName, insertValues, updateValues, where, model, options) {
    const wheres = [];
-    const attributes = Object.keys(valuesByField);
+    const attributes = Object.keys(insertValues);
    let indexes = [];
    let indexFields;
@@ -938,7 +950,7 @@ class QueryInterface {
      if (_.intersection(attributes, index).length === index.length) {
        where = {};
        for (const field of index) {
-          where[field] = valuesByField[field];
+          where[field] = insertValues[field];
        }
        wheres.push(where);
      }
@@ -949,15 +961,26 @@ class QueryInterface {
    options.type = QueryTypes.UPSERT;
    options.raw = true;
-    const sql = this.QueryGenerator.upsertQuery(tableName, valuesByField, updateValues, where, model, options);
+    const sql = this.QueryGenerator.upsertQuery(tableName, insertValues, updateValues, where, model, options);
-    return this.sequelize.query(sql, options).then(rowCount => {
+    return this.sequelize.query(sql, options).then(result => {
-      if (rowCount === undefined) {
+      switch (this.sequelize.options.dialect) {
-        return rowCount;
+        case 'postgres':
+          return [result.created, result.primary_key];
+        case 'mssql':
+          return [
+            result.$action === 'INSERT',
+            result[model.primaryKeyField]
+          ];
+        // MySQL returns 1 for inserted, 2 for updated
+        // http://dev.mysql.com/doc/refman/5.0/en/insert-on-duplicate.html.
+        case 'mysql':
+          return [result === 1, undefined];
+        default:
+          return [result, undefined];
      }
-      // MySQL returns 1 for inserted, 2 for updated http://dev.mysql.com/doc/refman/5.0/en/insert-on-duplicate.html. Postgres has been modded to do the same
-      return rowCount === 1;
    });
  }

--- a/test/integration/model/upsert.test.js
+++ b/test/integration/model/upsert.test.js
@@ -15,6 +15,14 @@ describe(Support.getTestDialectTeaser('Model'), () => {
    this.clock = sinon.useFakeTimers();
  });
+  after(function() {
+    this.clock.restore();
+  });
+  beforeEach(function() {
+    this.clock.reset();
+  });
  beforeEach(function() {
    this.User = this.sequelize.define('user', {
      username: DataTypes.STRING,
@@ -50,10 +58,6 @@ describe(Support.getTestDialectTeaser('Model'), () => {
    return this.sequelize.sync({ force: true });
  });
-  after(function() {
-    this.clock.restore();
-  });
  if (current.dialect.supports.upserts) {
    describe('upsert', () => {
      it('works with upsert on id', function() {
@@ -154,7 +158,6 @@ describe(Support.getTestDialectTeaser('Model'), () => {
            expect(created2).to.be.ok;
          }
          this.clock.tick(1000);
          // Update the first one
          return User.upsert({ a: 'a', b: 'b', username: 'doe' });
@@ -201,7 +204,6 @@ describe(Support.getTestDialectTeaser('Model'), () => {
            expect(created).to.be.ok;
          }
          this.clock.tick(1000);
          return this.User.upsert({ id: 42, username: 'doe', blob: new Buffer('andrea') });
        }).then(function(created) {
@@ -250,7 +252,6 @@ describe(Support.getTestDialectTeaser('Model'), () => {
            expect(created).to.be.ok;
          }
          this.clock.tick(1000);
          return this.ModelWithFieldPK.upsert({ userId: 42, foo: 'second' });
        }).then(function(created) {
@@ -274,7 +275,6 @@ describe(Support.getTestDialectTeaser('Model'), () => {
            expect(created).to.be.ok;
          }
          this.clock.tick(1000);
          return this.User.upsert({ id: 42, username: 'doe', foo: this.sequelize.fn('upper', 'mixedCase2') });
        }).then(function(created) {
@@ -525,6 +525,79 @@ describe(Support.getTestDialectTeaser('Model'), () => {
          });
        });
      }
+      if (current.dialect.supports.returnValues) {
+        describe('with returning option', () => {
+          it('works with upsert on id', function() {
+            return this.User.upsert({ id: 42, username: 'john' }, { returning: true }).spread((user, created) => {
+              expect(user.get('id')).to.equal(42);
+              expect(user.get('username')).to.equal('john');
+              expect(created).to.be.true;
+              return this.User.upsert({ id: 42, username: 'doe' }, { returning: true });
+            }).spread((user, created) => {
+              expect(user.get('id')).to.equal(42);
+              expect(user.get('username')).to.equal('doe');
+              expect(created).to.be.false;
+            });
+          });
+          it('works for table with custom primary key field', function() {
+            const User = this.sequelize.define('User', {
+              id: {
+                type: DataTypes.INTEGER,
+                autoIncrement: true,
+                primaryKey: true,
+                field: 'id_the_primary'
+              },
+              username: {
+                type: DataTypes.STRING
+              }
+            });
+            return User.sync({ force: true }).then(() => {
+              return User.upsert({ id: 42, username: 'john' }, { returning: true });
+            }).spread((user, created) => {
+              expect(user.get('id')).to.equal(42);
+              expect(user.get('username')).to.equal('john');
+              expect(created).to.be.true;
+              return User.upsert({ id: 42, username: 'doe' }, { returning: true });
+            }).spread((user, created) => {
+              expect(user.get('id')).to.equal(42);
+              expect(user.get('username')).to.equal('doe');
+              expect(created).to.be.false;
+            });
+          });
+          it('works for non incrementing primaryKey', function() {
+            const User = this.sequelize.define('User', {
+              id: {
+                type: DataTypes.STRING,
+                primaryKey: true,
+                field: 'id_the_primary'
+              },
+              username: {
+                type: DataTypes.STRING
+              }
+            });
+            return User.sync({ force: true }).then(() => {
+              return User.upsert({ id: 'surya', username: 'john' }, { returning: true });
+            }).spread((user, created) => {
+              expect(user.get('id')).to.equal('surya');
+              expect(user.get('username')).to.equal('john');
+              expect(created).to.be.true;
+              return User.upsert({ id: 'surya', username: 'doe' }, { returning: true });
+            }).spread((user, created) => {
+              expect(user.get('id')).to.equal('surya');
+              expect(user.get('username')).to.equal('doe');
+              expect(created).to.be.false;
+            });
+          });
+        });
+      }
    });
  }
 });
--- a/test/unit/model/upsert.test.js
+++ b/test/unit/model/upsert.test.js
@@ -9,96 +9,88 @@ const chai = require('chai'),
  DataTypes = require('../../../lib/data-types');
 describe(Support.getTestDialectTeaser('Model'), () => {
  if (current.dialect.supports.upserts) {
    describe('method upsert', () => {
-      const self = this;
+      before(function() {
-      const User = current.define('User', {
+        this.User = current.define('User', {
-        name: DataTypes.STRING,
+          name: DataTypes.STRING,
-        virtualValue: {
+          virtualValue: {
-          type: DataTypes.VIRTUAL,
+            type: DataTypes.VIRTUAL,
-          set(val) {
+            set(val) {
-            return this.value = val;
+              return this.value = val;
+            },
+            get() {
+              return this.value;
+            }
+          },
+          value: DataTypes.STRING,
+          secretValue: {
+            type: DataTypes.INTEGER,
+            allowNull: false
          },
-          get() {
+          createdAt: {
-            return this.value;
+            type: DataTypes.DATE,
+            field: 'created_at'
          }
-        },
+        });
-        value: DataTypes.STRING,
-        secretValue: {
-          type: DataTypes.INTEGER,
-          allowNull: false
-        },
-        createdAt: {
-          type: DataTypes.DATE,
-          field: 'created_at'
-        }
-      });
-      const UserNoTime = current.define('UserNoTime', {
-        name: DataTypes.STRING
-      }, {
-        timestamps: false
-      });
-      before(function() {
-        this.query = current.query;
-        current.query = sinon.stub().returns(Promise.resolve());
-        self.stub = sinon.stub(current.getQueryInterface(), 'upsert').callsFake(() => {
+        this.UserNoTime = current.define('UserNoTime', {
-          return User.build({});
+          name: DataTypes.STRING
+        }, {
+          timestamps: false
        });
      });
-      beforeEach(() => {
+      beforeEach(function() {
-        self.stub.reset();
+        this.sinon = sinon.sandbox.create();
-      });
-      after(function() {
+        this.query = this.sinon.stub(current, 'query').returns(Promise.resolve());
-        current.query = this.query;
+        this.stub = this.sinon.stub(current.getQueryInterface(), 'upsert').returns(Promise.resolve([true, undefined]));
-        self.stub.restore();
      });
+      afterEach(function() {
+        this.sinon.restore();
+      });
-      it('skip validations for missing fields', () => {
+      it('skip validations for missing fields', function() {
-        return expect(User.upsert({
+        return expect(this.User.upsert({
          name: 'Grumpy Cat'
        })).not.to.be.rejectedWith(current.ValidationError);
      });
-      it('creates new record with correct field names', () => {
+      it('creates new record with correct field names', function() {
-        return User
+        return this.User
          .upsert({
            name: 'Young Cat',
            virtualValue: 999
          })
          .then(() => {
-            expect(Object.keys(self.stub.getCall(0).args[1])).to.deep.equal([
+            expect(Object.keys(this.stub.getCall(0).args[1])).to.deep.equal([
              'name', 'value', 'created_at', 'updatedAt'
            ]);
          });
      });
-      it('creates new record with timestamps disabled', () => {
+      it('creates new record with timestamps disabled', function() {
-        return UserNoTime
+        return this.UserNoTime
          .upsert({
            name: 'Young Cat'
          })
          .then(() => {
-            expect(Object.keys(self.stub.getCall(0).args[1])).to.deep.equal([
+            expect(Object.keys(this.stub.getCall(0).args[1])).to.deep.equal([
              'name'
            ]);
          });
      });
-      it('updates all changed fields by default', () => {
+      it('updates all changed fields by default', function() {
-        return User
+        return this.User
          .upsert({
            name: 'Old Cat',
            virtualValue: 111
          })
          .then(() => {
-            expect(Object.keys(self.stub.getCall(0).args[2])).to.deep.equal([
+            expect(Object.keys(this.stub.getCall(0).args[2])).to.deep.equal([
              'name', 'value', 'updatedAt'
            ]);
          });