fix(postgres): update upsert regex to match the last RETURNING *

sequelize · Oct 12, 2019 · 1655b1f · 1655b1f
1 parent bd59b87
commit 1655b1f
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 2 deletions.
diff --git a/lib/dialects/postgres/query-generator.js b/lib/dialects/postgres/query-generator.js
@@ -347,8 +347,12 @@ class PostgresQueryGenerator extends AbstractQueryGenerator {
     const insert = this.insertQuery(tableName, insertValues, model.rawAttributes, upsertOptions);
     const update = this.updateQuery(tableName, updateValues, where, upsertOptions, model.rawAttributes);
 
-    insert.query = insert.query.replace('RETURNING *', `RETURNING ${primaryField} INTO primary_key`);
-    update.query = update.query.replace('RETURNING *', `RETURNING ${primaryField} INTO primary_key`);
+    if (options.returning) {
+      const returningRegex = /RETURNING \*(?![\s\S]*RETURNING \*)/;
+
+      insert.query = insert.query.replace(returningRegex, `RETURNING ${primaryField} INTO primary_key`);
+      update.query = update.query.replace(returningRegex, `RETURNING ${primaryField} INTO primary_key`);
+    }
 
     return this.exceptionFn(
       'sequelize_upsert',

diff --git a/test/unit/dialects/postgres/query-generator.test.js b/test/unit/dialects/postgres/query-generator.test.js
@@ -44,12 +44,14 @@ if (dialect.startsWith('postgres')) {
           expectation: 'CREATE DATABASE "myDatabase" ENCODING = \'UTF8\' LC_COLLATE = \'en_US.UTF-8\' LC_CTYPE = \'zh_TW.UTF-8\' TEMPLATE = \'template0\';'
         }
       ],
+
       dropDatabaseQuery: [
         {
           arguments: ['myDatabase'],
           expectation: 'DROP DATABASE IF EXISTS "myDatabase";'
         }
       ],
+
       arithmeticQuery: [
         {
           title: 'Should use the plus operator',
@@ -87,6 +89,7 @@ if (dialect.startsWith('postgres')) {
           expectation: 'UPDATE "myTable" SET "foo"="foo"- \'bar\''
         }
       ],
+
       attributesToSQL: [
         {
           arguments: [{ id: 'INTEGER' }],
@@ -1091,6 +1094,53 @@ if (dialect.startsWith('postgres')) {
         }
       ],
 
+      upsertQuery: [
+        {
+          arguments: [
+            'myTable',
+            { name: 'foo' },
+            { name: 'foo' },
+            { id: 2 },
+            { primaryKeyField: 'id' },
+            { returning: false }
+          ],
+          expectation: 'CREATE OR REPLACE FUNCTION pg_temp.sequelize_upsert(OUT created boolean, OUT primary_key text)  AS $func$ BEGIN INSERT INTO "myTable" ("name") VALUES (\'foo\'); created := true; EXCEPTION WHEN unique_violation THEN UPDATE "myTable" SET "name"=\'foo\' WHERE "id" = 2; created := false; END; $func$ LANGUAGE plpgsql; SELECT * FROM pg_temp.sequelize_upsert();'
+        },
+        {
+          arguments: [
+            'myTable',
+            { name: 'RETURNING *', json: '{"foo":"RETURNING *"}' },
+            { name: 'RETURNING *', json: '{"foo":"RETURNING *"}' },
+            { id: 2 },
+            { primaryKeyField: 'id' },
+            { returning: false }
+          ],
+          expectation: 'CREATE OR REPLACE FUNCTION pg_temp.sequelize_upsert(OUT created boolean, OUT primary_key text)  AS $func$ BEGIN INSERT INTO "myTable" ("name","json") VALUES (\'RETURNING *\',\'{"foo":"RETURNING *"}\'); created := true; EXCEPTION WHEN unique_violation THEN UPDATE "myTable" SET "name"=\'RETURNING *\',"json"=\'{"foo":"RETURNING *"}\' WHERE "id" = 2; created := false; END; $func$ LANGUAGE plpgsql; SELECT * FROM pg_temp.sequelize_upsert();'
+        },
+        {
+          arguments: [
+            'myTable',
+            { name: 'foo' },
+            { name: 'foo' },
+            { id: 2 },
+            { primaryKeyField: 'id' },
+            { returning: true }
+          ],
+          expectation: 'CREATE OR REPLACE FUNCTION pg_temp.sequelize_upsert(OUT created boolean, OUT primary_key text)  AS $func$ BEGIN INSERT INTO "myTable" ("name") VALUES (\'foo\') RETURNING "id" INTO primary_key; created := true; EXCEPTION WHEN unique_violation THEN UPDATE "myTable" SET "name"=\'foo\' WHERE "id" = 2 RETURNING "id" INTO primary_key; created := false; END; $func$ LANGUAGE plpgsql; SELECT * FROM pg_temp.sequelize_upsert();'
+        },
+        {
+          arguments: [
+            'myTable',
+            { name: 'RETURNING *', json: '{"foo":"RETURNING *"}' },
+            { name: 'RETURNING *', json: '{"foo":"RETURNING *"}' },
+            { id: 2 },
+            { primaryKeyField: 'id' },
+            { returning: true }
+          ],
+          expectation: 'CREATE OR REPLACE FUNCTION pg_temp.sequelize_upsert(OUT created boolean, OUT primary_key text)  AS $func$ BEGIN INSERT INTO "myTable" ("name","json") VALUES (\'RETURNING *\',\'{"foo":"RETURNING *"}\') RETURNING "id" INTO primary_key; created := true; EXCEPTION WHEN unique_violation THEN UPDATE "myTable" SET "name"=\'RETURNING *\',"json"=\'{"foo":"RETURNING *"}\' WHERE "id" = 2 RETURNING "id" INTO primary_key; created := false; END; $func$ LANGUAGE plpgsql; SELECT * FROM pg_temp.sequelize_upsert();'
+        }
+      ],
+
       removeIndexQuery: [
         {
           arguments: ['User', 'user_foo_bar'],