From 8a2bd24d5d1e75ef6a1e5da95c00deca73dcdfe4 Mon Sep 17 00:00:00 2001
From: Pavan Deolasee <pavan.deolasee@gmail.com>
Date: Mon, 26 Mar 2018 18:56:50 +0530
Subject: [PATCH v27 2/5] Add support for CTE

---
 src/backend/nodes/copyfuncs.c       |   1 +
 src/backend/nodes/equalfuncs.c      |   1 +
 src/backend/nodes/nodeFuncs.c       |   2 +
 src/backend/parser/gram.y           |  11 +--
 src/backend/parser/parse_merge.c    |   9 +++
 src/include/nodes/parsenodes.h      |   1 +
 src/test/regress/expected/merge.out |   3 -
 src/test/regress/expected/with.out  | 132 ++++++++++++++++++++++++++++++++++++
 src/test/regress/sql/with.sql       |  51 ++++++++++++++
 9 files changed, 203 insertions(+), 8 deletions(-)

diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 770ed3b1a8..c3efca3c45 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 5a0151eece..45ceba2830 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 68e2cec66e..7106765e2b 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_SelectStmt:
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index ebca5f3eb7..2f21571915 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 b6e0c46656..26a6692231 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"
@@ -203,6 +204,14 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 
 	qry->commandType = CMD_MERGE;
 
+	/* process the WITH clause independently of all else */
+	if (stmt->withClause)
+	{
+		qry->hasRecursive = stmt->withClause->recursive;
+		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 0c904f4d7f..36e6e2e976 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 05c4287078..411ccd0e66 100644
--- a/src/test/regress/expected/merge.out
+++ b/src/test/regress/expected/merge.out
@@ -1191,9 +1191,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 2a2085556b..543ca4f272 100644
--- a/src/test/regress/expected/with.out
+++ b/src/test/regress/expected/with.out
@@ -1904,6 +1904,138 @@ 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);
+-- 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, m_1.tableoid
+         Hash Cond: (m_1.k = o.k)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.tableoid, 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, m_1.tableoid
+         Hash Cond: (m_1.k = o.k)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.tableoid, 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, m_1.tableoid
+         Hash Cond: (m_1.k = merge_source_cte.a)
+         ->  Seq Scan on public.m m_1
+               Output: m_1.ctid, m_1.tableoid, 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 f85645efde..dd73b334de 100644
--- a/src/test/regress/sql/with.sql
+++ b/src/test/regress/sql/with.sql
@@ -862,6 +862,57 @@ 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);
+
+-- 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;
-- 
2.14.3 (Apple Git-98)

