WITH support in MERGE

Author: Peter Geoghegan
Recursive support removed, no tests
Docs added by me

10 files changed, 232 insertions, 9 deletions

diff --git a/doc/src/sgml/ref/merge.sgml b/doc/src/sgml/ref/merge.sgml
index a7d44a39b68..0ca3d26ff48 100644
--- a/doc/src/sgml/ref/merge.sgml
+++ b/doc/src/sgml/ref/merge.sgml
@@ -18,6 +18,7 @@ PostgreSQL documentation
 
  <refsynopsisdiv>
 <synopsis>
+[ WITH <replaceable class="parameter">with_query</replaceable> [, ...] ]
 MERGE INTO <replaceable class="parameter">target_table_name</replaceable> [ [ AS ] <replaceable class="parameter">target_alias</replaceable> ]
 USING <replaceable class="parameter">data_source</replaceable>
 ON <replaceable class="parameter">join_condition</replaceable>
@@ -391,6 +392,18 @@ DELETE
     </listitem>
    </varlistentry>
 
+   <varlistentry>
+    <term><replaceable class="parameter">with_query</replaceable></term>
+    <listitem>
+     <para>
+      The <literal>WITH</literal> clause allows you to specify one or more
+      subqueries that can be referenced by name in the <command>MERGE</command>
+      query. See <xref linkend="queries-with"/> and <xref linkend="sql-select"/>
+      for details.
+     </para>
+    </listitem>
+   </varlistentry>
+
   </variablelist>
  </refsect1>
 
@@ -597,7 +610,7 @@ WHEN MATCHED THEN
     This command conforms to the <acronym>SQL</acronym> standard.
   </para>
    <para>
-    The DO NOTHING action is an extension to the <acronym>SQL</acronym> standard.
+    The WITH clause and DO NOTHING action are extensions to the <acronym>SQL</acronym> standard.
   </para>
  </refsect1>
 </refentry>
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 770ed3b1a88..c3efca3c452 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -3055,6 +3055,7 @@ _copyMergeStmt(const MergeStmt *from)
 	COPY_NODE_FIELD(source_relation);
 	COPY_NODE_FIELD(join_condition);
 	COPY_NODE_FIELD(mergeActionList);
+	COPY_NODE_FIELD(withClause);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
index 5a0151eece5..45ceba28309 100644
--- a/src/backend/nodes/equalfuncs.c
+++ b/src/backend/nodes/equalfuncs.c
@@ -1051,6 +1051,7 @@ _equalMergeStmt(const MergeStmt *a, const MergeStmt *b)
 	COMPARE_NODE_FIELD(source_relation);
 	COMPARE_NODE_FIELD(join_condition);
 	COMPARE_NODE_FIELD(mergeActionList);
+	COMPARE_NODE_FIELD(withClause);
 
 	return true;
 }
diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index 3c302db0575..4157e7eb9a3 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -3446,6 +3446,8 @@ raw_expression_tree_walker(Node *node,
 					return true;
 				if (walker(stmt->mergeActionList, context))
 					return true;
+				if (walker(stmt->withClause, context))
+					return true;
 			}
 			break;
 		case T_MergeAction:
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 583ee321e1d..b879358de16 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -11105,17 +11105,18 @@ set_target_list:
  *****************************************************************************/
 
 MergeStmt:
-			MERGE INTO relation_expr_opt_alias
+			opt_with_clause MERGE INTO relation_expr_opt_alias
 			USING table_ref
 			ON a_expr
 			merge_when_list
 				{
 					MergeStmt *m = makeNode(MergeStmt);
 
-					m->relation = $3;
-					m->source_relation = $5;
-					m->join_condition = $7;
-					m->mergeActionList = $8;
+					m->withClause = $1;
+					m->relation = $4;
+					m->source_relation = $6;
+					m->join_condition = $8;
+					m->mergeActionList = $9;
 
 					$$ = (Node *)m;
 				}
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index d6d6876961c..eb4c615ce1c 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -24,6 +24,7 @@
 #include "parser/parsetree.h"
 #include "parser/parser.h"
 #include "parser/parse_clause.h"
+#include "parser/parse_cte.h"
 #include "parser/parse_merge.h"
 #include "parser/parse_relation.h"
 #include "parser/parse_target.h"
@@ -202,6 +203,19 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->hasRecursive = false;
+
+	/* process the WITH clause independently of all else */
+	if (stmt->withClause)
+	{
+		if (stmt->withClause->recursive)
+			ereport(ERROR,
+					(errcode(ERRCODE_SYNTAX_ERROR),
+					 errmsg("WITH RECURSIVE is not supported for MERGE statement")));
+
+		qry->cteList = transformWithClause(pstate, stmt->withClause);
+		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+	}
 
 	/*
 	 * Check WHEN clauses for permissions and sanity
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index d005beeba82..699fa77bc70 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1519,6 +1519,7 @@ typedef struct MergeStmt
 	Node	   *source_relation;	/* source relation */
 	Node	   *join_condition; /* join condition between source and target */
 	List	   *mergeActionList;	/* list of MergeAction(s) */
