From 94d4c8aa6fbce7015cce3c029a19e81b75818e6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mikael=20Lepist=C3=B6?= <mikael.lepisto@vincit.fi>
Date: Mon, 30 Oct 2017 01:41:43 +0200
Subject: [PATCH] Allow update queries in with statements. Fixes #2263.

---
 src/formatter.js           |  9 ++++++++-
 src/query/builder.js       | 30 +++++++++---------------------
 src/query/compiler.js      |  5 -----
 test/unit/query/builder.js | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/src/formatter.js b/src/formatter.js
index b54d2d38c8..61f40caa1b 100644
--- a/src/formatter.js
+++ b/src/formatter.js
@@ -87,6 +87,13 @@ export default class Formatter {
     }
   }
 
+  /**
+   * Creates SQL for a parameter, which might be passed to where() or .with() or
+   * pretty much anywhere in API.
+   *
+   * @param query Callback (for where or complete builder), Raw or QueryBuilder
+   * @param method Optional at least 'select' or 'update' are valid
+   */
   rawOrFn(value, method) {
     if (typeof value === 'function') {
       return this.outputQuery(this.compileCallback(value, method));
@@ -149,7 +156,7 @@ export default class Formatter {
     compiler.formatter = this;
 
     // Return the compiled & parameterized sql.
-    return compiler.toSQL(method || 'select');
+    return compiler.toSQL(method || builder._method || 'select');
   }
 
   // Ensures the query is aliased if necessary.
diff --git a/src/query/builder.js b/src/query/builder.js
index b033d5a0c9..0b5f7ea7e2 100644
--- a/src/query/builder.js
+++ b/src/query/builder.js
@@ -71,39 +71,27 @@ assign(Builder.prototype, {
   // With
   // ------
 
-  with(alias, statement, bindings) {
+  with(alias, statement) {
     if(typeof alias !== 'string') {
       throw new Error('with() first argument must be a string');
     }
-    if (typeof statement === 'function') {
+    if (
+      typeof statement === 'function' ||
+      statement instanceof Builder ||
+      statement instanceof Raw
+    ) {
       return this.withWrapped(alias, statement);
     }
-    // Allow a raw statement to be passed along to the query.
-    if (statement instanceof Raw && arguments.length >= 2) {
-      return this.withRaw(alias, statement, bindings);
-    }
-    throw new Error('with() second argument must be a function or a raw');
-  },
-
-  // Adds a raw `with` clause to the query.
-  withRaw(alias, sql, bindings) {
-    const raw = (sql instanceof Raw ? sql : this.client.raw(sql, bindings));
-    this._statements.push({
-      grouping: 'with',
-      type: 'withRaw',
-      alias: alias,
-      value: raw
-    });
-    return this;
+    throw new Error('with() second argument must be a function / QueryBuilder or a raw');
   },
 
   // Helper for compiling any advanced `with` queries.
-  withWrapped(alias, callback) {
+  withWrapped(alias, query) {
     this._statements.push({
       grouping: 'with',
       type: 'withWrapped',
       alias: alias,
-      value: callback
+      value: query
     });
     return this;
   },
diff --git a/src/query/compiler.js b/src/query/compiler.js
index 86379c9118..cf4b7f1b5f 100644
--- a/src/query/compiler.js
+++ b/src/query/compiler.js
@@ -556,11 +556,6 @@ assign(QueryCompiler.prototype, {
     return val && this.formatter.columnize(statement.alias) + ' as (' + val + ')' || '';
   },
 
-  withRaw(statement) {
-    return this.formatter.columnize(statement.alias) + ' as (' +
-      this.formatter.unwrapRaw(statement.value) + ')';
-  },
-
   // Determines whether to add a "not" prefix to the where clause.
   _not(statement, str) {
     if (statement.not) return `not ${str}`;
diff --git a/test/unit/query/builder.js b/test/unit/query/builder.js
index f0d3fca5e6..3e97f47ffd 100644
--- a/test/unit/query/builder.js
+++ b/test/unit/query/builder.js
@@ -4587,6 +4587,44 @@ describe("QueryBuilder", function() {
     });
   });
 
+  describe("#2263, update / delete queries in with syntax", () => {
+    it("with update query passed as raw", () => {
+      testquery(qb().with('update1', raw('??', [qb().from('accounts').update({ name: 'foo' })])).from('accounts'), {
+        postgres: `with "update1" as (update "accounts" set "name" = 'foo') select * from "accounts"`,
+      });
+    });
+
+    it("with update query passed as query builder", () => {
+      testquery(qb().with('update1', qb().from('accounts').update({ name: 'foo' })).from('accounts'), {
+       postgres: `with "update1" as (update "accounts" set "name" = 'foo') select * from "accounts"`,
+      });
+    });
+
+    it("with update query passed as callback", () => {
+      testquery(qb().with('update1', builder => builder.from('accounts').update({ name: 'foo' })).from('accounts'), {
+        postgres: `with "update1" as (update "accounts" set "name" = 'foo') select * from "accounts"`,
+      });
+    });
+
+    it("with delete query passed as raw", () => {
+      testquery(qb().with('delete1', raw('??', [qb().delete().from('accounts').where('id', 1)])).from('accounts'), {
+        postgres: `with "delete1" as (delete from "accounts" where "id" = 1) select * from "accounts"`,
+      });
+    });
+
+    it("with delete query passed as query builder", () => {
+      testquery(qb().with('delete1', builder => builder.delete().from('accounts').where('id', 1)).from('accounts'), {
+        postgres: `with "delete1" as (delete from "accounts" where "id" = 1) select * from "accounts"`,
+      });
+    });
+
+    it("with delete query passed as callback", () => {
+      testquery(qb().with('delete1', qb().delete().from('accounts').where('id', 1)).from('accounts'), {
+        postgres: `with "delete1" as (delete from "accounts" where "id" = 1) select * from "accounts"`,
+      });
+    });
+  });
+
   it('#1710, properly escapes arrays in where clauses in postgresql', function() {
     testquery(qb().select('*').from('sometable').where('array_field', '&&', [7]), {
       postgres: "select * from \"sometable\" where \"array_field\" && '{7}'"