Match bulkCreate upsert results to instances by conflict key

.github/workflows/ignore_tests/shared/model_bulk_create.txt

@@ -1,4 +1,3 @@
 should make the auto incremented values available on the returned instances
 should make the auto incremented values available on the returned instances with custom fields
 should return auto increment primary key values
-should upsert existing memberships

source/patches-v6.js

@@ -37,3 +37,113 @@ PostgresQueryInterface.prototype.removeConstraint = async function (
     } else throw error;
   }
 };
+
+// Model.bulkCreate assigns the rows returned by bulkInsert to its instances by
+// position. CockroachDB does not return the rows of
+// INSERT ... ON CONFLICT DO UPDATE ... RETURNING in the order of the VALUES
+// list: the conflict check is planned as a join against the conflict index,
+// and rows that were updated can be returned before rows that were inserted.
+// Reorder the returned rows so that they line up with the input records again,
+// using the conflict target (upsertKeys) to identify each row.
+PostgresQueryInterface.prototype.__bulkInsert =
+  PostgresQueryInterface.prototype.bulkInsert;
+
+PostgresQueryInterface.prototype.bulkInsert = async function (
+  tableName,
+  records,
+  options,
+  attributes
+) {
+  const results = await this.__bulkInsert(
+    tableName,
+    records,
+    options,
+    attributes
+  );
+  if (
+    !options ||
+    !options.updateOnDuplicate ||
+    !Array.isArray(options.upsertKeys) ||
+    options.upsertKeys.length === 0
+  ) {
+    return results;
+  }
+  return matchRowsToRecords(results, records, options.upsertKeys, attributes);
+};
+
+// Returns a string that identifies a value of a conflict target column, or
+// undefined if the value cannot be compared with the value that the database
+// returns (e.g. Sequelize.fn or Sequelize.literal).
+function conflictKeyPart(value, attribute) {
+  if (value instanceof Date) return `date:${value.getTime()}`;
+  if (Buffer.isBuffer(value)) return `bytes:${value.toString('hex')}`;
+  if (typeof value === 'string') {
+    // CockroachDB returns UUIDs in lowercase.
+    const isUUID = attribute && attribute.type && attribute.type.key === 'UUID';
+    return isUUID ? value.toLowerCase() : value;
+  }
+  if (['number', 'bigint', 'boolean'].includes(typeof value)) {
+    return String(value);
+  }
+  return undefined;
+}
+
+// Returns the conflict key of a record or returned row, null if the row can
+// never conflict (a conflict target column is NULL or DEFAULT), or undefined if
+// the key cannot be computed.
+function conflictKey(row, upsertKeys, attributes) {
+  const parts = [];
+  for (const key of upsertKeys) {
+    const value = row[key];
+    if (value === null || value === undefined) return null;
+    const part = conflictKeyPart(value, attributes && attributes[key]);
+    if (part === undefined) return undefined;
+    parts.push(part);
+  }
+  return JSON.stringify(parts);
+}
+
+// Returns the rows reordered to match the records, or the rows unchanged if
+// they cannot be matched unambiguously.
+function matchRowsToRecords(rows, records, upsertKeys, attributes) {
+  if (!Array.isArray(rows) || rows.length !== records.length) return rows;
+
+  // Records whose conflict key is NULL or DEFAULT are always inserted. Their
+  // returned rows are matched in order, since CockroachDB returns inserted
+  // rows in the order of the VALUES list.
+  const unkeyedRecords = [];
+  const recordsByKey = new Map();
+  for (let i = 0; i < records.length; i++) {
+    const key = conflictKey(records[i], upsertKeys, attributes);
+    if (key === undefined) return rows;
+    if (key === null) {
+      unkeyedRecords.push(i);
+    } else {
+      if (!recordsByKey.has(key)) recordsByKey.set(key, []);
+      // Duplicate keys are possible with a partial unique index, for records
+      // that fall outside of the index predicate and are always inserted.
+      recordsByKey.get(key).push(i);
+    }
+  }
+
+  const matched = new Array(records.length);
+  for (const row of rows) {
+    const key = conflictKey(row, upsertKeys, attributes);
+    const candidates = recordsByKey.get(key);
+    let i;
+    if (candidates && candidates.length > 0) {
+      i = candidates.shift();
+    } else if (unkeyedRecords.length > 0) {
+      i = unkeyedRecords.shift();
+    } else {
+      return rows;
+    }
+    matched[i] = row;
+  }
+  // If a keyed record did not get its own row, some row was matched to the
+  // wrong record, e.g. because the database normalized the key value.
+  for (const candidates of recordsByKey.values()) {
+    if (candidates.length > 0) return rows;
+  }
+  return matched;
+}

tests/model_bulk_create_test.js

@@ -171,5 +171,145 @@ describe('Model', () => {
         expect(ms[0].id < ms[1].id).to.be.true;
       });
     });