+	WithClause *withClause;		/* WITH clause */
 } MergeStmt;
 
 typedef struct MergeAction
diff --git a/src/test/regress/expected/merge.out b/src/test/regress/expected/merge.out
index b7e42515e21..389eeedf289 100644
--- a/src/test/regress/expected/merge.out
+++ b/src/test/regress/expected/merge.out
@@ -1210,9 +1210,6 @@ WHEN NOT MATCHED THEN
 WHEN MATCHED AND tid < 2 THEN
 	DELETE
 ;
-ERROR:  syntax error at or near "MERGE"
-LINE 4: MERGE INTO sq_target t
-        ^
 ROLLBACK;
 -- RETURNING
 BEGIN;
diff --git a/src/test/regress/expected/with.out b/src/test/regress/expected/with.out
index 2a2085556bb..350a34d9870 100644
--- a/src/test/regress/expected/with.out
+++ b/src/test/regress/expected/with.out
@@ -1904,6 +1904,143 @@ RETURNING k, v;
 (0 rows)
 
 DROP TABLE withz;
+-- WITH referenced by MERGE statement
+CREATE TABLE m AS SELECT i AS k, (i || ' v')::text v FROM generate_series(1, 16, 3) i;
+ALTER TABLE m ADD UNIQUE (k);
+WITH RECURSIVE cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+ERROR:  WITH RECURSIVE is not supported for MERGE statement
+-- Basic:
+WITH cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+-- Examine
+SELECT * FROM m where k = 0;
+ k |          v           
+---+----------------------
+ 0 | merge source SubPlan
+(1 row)
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+                            QUERY PLAN                             
+-------------------------------------------------------------------
+ Merge on public.m
+   CTE cte_basic
+     ->  Result
+           Output: 1, 'cte_basic val'::text
+   ->  Hash Right Join
+         Output: o.k, o.v, o.*, m_1.ctid
+         Hash Cond: (m_1.k = o.k)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.k
+         ->  Hash
+               Output: o.k, o.v, o.*
+               ->  Subquery Scan on o
+                     Output: o.k, o.v, o.*
+                     ->  Result
+                           Output: 0, 'merge source SubPlan'::text
+   SubPlan 2
+     ->  Limit
+           Output: ((cte_basic.b || ' merge update'::text))
+           ->  CTE Scan on cte_basic
+                 Output: (cte_basic.b || ' merge update'::text)
+                 Filter: (cte_basic.a = m.k)
+(21 rows)
+
+-- InitPlan
+WITH cte_init AS (SELECT 1 a, 'cte_init val' b)
+MERGE INTO m USING (select 1 k, 'merge source InitPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_init WHERE a = 1 LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+-- Examine
+SELECT * FROM m where k = 1;
+ k |             v             
+---+---------------------------
+ 1 | cte_init val merge update
+(1 row)
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH cte_init AS (SELECT 1 a, 'cte_init val' b)
+MERGE INTO m USING (select 1 k, 'merge source InitPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_init WHERE a = 1 LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+                             QUERY PLAN                             
+--------------------------------------------------------------------
+ Merge on public.m
+   CTE cte_init
+     ->  Result
+           Output: 1, 'cte_init val'::text
+   InitPlan 2 (returns $1)
+     ->  Limit
+           Output: ((cte_init.b || ' merge update'::text))
+           ->  CTE Scan on cte_init
+                 Output: (cte_init.b || ' merge update'::text)
+                 Filter: (cte_init.a = 1)
+   ->  Hash Right Join
+         Output: o.k, o.v, o.*, m_1.ctid
+         Hash Cond: (m_1.k = o.k)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.k
+         ->  Hash
+               Output: o.k, o.v, o.*
+               ->  Subquery Scan on o
+                     Output: o.k, o.v, o.*
+                     ->  Result
+                           Output: 1, 'merge source InitPlan'::text
+(21 rows)
+
+-- MERGE source comes from CTE:
+WITH merge_source_cte AS (SELECT 15 a, 'merge_source_cte val' b)
+MERGE INTO m USING (select * from merge_source_cte) o ON m.k=o.a
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || merge_source_cte.*::text || ' merge update' FROM merge_source_cte WHERE a = 15)
+WHEN NOT MATCHED THEN INSERT VALUES(o.a, o.b || (SELECT merge_source_cte.*::text || ' merge insert' FROM merge_source_cte));
+-- Examine
+SELECT * FROM m where k = 15;
+ k  |                              v                               
+----+--------------------------------------------------------------
+ 15 | merge_source_cte val(15,"merge_source_cte val") merge insert
+(1 row)
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH merge_source_cte AS (SELECT 15 a, 'merge_source_cte val' b)
+MERGE INTO m USING (select * from merge_source_cte) o ON m.k=o.a
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || merge_source_cte.*::text || ' merge update' FROM merge_source_cte WHERE a = 15)
+WHEN NOT MATCHED THEN INSERT VALUES(o.a, o.b || (SELECT merge_source_cte.*::text || ' merge insert' FROM merge_source_cte));
+                                                  QUERY PLAN                                                   
+---------------------------------------------------------------------------------------------------------------
+ Merge on public.m
+   CTE merge_source_cte
+     ->  Result
+           Output: 15, 'merge_source_cte val'::text
+   InitPlan 2 (returns $1)
+     ->  CTE Scan on merge_source_cte merge_source_cte_1
+           Output: ((merge_source_cte_1.b || (merge_source_cte_1.*)::text) || ' merge update'::text)
+           Filter: (merge_source_cte_1.a = 15)
+   InitPlan 3 (returns $2)
+     ->  CTE Scan on merge_source_cte merge_source_cte_2
+           Output: ((merge_source_cte_2.*)::text || ' merge insert'::text)
+   ->  Hash Right Join
+         Output: merge_source_cte.a, merge_source_cte.b, ROW(merge_source_cte.a, merge_source_cte.b), m_1.ctid
+         Hash Cond: (m_1.k = merge_source_cte.a)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.k
+         ->  Hash
+               Output: merge_source_cte.a, merge_source_cte.b
+               ->  CTE Scan on merge_source_cte
+                     Output: merge_source_cte.a, merge_source_cte.b
+(20 rows)
+
+DROP TABLE m;
 -- check that run to completion happens in proper ordering
 TRUNCATE TABLE y;
 INSERT INTO y SELECT generate_series(1, 3);
