Subject: MERGE ... RETURNING

From: Dean Rasheed <[email-masked]>
Date: Sun, 8 Jan 2023 12:28:02 +0000

I've been thinking about adding RETURNING support to MERGE in order to
let the user see what changed.

I considered allowing a separate RETURNING list at the end of each
action, but rapidly dismissed that idea. Firstly, it introduces
shift/reduce conflicts to the grammar. These can be resolved by making
the "AS" before column aliases non-optional, but that's pretty ugly,
and there may be a better way. More serious drawbacks are that this
syntax is much more cumbersome for the end user, having to repeat the
RETURNING clause several times, and the implementation is likely to be
pretty complex, so I didn't pursue it.

A much simpler approach is to just have a single RETURNING list at the
end of the command. That's much easier to implement, and easier for
the end user. The main drawback is that it's impossible for the user
to work out from the values returned which action was actually taken,
and I think that's a pretty essential piece of information (at least
it seems pretty limiting to me, not being able to work that out).

So playing around with it (and inspired by the WITH ORDINALITY syntax
for SRFs), I had the idea of allowing "WITH WHEN CLAUSE" at the end of
the returning list, which adds an integer column to the list, whose
value is set to the index of the when clause executed, as in the
attached very rough patch.

So, quoting an example from the tests, this allows things like:

WITH t AS (
  MERGE INTO sq_target t USING v ON tid = sid
    WHEN MATCHED AND tid > 2 THEN UPDATE SET balance = t.balance + delta
    WHEN NOT MATCHED THEN INSERT (balance, tid) VALUES (balance + delta, sid)
    WHEN MATCHED AND tid < 2 THEN DELETE
    RETURNING t.* WITH WHEN CLAUSE
)
SELECT CASE when_clause
         WHEN 1 THEN 'UPDATE'
         WHEN 2 THEN 'INSERT'
         WHEN 3 THEN 'DELETE'
       END, *
FROM t;

  case  | tid | balance | when_clause
--------+-----+---------+-------------
 INSERT |  -1 |     -11 |           2
 DELETE |   1 |     100 |           3
(2 rows)

1 row is returned for each merge action executed (other than DO
NOTHING actions), and as usual, the values represent old target values
for DELETE actions, and new target values for INSERT/UPDATE actions.

It's also possible to return the source values, and a bare "*" in the
returning list expands to all the source columns, followed by all the
target columns.

The name of the added column, if included, can be changed by
specifying "WITH WHEN CLAUSE [AS] col_alias". I chose the syntax "WHEN
CLAUSE" and "when_clause" as the default column name because those
match the existing terminology used in the docs.

Anyway, this feels like a good point to stop playing around and get
feedback on whether this seems useful, or if anyone has other ideas.

Regards,
Dean
diff --git a/doc/src/sgml/ref/merge.sgml b/doc/src/sgml/ref/merge.sgml
new file mode 100644
index 0995fe0..4fc0a65
--- a/doc/src/sgml/ref/merge.sgml
+++ b/doc/src/sgml/ref/merge.sgml
@@ -25,6 +25,8 @@ PostgreSQL documentation
 MERGE INTO [ ONLY ] <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>
 <replaceable class="parameter">when_clause</replaceable> [...]
+[ RETURNING * | <replaceable class="parameter">output_expression</replaceable> [ [ AS ] <replaceable class="parameter">output_name</replaceable> ] [, ...]
+  [ WITH WHEN CLAUSE [ [ AS ] <replaceable class="parameter">col_alias</replaceable> ] ] ]
 
 <phrase>where <replaceable class="parameter">data_source</replaceable> is:</phrase>
 
diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
new file mode 100644
index e34f583..aa3cca0
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -274,12 +274,6 @@ DoCopy(ParseState *pstate, const CopyStm
 	{
 		Assert(stmt->query);
 
-		/* MERGE is allowed by parser, but unimplemented. Reject for now */
-		if (IsA(stmt->query, MergeStmt))
-			ereport(ERROR,
-					errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("MERGE not supported in COPY"));
-
 		query = makeNode(RawStmt);
 		query->stmt = stmt->query;
 		query->stmt_location = stmt_location;
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
new file mode 100644
index 8043b4e..e02d7d0
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -510,7 +510,8 @@ BeginCopyTo(ParseState *pstate,
 		{
 			Assert(query->commandType == CMD_INSERT ||
 				   query->commandType == CMD_UPDATE ||
-				   query->commandType == CMD_DELETE);
+				   query->commandType == CMD_DELETE ||
+				   query->commandType == CMD_MERGE);
 
 			ereport(ERROR,
 					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
new file mode 100644
index 651ad24..3391269
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -612,6 +612,9 @@ ExecInitPartitionInfo(ModifyTableState *
 	 * case or in the case of UPDATE tuple routing where we didn't find a
 	 * result rel to reuse.
 	 */
+
+	/* XXX: What about the MERGE case ??? */
+
 	if (node && node->returningLists != NIL)
 	{
 		TupleTableSlot *slot;
@@ -877,6 +880,7 @@ ExecInitPartitionInfo(ModifyTableState *
 		List	   *firstMergeActionList = linitial(node->mergeActionLists);
 		ListCell   *lc;
 		ExprContext *econtext = mtstate->ps.ps_ExprContext;
+		int			action_idx = 1;
 
 		if (part_attmap == NULL)
 			part_attmap =
@@ -897,6 +901,7 @@ ExecInitPartitionInfo(ModifyTableState *
 			/* Generate the action's state for this relation */
 			action_state = makeNode(MergeActionState);
 			action_state->mas_action = action;
+			action_state->mas_action_idx = action_idx++;
 
 			/* And put the action in the appropriate list */
 			if (action->matched)
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
new file mode 100644
index 50e06ec..7e4717a
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -1665,8 +1665,8 @@ check_sql_fn_retval(List *queryTreeLists
 
 	/*
 	 * If it's a plain SELECT, it returns whatever the targetlist says.
-	 * Otherwise, if it's INSERT/UPDATE/DELETE with RETURNING, it returns
-	 * that. Otherwise, the function return type must be VOID.
+	 * Otherwise, if it's INSERT/UPDATE/DELETE/MERGE with RETURNING, it
+	 * returns that. Otherwise, the function return type must be VOID.
 	 *
 	 * Note: eventually replace this test with QueryReturnsTuples?	We'd need
 	 * a more general method of determining the output type, though.  Also, it
@@ -1684,7 +1684,8 @@ check_sql_fn_retval(List *queryTreeLists
 	else if (parse &&
 			 (parse->commandType == CMD_INSERT ||
 			  parse->commandType == CMD_UPDATE ||
-			  parse->commandType == CMD_DELETE) &&
+			  parse->commandType == CMD_DELETE ||
+			  parse->commandType == CMD_MERGE) &&
 			 parse->returningList)
 	{
 		tlist = parse->returningList;
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
new file mode 100644
index 687a542..64f6531
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -36,8 +36,7 @@
  *		RETURNING tuple after completing each row insert, update, or delete.
  *		It must be called again to continue the operation.  Without RETURNING,
  *		we just loop within the node until all the work is done, then
- *		return NULL.  This avoids useless call/return overhead.  (MERGE does
- *		not support RETURNING.)
+ *		return NULL.  This avoids useless call/return overhead.
  */
 
 #include "postgres.h"
@@ -172,13 +171,14 @@ static TupleTableSlot *ExecMerge(ModifyT
 								 ItemPointer tupleid,
 								 bool canSetTag);
 static void ExecInitMerge(ModifyTableState *mtstate, EState *estate);
-static bool ExecMergeMatched(ModifyTableContext *context,
-							 ResultRelInfo *resultRelInfo,
-							 ItemPointer tupleid,
-							 bool canSetTag);
-static void ExecMergeNotMatched(ModifyTableContext *context,
-								ResultRelInfo *resultRelInfo,
-								bool canSetTag);
+static TupleTableSlot *ExecMergeMatched(ModifyTableContext *context,
+										ResultRelInfo *resultRelInfo,
+										ItemPointer tupleid,
+										bool canSetTag,
+										bool *matched);
+static TupleTableSlot *ExecMergeNotMatched(ModifyTableContext *context,
+										   ResultRelInfo *resultRelInfo,
+										   bool canSetTag);
 static TupleTableSlot *mergeGetUpdateNewTuple(ResultRelInfo *relinfo,
 											  TupleTableSlot *planSlot,
 											  TupleTableSlot *oldSlot,
@@ -2692,6 +2692,7 @@ static TupleTableSlot *
 ExecMerge(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
 		  ItemPointer tupleid, bool canSetTag)
 {
+	TupleTableSlot *rslot = NULL;
 	bool		matched;
 
 	/*-----
@@ -2739,7 +2740,8 @@ ExecMerge(ModifyTableContext *context, R
 	 */
 	matched = tupleid != NULL;
 	if (matched)
-		matched = ExecMergeMatched(context, resultRelInfo, tupleid, canSetTag);
+		rslot = ExecMergeMatched(context, resultRelInfo, tupleid, canSetTag,
+								 &matched);
 
 	/*
 	 * Either we were dealing with a NOT MATCHED tuple or ExecMergeMatched()
@@ -2747,10 +2749,9 @@ ExecMerge(ModifyTableContext *context, R
 	 * matches.
 	 */
 	if (!matched)
-		ExecMergeNotMatched(context, resultRelInfo, canSetTag);
+		rslot = ExecMergeNotMatched(context, resultRelInfo, canSetTag);
 
-	/* No RETURNING support yet */
-	return NULL;
+	return rslot;
 }
 
 /*
@@ -2760,8 +2761,8 @@ ExecMerge(ModifyTableContext *context, R
  * We start from the first WHEN MATCHED action and check if the WHEN quals
  * pass, if any. If the WHEN quals for the first action do not pass, we
  * check the second, then the third and so on. If we reach to the end, no
- * action is taken and we return true, indicating that no further action is
- * required for this tuple.
+ * action is taken and "matched" is set to true, indicating that no further
+ * action is required for this tuple.
  *
  * If we do find a qualifying action, then we attempt to execute the action.
  *
@@ -2770,15 +2771,17 @@ ExecMerge(ModifyTableContext *context, R
  * with individual actions are evaluated by this routine via ExecQual, while
  * EvalPlanQual checks for the join quals. If EvalPlanQual tells us that the
  * updated tuple still passes the join quals, then we restart from the first
- * action to look for a qualifying action. Otherwise, we return false --
- * meaning that a NOT MATCHED action must now be executed for the current
- * source tuple.
+ * action to look for a qualifying action. Otherwise, "matched" is set to
+ * false -- meaning that a NOT MATCHED action must now be executed for the
+ * current source tuple.
  */
-static bool
+static TupleTableSlot *
 ExecMergeMatched(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
-				 ItemPointer tupleid, bool canSetTag)
+				 ItemPointer tupleid, bool canSetTag, bool *matched)
 {
 	ModifyTableState *mtstate = context->mtstate;
+	ModifyTable *node = (ModifyTable *) mtstate->ps.plan;
+	TupleTableSlot *rslot = NULL;
 	TupleTableSlot *newslot;
 	EState	   *estate = context->estate;
 	ExprContext *econtext = mtstate->ps.ps_ExprContext;
@@ -2790,7 +2793,10 @@ ExecMergeMatched(ModifyTableContext *con
 	 * If there are no WHEN MATCHED actions, we are done.
 	 */
 	if (resultRelInfo->ri_matchedMergeAction == NIL)
-		return true;
+	{
+		*matched = true;
+		return NULL;
+	}
 
 	/*
 	 * Make tuple and any needed join variables available to ExecQual and
@@ -2952,7 +2958,8 @@ lmerge_matched:
 				 * If the tuple was already deleted, return to let caller
 				 * handle it under NOT MATCHED clauses.
 				 */
-				return false;
+				*matched = false;
+				return NULL;
 
 			case TM_Updated:
 				{
@@ -3020,13 +3027,19 @@ lmerge_matched:
 							 * NOT MATCHED actions.
 							 */
 							if (TupIsNull(epqslot))
-								return false;
+							{
+								*matched = false;
+								return NULL;
+							}
 
 							(void) ExecGetJunkAttribute(epqslot,
 														resultRelInfo->ri_RowIdAttNo,
 														&isNull);
 							if (isNull)
-								return false;
+							{
+								*matched = false;
+								return NULL;
+							}
 
 							/*
 							 * When a tuple was updated and migrated to
@@ -3061,7 +3074,8 @@ lmerge_matched:
 							 * tuple already deleted; tell caller to run NOT
 							 * MATCHED actions
 							 */
-							return false;
+							*matched = false;
+							return NULL;
 
 						case TM_SelfModified:
 
@@ -3081,13 +3095,14 @@ lmerge_matched:
 										(errcode(ERRCODE_TRIGGERED_DATA_CHANGE_VIOLATION),
 										 errmsg("tuple to be updated or deleted was already modified by an operation triggered by the current command"),
 										 errhint("Consider using an AFTER trigger instead of a BEFORE trigger to propagate changes to other rows.")));
-							return false;
+							*matched = false;
+							return NULL;
 
 						default:
 							/* see table_tuple_lock call in ExecDelete() */
 							elog(ERROR, "unexpected table_tuple_lock status: %u",
 								 result);
-							return false;
+							return NULL;
 					}
 				}
 
@@ -3099,6 +3114,39 @@ lmerge_matched:
 				break;
 		}
 
+		/* Process RETURNING if present */
+		if (resultRelInfo->ri_projectReturning)
+		{
+			switch (commandType)
+			{
+				case CMD_UPDATE:
+					rslot = ExecProcessReturning(resultRelInfo, newslot,
+												 context->planSlot);
+					break;
+
+				case CMD_DELETE:
+					rslot = ExecProcessReturning(resultRelInfo,
+												 resultRelInfo->ri_oldTupleSlot,
+												 context->planSlot);
+					break;
+
+				case CMD_NOTHING:
+					break;
+
+				default:
+					elog(ERROR, "unknown action in MERGE WHEN MATCHED clause");
+			}
+		}
+
+		/* Set RETURNING WITH WHEN CLAUSE result, if required */
+		if (rslot && node->hasReturningWhenClause)
+		{
+			int			attnum = rslot->tts_tupleDescriptor->natts - 1;
+
+			rslot->tts_values[attnum] = Int32GetDatum(relaction->mas_action_idx);
+			rslot->tts_isnull[attnum] = false;
+		}
+
 		/*
 		 * We've activated one of the WHEN clauses, so we don't search
 		 * further. This is required behaviour, not an optimization.
@@ -3109,19 +3157,23 @@ lmerge_matched:
 	/*
 	 * Successfully executed an action or no qualifying action was found.
 	 */
-	return true;
+	*matched = true;
+
+	return rslot;
 }
 
 /*
  * Execute the first qualifying NOT MATCHED action.
  */
-static void
+static TupleTableSlot *
 ExecMergeNotMatched(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
 					bool canSetTag)
 {
 	ModifyTableState *mtstate = context->mtstate;
+	ModifyTable *node = (ModifyTable *) mtstate->ps.plan;
 	ExprContext *econtext = mtstate->ps.ps_ExprContext;
 	List	   *actionStates = NIL;
+	TupleTableSlot *rslot = NULL;
 	ListCell   *l;
 
 	/*
@@ -3173,9 +3225,18 @@ ExecMergeNotMatched(ModifyTableContext *
 				newslot = ExecProject(action->mas_proj);
 				context->relaction = action;
 
-				(void) ExecInsert(context, mtstate->rootResultRelInfo, newslot,
-								  canSetTag, NULL, NULL);
+				rslot = ExecInsert(context, mtstate->rootResultRelInfo,
+								   newslot, canSetTag, NULL, NULL);
 				mtstate->mt_merge_inserted += 1;
+
+				/* Set RETURNING WITH WHEN CLAUSE result, if requested */
+				if (rslot && node->hasReturningWhenClause)
+				{
+					int			attnum = rslot->tts_tupleDescriptor->natts - 1;
+
+					rslot->tts_values[attnum] = Int32GetDatum(action->mas_action_idx);
+					rslot->tts_isnull[attnum] = false;
+				}
 				break;
 			case CMD_NOTHING:
 				/* Do nothing */
@@ -3190,6 +3251,8 @@ ExecMergeNotMatched(ModifyTableContext *
 		 */
 		break;
 	}
+
+	return rslot;
 }
 
 /*
@@ -3227,6 +3290,7 @@ ExecInitMerge(ModifyTableState *mtstate,
 		List	   *mergeActionList = lfirst(lc);
 		TupleDesc	relationDesc;
 		ListCell   *l;
+		int			action_idx;
 
 		resultRelInfo = mtstate->resultRelInfo + i;
 		i++;
@@ -3236,6 +3300,7 @@ ExecInitMerge(ModifyTableState *mtstate,
 		if (unlikely(!resultRelInfo->ri_projectNewInfoValid))
 			ExecInitMergeTupleSlots(mtstate, resultRelInfo);
 
+		action_idx = 1;
 		foreach(l, mergeActionList)
 		{
 			MergeAction *action = (MergeAction *) lfirst(l);
@@ -3250,6 +3315,7 @@ ExecInitMerge(ModifyTableState *mtstate,
 			 */
 			action_state = makeNode(MergeActionState);
 			action_state->mas_action = action;
+			action_state->mas_action_idx = action_idx++;
 			action_state->mas_whenqual = ExecInitQual((List *) action->qual,
 													  &mtstate->ps);
 
@@ -3671,8 +3737,17 @@ ExecModifyTable(PlanState *pstate)
 				{
 					EvalPlanQualSetSlot(&node->mt_epqstate, context.planSlot);
 
-					ExecMerge(&context, node->resultRelInfo, NULL, node->canSetTag);
-					continue;	/* no RETURNING support yet */
+					slot = ExecMerge(&context, node->resultRelInfo, NULL,
+									 node->canSetTag);
+
+					/*
+					 * If we got a RETURNING result, return it to the caller.
+					 * We'll continue the work on next call.
+					 */
+					if (slot)
+						return slot;
+
+					continue;	/* continue with the next tuple */
 				}
 
 				elog(ERROR, "tableoid is NULL");
@@ -3749,8 +3824,17 @@ ExecModifyTable(PlanState *pstate)
 					{
 						EvalPlanQualSetSlot(&node->mt_epqstate, context.planSlot);
 
-						ExecMerge(&context, node->resultRelInfo, NULL, node->canSetTag);
-						continue;	/* no RETURNING support yet */
+						slot = ExecMerge(&context, node->resultRelInfo, NULL,
+										 node->canSetTag);
+
+						/*
+						 * If we got a RETURNING result, return it to the
+						 * caller.  We'll continue the work on next call.
+						 */
+						if (slot)
+							return slot;
+
+						continue;	/* continue with the next tuple */
 					}
 
 					elog(ERROR, "ctid is NULL");
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
new file mode 100644
index aa24958..90b7fdc
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -7056,6 +7056,7 @@ make_modifytable(PlannerInfo *root, Plan
 	node->rowMarks = rowMarks;
 	node->mergeActionLists = mergeActionLists;
 	node->epqParam = epqParam;
+	node->hasReturningWhenClause = root->parse->hasReturningWhenClause;
 
 	/*
 	 * For each result relation that is a foreign table, allow the FDW to
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
new file mode 100644
index 5b90974..50285a2
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -73,7 +73,6 @@ static void determineRecursiveColTypes(P
 									   Node *larg, List *nrtargetlist);
 static Query *transformReturnStmt(ParseState *pstate, ReturnStmt *stmt);
 static Query *transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt);
-static List *transformReturningList(ParseState *pstate, List *returningList);
 static Query *transformPLAssignStmt(ParseState *pstate,
 									PLAssignStmt *stmt);
 static Query *transformDeclareCursorStmt(ParseState *pstate,
@@ -2501,7 +2500,7 @@ transformUpdateTargetList(ParseState *ps
  * transformReturningList -
  *	handle a RETURNING clause in INSERT/UPDATE/DELETE
  */
-static List *
+List *
 transformReturningList(ParseState *pstate, List *returningList)
 {
 	List	   *rlist;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
new file mode 100644
index a013838..cf8599a
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -151,6 +151,13 @@ typedef struct KeyActions
 	KeyAction *deleteAction;
 } KeyActions;
 
+/* Private struct for the result of merge_returning_clause */
+typedef struct MergeReturningClause
+{
+	List	   *returningList;
+	char	   *returningWhenClause;
+} MergeReturningClause;
+
 /* ConstraintAttributeSpec yields an integer bitmask of these flags: */
 #define CAS_NOT_DEFERRABLE			0x01
 #define CAS_DEFERRABLE				0x02
@@ -274,6 +281,7 @@ static Node *makeRecursiveViewSelect(cha
 	SetQuantifier setquantifier;
 	struct GroupClause *groupclause;
 	MergeWhenClause *mergewhen;
+	struct MergeReturningClause *mergereturning;
 	struct KeyActions *keyactions;
 	struct KeyAction *keyaction;
 }
@@ -511,6 +519,7 @@ static Node *makeRecursiveViewSelect(cha
 %type <infer>	opt_conf_expr
 %type <onconflict> opt_on_conflict
 %type <mergewhen>	merge_insert merge_update merge_delete
+%type <mergereturning> merge_returning_clause
 
 %type <node>	merge_when_clause opt_merge_when_condition
 %type <list>	merge_when_list
@@ -677,7 +686,7 @@ static Node *makeRecursiveViewSelect(cha
 	BOOLEAN_P BOTH BREADTH BY
 
 	CACHE CALL CALLED CASCADE CASCADED CASE CAST CATALOG_P CHAIN CHAR_P
-	CHARACTER CHARACTERISTICS CHECK CHECKPOINT CLASS CLOSE
+	CHARACTER CHARACTERISTICS CHECK CHECKPOINT CLASS CLAUSE CLOSE
 	CLUSTER COALESCE COLLATE COLLATION COLUMN COLUMNS COMMENT COMMENTS COMMIT
 	COMMITTED COMPRESSION CONCURRENTLY CONFIGURATION CONFLICT
 	CONNECTION CONSTRAINT CONSTRAINTS CONTENT_P CONTINUE_P CONVERSION_P COPY
@@ -12241,6 +12250,7 @@ MergeStmt:
 			USING table_ref
 			ON a_expr
 			merge_when_list
+			merge_returning_clause
 				{
 					MergeStmt  *m = makeNode(MergeStmt);
 
@@ -12249,6 +12259,8 @@ MergeStmt:
 					m->sourceRelation = $6;
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
+					m->returningList = $10 ? $10->returningList : NIL;
+					m->returningWhenClause = $10 ? $10->returningWhenClause : NULL;
 
 					$$ = (Node *) m;
 				}
@@ -12389,6 +12401,45 @@ merge_values_clause:
 				}
 		;
 
+merge_returning_clause:
+			RETURNING target_list WITH WHEN CLAUSE AS ColLabel
+				{
+					MergeReturningClause *n = palloc(sizeof(MergeReturningClause));
+
+					n->returningList = $2;
+					n->returningWhenClause = $7;
+					$$ = n;
+				}
+			| RETURNING target_list WITH WHEN CLAUSE ColLabel
+				{
+					MergeReturningClause *n = palloc(sizeof(MergeReturningClause));
+
+					n->returningList = $2;
+					n->returningWhenClause = $6;
+					$$ = n;
+				}
+			| RETURNING target_list WITH WHEN CLAUSE
+				{
+					MergeReturningClause *n = palloc(sizeof(MergeReturningClause));
+
+					n->returningList = $2;
+					n->returningWhenClause = pstrdup("when_clause");
+					$$ = n;
+				}
+			| RETURNING target_list
+				{
+					MergeReturningClause *n = palloc(sizeof(MergeReturningClause));
+
+					n->returningList = $2;
+					n->returningWhenClause = NULL;
+					$$ = n;
+				}
+			| /* EMPTY */
+				{
+					$$ = NULL;
+				}
+		;
+
 /*****************************************************************************
  *
  *		QUERY:
@@ -16744,6 +16795,7 @@ unreserved_keyword:
 			| CHARACTERISTICS
 			| CHECKPOINT
 			| CLASS
+			| CLAUSE
 			| CLOSE
 			| CLUSTER
 			| COLUMNS
@@ -17269,6 +17321,7 @@ bare_label_keyword:
 			| CHECK
 			| CHECKPOINT
 			| CLASS
+			| CLAUSE
 			| CLOSE
 			| CLUSTER
 			| COALESCE
diff --git a/src/backend/parser/parse_cte.c b/src/backend/parser/parse_cte.c
new file mode 100644
index c5b1a49..ad3c525
--- a/src/backend/parser/parse_cte.c
+++ b/src/backend/parser/parse_cte.c
@@ -126,13 +126,6 @@ transformWithClause(ParseState *pstate,
 		CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc);
 		ListCell   *rest;
 
-		/* MERGE is allowed by parser, but unimplemented. Reject for now */
-		if (IsA(cte->ctequery, MergeStmt))
-			ereport(ERROR,
-					errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("MERGE not supported in WITH query"),
-					parser_errposition(pstate, cte->location));
-
 		for_each_cell(rest, withClause->ctes, lnext(withClause->ctes, lc))
 		{
 			CommonTableExpr *cte2 = (CommonTableExpr *) lfirst(rest);
@@ -153,7 +146,8 @@ transformWithClause(ParseState *pstate,
 			/* must be a data-modifying statement */
 			Assert(IsA(cte->ctequery, InsertStmt) ||
 				   IsA(cte->ctequery, UpdateStmt) ||
-				   IsA(cte->ctequery, DeleteStmt));
+				   IsA(cte->ctequery, DeleteStmt) ||
+				   IsA(cte->ctequery, MergeStmt));
 
 			pstate->p_hasModifyingCTE = true;
 		}
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
new file mode 100644
index 611dfce..0ba114c
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -233,6 +233,34 @@ transformMergeStmt(ParseState *pstate, M
 	 */
 	qry->jointree = makeFromExpr(pstate->p_joinlist, joinExpr);
 
+	/* Transform the RETURNING list, if any */
+	qry->returningList = transformReturningList(pstate, stmt->returningList);
+
+	/* Add integer column for RETURNING ... WITH WHEN CLAUSE, if specified */
+	if (stmt->returningWhenClause)
+	{
+		Expr	   *null_expr;
+		TargetEntry *tle;
+
+		null_expr = (Expr *) makeConst(INT4OID,
+									   -1,
+									   InvalidOid,
+									   sizeof(int32),
+									   (Datum) 0,
+									   true,
+									   true);
+
+		tle = makeTargetEntry(null_expr,
+							  list_length(qry->returningList) + 1,
+							  stmt->returningWhenClause,
+							  false);
+
+		qry->returningList = lappend(qry->returningList, tle);
+		qry->hasReturningWhenClause = true;
+	}
+	else
+		qry->hasReturningWhenClause = false;
+
 	/*
 	 * We now have a good query shape, so now look at the WHEN conditions and
 	 * action targetlists.
@@ -390,9 +418,6 @@ transformMergeStmt(ParseState *pstate, M
 
 	qry->mergeActionList = mergeActionList;
 
-	/* RETURNING could potentially be added in the future, but not in SQL std */
-	qry->returningList = NULL;
-
 	qry->hasTargetSRFs = false;
 	qry->hasSubLinks = pstate->p_hasSubLinks;
 
diff --git a/src/backend/parser/parse_relation.c b/src/backend/parser/parse_relation.c
new file mode 100644
index 5389a0e..162a305
--- a/src/backend/parser/parse_relation.c
+++ b/src/backend/parser/parse_relation.c
@@ -2313,9 +2313,10 @@ addRangeTableEntryForCTE(ParseState *pst
 		cte->cterefcount++;
 
 	/*
-	 * We throw error if the CTE is INSERT/UPDATE/DELETE without RETURNING.
-	 * This won't get checked in case of a self-reference, but that's OK
-	 * because data-modifying CTEs aren't allowed to be recursive anyhow.
+	 * We throw error if the CTE is INSERT/UPDATE/DELETE/MERGE without
+	 * RETURNING.  This won't get checked in case of a self-reference, but
+	 * that's OK because data-modifying CTEs aren't allowed to be recursive
+	 * anyhow.
 	 */
 	if (IsA(cte->ctequery, Query))
 	{
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
new file mode 100644
index 1960dad..ac5a0c7
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -3640,7 +3640,8 @@ RewriteQuery(Query *parsetree, List *rew
 			if (!(ctequery->commandType == CMD_SELECT ||
 				  ctequery->commandType == CMD_UPDATE ||
 				  ctequery->commandType == CMD_INSERT ||
-				  ctequery->commandType == CMD_DELETE))
+				  ctequery->commandType == CMD_DELETE ||
+				  ctequery->commandType == CMD_MERGE))
 			{
 				/*
 				 * Currently it could only be NOTIFY; this error message will
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
new file mode 100644
index c7d9d96..d393f51
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -2131,11 +2131,10 @@ QueryReturnsTuples(Query *parsetree)
 		case CMD_SELECT:
 			/* returns tuples */
 			return true;
-		case CMD_MERGE:
-			return false;
 		case CMD_INSERT:
 		case CMD_UPDATE:
 		case CMD_DELETE:
+		case CMD_MERGE:
 			/* the forms with RETURNING return tuples */
 			if (parsetree->returningList)
 				return true;
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
new file mode 100644
index 38f9b10..024e75a
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -962,13 +962,17 @@ PrintQueryResult(PGresult *result, bool
 			else
 				success = true;
 
-			/* if it's INSERT/UPDATE/DELETE RETURNING, also print status */
+			/*
+			 * If it's INSERT/UPDATE/DELETE/MERGE RETURNING, also print
+			 * status.
+			 */
 			if (last || pset.show_all_results)
 			{
 				cmdstatus = PQcmdStatus(result);
 				if (strncmp(cmdstatus, "INSERT", 6) == 0 ||
 					strncmp(cmdstatus, "UPDATE", 6) == 0 ||
-					strncmp(cmdstatus, "DELETE", 6) == 0)
+					strncmp(cmdstatus, "DELETE", 6) == 0 ||
+					strncmp(cmdstatus, "MERGE", 5) == 0)
 					PrintQueryStatus(result, printStatusFout);
 			}
 
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
new file mode 100644
index 20f4c8b..3a25c93
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -415,6 +415,7 @@ typedef struct MergeActionState
 	NodeTag		type;
 
 	MergeAction *mas_action;	/* associated MergeAction node */
+	int			mas_action_idx;	/* 1-based index of MergeAction node */
 	ProjectionInfo *mas_proj;	/* projection of the action's targetlist for
 								 * this rel */
 	ExprState  *mas_whenqual;	/* WHEN [NOT] MATCHED AND conditions */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
new file mode 100644
index cfeca96..c8fc76c
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -160,6 +160,7 @@ typedef struct Query
 
 	List	   *mergeActionList;	/* list of actions for MERGE (only) */
 	bool		mergeUseOuterJoin;	/* whether to use outer join */
+	bool		hasReturningWhenClause;	/* has RETURNING WITH WHEN CLAUSE */
 
 	List	   *targetList;		/* target list (of TargetEntry) */
 
@@ -1713,6 +1714,9 @@ typedef struct MergeStmt
 	Node	   *sourceRelation; /* source relation */
 	Node	   *joinCondition;	/* join condition between source and target */
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
+	List	   *returningList;	/* list of expressions to return */
+	char	   *returningWhenClause;	/* RETURNING ... WITH WHEN CLAUSE
+										 * column alias */
 	WithClause *withClause;		/* WITH clause */
 } MergeStmt;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
new file mode 100644
index c1234fc..c8c8b97
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -254,6 +254,7 @@ typedef struct ModifyTable
 	List	   *exclRelTlist;	/* tlist of the EXCLUDED pseudo relation */
 	List	   *mergeActionLists;	/* per-target-table lists of actions for
 									 * MERGE */
+	bool		hasReturningWhenClause;	/* has RETURNING WITH WHEN CLAUSE */
 } ModifyTable;
 
 struct PartitionPruneInfo;		/* forward reference to struct below */
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
new file mode 100644
index c97be6e..34cc1ec
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -44,6 +44,7 @@ extern List *transformInsertRow(ParseSta
 								bool strip_indirection);
 extern List *transformUpdateTargetList(ParseState *pstate,
 									   List *origTlist);
+extern List *transformReturningList(ParseState *pstate, List *returningList);
 extern Query *transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree);
 extern Query *transformStmt(ParseState *pstate, Node *parseTree);
 
diff --git a/src/include/parser/kwlist.h b/src/include/parser/kwlist.h
new file mode 100644
index bb36213..e3860fa
--- a/src/include/parser/kwlist.h
+++ b/src/include/parser/kwlist.h
@@ -79,6 +79,7 @@ PG_KEYWORD("characteristics", CHARACTERI
 PG_KEYWORD("check", CHECK, RESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("checkpoint", CHECKPOINT, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("class", CLASS, UNRESERVED_KEYWORD, BARE_LABEL)
+PG_KEYWORD("clause", CLAUSE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("close", CLOSE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("cluster", CLUSTER, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("coalesce", COALESCE, COL_NAME_KEYWORD, BARE_LABEL)
diff --git a/src/test/regress/expected/merge.out b/src/test/regress/expected/merge.out
new file mode 100644
index 6c8a18f..0adf7c1
--- a/src/test/regress/expected/merge.out
+++ b/src/test/regress/expected/merge.out
@@ -123,20 +123,20 @@ ON tid = tid
 WHEN MATCHED THEN DO NOTHING;
 ERROR:  name "target" specified more than once
 DETAIL:  The name is used both as MERGE target table and data source.
--- used in a CTE
+-- used in a CTE without RETURNING
 WITH foo AS (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE
 ) SELECT * FROM foo;
-ERROR:  MERGE not supported in WITH query
-LINE 1: WITH foo AS (
-             ^
--- used in COPY
+ERROR:  WITH query "foo" does not have a RETURNING clause
+LINE 4: ) SELECT * FROM foo;
+                        ^
+-- used in COPY without RETURNING
 COPY (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE
 ) TO stdout;
-ERROR:  MERGE not supported in COPY
+ERROR:  COPY query must have a RETURNING clause
 -- unsupported relation types
 -- view
 CREATE VIEW tv AS SELECT * FROM target;
@@ -1299,11 +1299,69 @@ WHEN NOT MATCHED THEN
 	INSERT (balance, tid) VALUES (balance + delta, sid)
 WHEN MATCHED AND tid < 2 THEN
 	DELETE
-RETURNING *;
-ERROR:  syntax error at or near "RETURNING"
-LINE 10: RETURNING *;
-         ^
+RETURNING t.*;
+ tid | balance 
+-----+---------
+  -1 |     -11
+   1 |     100
+(2 rows)
+
+ROLLBACK;
+-- RETURNING ... WITH WHEN CLAUSE
+BEGIN;
+INSERT INTO sq_source (sid, balance, delta) VALUES (-1, -1, -10);
+WITH t AS (
+    MERGE INTO sq_target t
+    USING v
+    ON tid = sid
+    WHEN MATCHED AND tid > 2 THEN
+        UPDATE SET balance = t.balance + delta
+    WHEN NOT MATCHED THEN
+        INSERT (balance, tid) VALUES (balance + delta, sid)
+    WHEN MATCHED AND tid < 2 THEN
+        DELETE
+    RETURNING t.* WITH WHEN CLAUSE
+)
+SELECT CASE when_clause
+         WHEN 1 THEN 'UPDATE'
+         WHEN 2 THEN 'INSERT'
+         WHEN 3 THEN 'DELETE'
+       END, *
+FROM t;
+  case  | tid | balance | when_clause 
+--------+-----+---------+-------------
+ INSERT |  -1 |     -11 |           2
+ DELETE |   1 |     100 |           3
+(2 rows)
+
 ROLLBACK;
+-- RETURNING ... WITH WHEN CLAUSE [AS] alias
+MERGE INTO sq_target t
+USING sq_source s
+ON tid = sid
+WHEN MATCHED AND tid = 1 THEN
+    DO NOTHING
+WHEN MATCHED THEN
+    UPDATE SET balance = t.balance + delta
+RETURNING t.* WITH WHEN CLAUSE AS wc;
+ tid | balance | wc 
+-----+---------+----
+   2 |     220 |  2
+(1 row)
+
+MERGE INTO sq_target t
+USING sq_source s
+ON tid = sid
+WHEN MATCHED AND tid = 2 THEN
+    UPDATE SET balance = t.balance - delta
+WHEN MATCHED THEN
+    DO NOTHING
+RETURNING t.* WITH WHEN CLAUSE wc;
+ tid | balance | wc 
+-----+---------+----
+   2 |     200 |  1
+(1 row)
+
 -- EXPLAIN
 CREATE TABLE ex_mtarget (a int, b int)
   WITH (autovacuum_enabled=off);
diff --git a/src/test/regress/sql/merge.sql b/src/test/regress/sql/merge.sql
new file mode 100644
index 98fe104..42f2483
--- a/src/test/regress/sql/merge.sql
+++ b/src/test/regress/sql/merge.sql
@@ -86,12 +86,12 @@ MERGE INTO target
 USING target
 ON tid = tid
 WHEN MATCHED THEN DO NOTHING;
--- used in a CTE
+-- used in a CTE without RETURNING
 WITH foo AS (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE
 ) SELECT * FROM foo;
--- used in COPY
+-- used in COPY without RETURNING
 COPY (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE
@@ -854,9 +854,51 @@ WHEN NOT MATCHED THEN
 	INSERT (balance, tid) VALUES (balance + delta, sid)
 WHEN MATCHED AND tid < 2 THEN
 	DELETE
-RETURNING *;
+RETURNING t.*;
+ROLLBACK;
+
+-- RETURNING ... WITH WHEN CLAUSE
+BEGIN;
+INSERT INTO sq_source (sid, balance, delta) VALUES (-1, -1, -10);
+WITH t AS (
+    MERGE INTO sq_target t
+    USING v
+    ON tid = sid
+    WHEN MATCHED AND tid > 2 THEN
+        UPDATE SET balance = t.balance + delta
+    WHEN NOT MATCHED THEN
+        INSERT (balance, tid) VALUES (balance + delta, sid)
+    WHEN MATCHED AND tid < 2 THEN
+        DELETE
+    RETURNING t.* WITH WHEN CLAUSE
+)
+SELECT CASE when_clause
+         WHEN 1 THEN 'UPDATE'
+         WHEN 2 THEN 'INSERT'
+         WHEN 3 THEN 'DELETE'
+       END, *
+FROM t;
 ROLLBACK;
 
+-- RETURNING ... WITH WHEN CLAUSE [AS] alias
+MERGE INTO sq_target t
+USING sq_source s
+ON tid = sid
+WHEN MATCHED AND tid = 1 THEN
+    DO NOTHING
+WHEN MATCHED THEN
+    UPDATE SET balance = t.balance + delta
+RETURNING t.* WITH WHEN CLAUSE AS wc;
+
+MERGE INTO sq_target t
+USING sq_source s
+ON tid = sid
+WHEN MATCHED AND tid = 2 THEN
+    UPDATE SET balance = t.balance - delta
+WHEN MATCHED THEN
+    DO NOTHING
+RETURNING t.* WITH WHEN CLAUSE wc;
+
 -- EXPLAIN
 CREATE TABLE ex_mtarget (a int, b int)
   WITH (autovacuum_enabled=off);