+
+    // CockroachDB does not return the rows of INSERT ... ON CONFLICT DO UPDATE
+    // ... RETURNING in the order of the VALUES list; rows that were updated
+    // can come back before rows that were inserted. Sequelize assigns the
+    // returned rows to instances by position, so the adapter must reorder them.
+    describe('updateOnDuplicate return values', () => {
+      // Checks that each returned instance matches both its input and the
+      // row that is stored in the database under the instance's primary key.
+      async function expectInstancesMatchRows(Model, input, results, fields) {
+        expect(results).to.have.length(input.length);
+        const pk = Model.primaryKeyAttribute;
+        const rows = await Model.findAll({ paranoid: false });
+        const rowsByPk = new Map(rows.map(row => [String(row[pk]), row]));
+        results.forEach((result, i) => {
+          for (const field of fields) {
+            expect(result[field]).to.eql(input[i][field]);
+          }
+          const row = rowsByPk.get(String(result[pk]));
+          expect(row, `row for ${pk}=${result[pk]}`).to.be.ok;
+          for (const field of fields) {
+            expect(row[field]).to.eql(result[field]);
+          }
+        });
+      }
+
+      it('should match rows to instances when the primary key conflicts', async function () {
+        await this.Student.bulkCreate([
+          { no: 2, name: 'old 2' },
+          { no: 4, name: 'old 4' }
+        ]);
+
+        const input = [1, 2, 3, 4, 5, 6].map(no => ({ no, name: `new ${no}` }));
+        const results = await this.Student.bulkCreate(input, {
+          updateOnDuplicate: ['name']
+        });
+
+        await expectInstancesMatchRows(this.Student, input, results, [
+          'no',
+          'name'
+        ]);
+      });
+
+      it('should match rows to instances when a unique key conflicts', async function () {
+        const existing = await this.User.bulkCreate([
+          { uniqueName: 'b', username: 'old b' },
+          { uniqueName: 'd', username: 'old d' }
+        ]);
+
+        const input = ['a', 'b', 'c', 'd', 'e'].map(uniqueName => ({
+          uniqueName,
+          username: `new ${uniqueName}`
+        }));
+        const results = await this.User.bulkCreate(input, {
+          updateOnDuplicate: ['username']
+        });
+
+        await expectInstancesMatchRows(this.User, input, results, [
+          'uniqueName',
+          'username'
+        ]);
+        // Updated rows keep their original primary keys.
+        expect(results[1].id).to.eql(existing[0].id);
+        expect(results[3].id).to.eql(existing[1].id);
+      });
+
+      it('should match rows to instances when some rows omit the conflict key', async function () {
+        const [existing] = await this.Account.bulkCreate([
+          { accountName: 'old' }
+        ]);
+
+        const input = [
+          { accountName: 'new 1' },
+          { id: existing.id, accountName: 'updated' },
+          { accountName: 'new 2' },
+          { accountName: 'new 3' }
+        ];
+        const results = await this.Account.bulkCreate(input, {
+          updateOnDuplicate: ['accountName']
+        });
+
+        await expectInstancesMatchRows(this.Account, input, results, [
+          'accountName'
+        ]);
+        expect(results[1].id).to.eql(existing.id);
+      });
+
+      it('should match rows to instances with a partial unique index and conflictWhere', async function () {
+        const Memberships = this.sequelize.define(
+          'memberships',
+          {
+            user_id: DataTypes.INTEGER,
+            foreign_id: DataTypes.INTEGER,
+            time_deleted: DataTypes.DATE
+          },
+          {
+            createdAt: false,
+            updatedAt: false,
+            deletedAt: 'time_deleted',
+            indexes: [
+              {
+                fields: ['user_id', 'foreign_id'],
+                unique: true,
+                where: { time_deleted: null }
+              }
+            ]
+          }
+        );
+        await Memberships.sync({ force: true });
+        const options = {
+          conflictWhere: { time_deleted: null },
+          conflictAttributes: ['user_id', 'foreign_id'],
+          updateOnDuplicate: ['user_id', 'foreign_id', 'time_deleted']
+        };
+
+        // Odd rows are soft-deleted, so they are not covered by the partial
+        // unique index and will be inserted again below.
+        await Memberships.bulkCreate(
+          new Array(10).fill().map((_, i) => ({
+            user_id: i + 1,
+            foreign_id: i + 20,
+            time_deleted: i % 2 ? new Date() : null
+          })),
+          options
+        );
+
+        const input = new Array(10).fill().map((_, i) => ({
+          user_id: i + 1,
+          foreign_id: i + 20,
+          time_deleted: null
+        }));
+        const results = await Memberships.bulkCreate(input, options);
+
+        await expectInstancesMatchRows(Memberships, input, results, [
+          'user_id',
+          'foreign_id',
+          'time_deleted'
+        ]);
+        expect(await Memberships.count({ paranoid: false })).to.eq(15);
+      });
+    });
   });
 });