diff --git a/src/test/regress/sql/with.sql b/src/test/regress/sql/with.sql
index f85645efdee..c6b197c3275 100644
--- a/src/test/regress/sql/with.sql
+++ b/src/test/regress/sql/with.sql
@@ -862,6 +862,62 @@ RETURNING k, v;
 
 DROP TABLE withz;
 
+-- WITH referenced by MERGE statement
+CREATE TABLE m AS SELECT i AS k, (i || ' v')::text v FROM generate_series(1, 16, 3) i;
+ALTER TABLE m ADD UNIQUE (k);
+
+WITH RECURSIVE cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+
+-- Basic:
+WITH cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+-- Examine
+SELECT * FROM m where k = 0;
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH cte_basic AS (SELECT 1 a, 'cte_basic val' b)
+MERGE INTO m USING (select 0 k, 'merge source SubPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_basic WHERE cte_basic.a = m.k LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+
+-- InitPlan
+WITH cte_init AS (SELECT 1 a, 'cte_init val' b)
+MERGE INTO m USING (select 1 k, 'merge source InitPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_init WHERE a = 1 LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+-- Examine
+SELECT * FROM m where k = 1;
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH cte_init AS (SELECT 1 a, 'cte_init val' b)
+MERGE INTO m USING (select 1 k, 'merge source InitPlan' v) o ON m.k=o.k
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || ' merge update' FROM cte_init WHERE a = 1 LIMIT 1)
+WHEN NOT MATCHED THEN INSERT VALUES(o.k, o.v);
+
+-- MERGE source comes from CTE:
+WITH merge_source_cte AS (SELECT 15 a, 'merge_source_cte val' b)
+MERGE INTO m USING (select * from merge_source_cte) o ON m.k=o.a
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || merge_source_cte.*::text || ' merge update' FROM merge_source_cte WHERE a = 15)
+WHEN NOT MATCHED THEN INSERT VALUES(o.a, o.b || (SELECT merge_source_cte.*::text || ' merge insert' FROM merge_source_cte));
+-- Examine
+SELECT * FROM m where k = 15;
+
+-- See EXPLAIN output for same query:
+EXPLAIN (VERBOSE, COSTS OFF)
+WITH merge_source_cte AS (SELECT 15 a, 'merge_source_cte val' b)
+MERGE INTO m USING (select * from merge_source_cte) o ON m.k=o.a
+WHEN MATCHED THEN UPDATE SET v = (SELECT b || merge_source_cte.*::text || ' merge update' FROM merge_source_cte WHERE a = 15)
+WHEN NOT MATCHED THEN INSERT VALUES(o.a, o.b || (SELECT merge_source_cte.*::text || ' merge insert' FROM merge_source_cte));
+
+DROP TABLE m;
+
 -- check that run to completion happens in proper ordering
 
 TRUNCATE TABLE y;