From 1655b1f565c74a34295782d02c2255550871b631 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Proen=C3=A7a?= <ricardo.proenca@uphold.com>
Date: Fri, 11 Oct 2019 10:27:48 +0100
Subject: [PATCH] fix(postgres): update upsert regex to match the last
 `RETURNING *`

---
 lib/dialects/postgres/query-generator.js      |  8 ++-
 .../dialects/postgres/query-generator.test.js | 50 +++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/lib/dialects/postgres/query-generator.js b/lib/dialects/postgres/query-generator.js
index 53511c017caa..9077f28c4db1 100755
--- 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
index d25043d5360d..9bd43140a1a0 100644
--- 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'],