2 files changed, 140 insertions, 7 deletions

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 4eeed580b16..29ae27e5e32 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -1643,6 +1643,10 @@ matchLocks(CmdType event,
 	if (rulelocks == NULL)
 		return NIL;
 
+	/* No rule support for MERGE */
+	if (parsetree->commandType == CMD_MERGE)
+		return NIL;
+
 	if (parsetree->commandType != CMD_SELECT)
 	{
 		if (parsetree->resultRelation != varno)
@@ -3671,8 +3675,8 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 	}
 
 	/*
-	 * If the statement is an insert, update, or delete, adjust its targetlist
-	 * as needed, and then fire INSERT/UPDATE/DELETE rules on it.
+	 * If the statement is an insert, update, delete, or merge, adjust its
+	 * targetlist as needed, and then fire INSERT/UPDATE/DELETE rules on it.
 	 *
 	 * SELECT rules are handled later when we have all the queries that should
 	 * get executed.  Also, utilities aren't rewritten at all (do we still
@@ -3770,6 +3774,7 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 		}
 		else if (event == CMD_UPDATE)
 		{
+			Assert(parsetree->override == OVERRIDING_NOT_SET);
 			parsetree->targetList =
 				rewriteTargetListIU(parsetree->targetList,
 									parsetree->commandType,
@@ -3780,6 +3785,38 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 			/* Also populate extraUpdatedCols (for generated columns) */
 			fill_extraUpdatedCols(rt_entry, rt_entry_relation);
 		}
+		else if (event == CMD_MERGE)
+		{
+			Assert(parsetree->override == OVERRIDING_NOT_SET);
+
+			/*
+			 * Rewrite each action targetlist separately
+			 */
+			foreach(lc1, parsetree->mergeActionList)
+			{
+				MergeAction *action = (MergeAction *) lfirst(lc1);
+
+				switch (action->commandType)
+				{
+					case CMD_NOTHING:
+					case CMD_DELETE:	/* Nothing to do here */
+						break;
+					case CMD_UPDATE:
+					case CMD_INSERT:
+						/* XXX is it possible to have a VALUES clause? */
+						action->targetList =
+							rewriteTargetListIU(action->targetList,
+												action->commandType,
+												action->override,
+												rt_entry_relation,
+												NULL, 0, NULL);
+						break;
+					default:
+						elog(ERROR, "unrecognized commandType: %d", action->commandType);
+						break;
+				}
+			}
+		}
 		else if (event == CMD_DELETE)
 		{
 			/* Nothing to do here */
diff --git a/src/backend/rewrite/rowsecurity.c b/src/backend/rewrite/rowsecurity.c
index f0a046d65a6..a233dd47585 100644
--- a/src/backend/rewrite/rowsecurity.c
+++ b/src/backend/rewrite/rowsecurity.c
@@ -232,15 +232,17 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 						   hasSubLinks);
 
 	/*
-	 * Similar to above, during an UPDATE or DELETE, if SELECT rights are also
-	 * required (eg: when a RETURNING clause exists, or the user has provided
-	 * a WHERE clause which involves columns from the relation), we collect up
-	 * CMD_SELECT policies and add them via add_security_quals first.
+	 * Similar to above, during an UPDATE, DELETE, or MERGE, if SELECT rights
+	 * are also required (eg: when a RETURNING clause exists, or the user has
+	 * provided a WHERE clause which involves columns from the relation), we
+	 * collect up CMD_SELECT policies and add them via add_security_quals
+	 * first.
 	 *
 	 * This way, we filter out any records which are not visible through an
 	 * ALL or SELECT USING policy.
 	 */
-	if ((commandType == CMD_UPDATE || commandType == CMD_DELETE) &&
+	if ((commandType == CMD_UPDATE || commandType == CMD_DELETE ||
+		 commandType == CMD_MERGE) &&
 		rte->requiredPerms & ACL_SELECT)
 	{
 		List	   *select_permissive_policies;
@@ -380,6 +382,92 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 		}
 	}
 
+	/*
+	 * FOR MERGE, we fetch policies for UPDATE, DELETE and INSERT (and ALL)
+	 * and set them up so that we can enforce the appropriate policy depending
+	 * on the final action we take.
+	 *
+	 * We already fetched the SELECT policies above.
+	 *
+	 * We don't push the UPDATE/DELETE USING quals to the RTE because we don't
+	 * really want to apply them while scanning the relation since we don't
+	 * know whether we will be doing an UPDATE or a DELETE at the end. We
+	 * apply the respective policy once we decide the final action on the
+	 * target tuple.
+	 *
+	 * XXX We are setting up USING quals as WITH CHECK. If RLS prohibits
+	 * UPDATE/DELETE on the target row, we shall throw an error instead of
+	 * silently ignoring the row. This is different than how normal
+	 * UPDATE/DELETE works and more in line with INSERT ON CONFLICT DO UPDATE
+	 * handling.
+	 */
+	if (commandType == CMD_MERGE)
+	{
+		List	   *merge_permissive_policies;
+		List	   *merge_restrictive_policies;
+
+		/*
+		 * Fetch the UPDATE policies and set them up to execute on the
+		 * existing target row before doing UPDATE.
+		 */
+		get_policies_for_relation(rel, CMD_UPDATE, user_id,
+								  &merge_permissive_policies,
+								  &merge_restrictive_policies);
+
+		/*
+		 * WCO_RLS_MERGE_UPDATE_CHECK is used to check UPDATE USING quals on
+		 * the existing target row.
+		 */
+		add_with_check_options(rel, rt_index,
+							   WCO_RLS_MERGE_UPDATE_CHECK,
+							   merge_permissive_policies,
+							   merge_restrictive_policies,
+							   withCheckOptions,
+							   hasSubLinks,
+							   true);
+
+		/*
+		 * Same with DELETE policies.
+		 */
+		get_policies_for_relation(rel, CMD_DELETE, user_id,
+								  &merge_permissive_policies,
+								  &merge_restrictive_policies);
+
+		add_with_check_options(rel, rt_index,
+							   WCO_RLS_MERGE_DELETE_CHECK,
+							   merge_permissive_policies,
+							   merge_restrictive_policies,
+							   withCheckOptions,
+							   hasSubLinks,
+							   true);
+
+		/*
+		 * No special handling is required for INSERT policies. They will be
+		 * checked and enforced during ExecInsert(). But we must add them to
+		 * withCheckOptions.
+		 */
+		get_policies_for_relation(rel, CMD_INSERT, user_id,
+								  &merge_permissive_policies,
+								  &merge_restrictive_policies);
+
+		add_with_check_options(rel, rt_index,
+							   WCO_RLS_INSERT_CHECK,
+							   merge_permissive_policies,
+							   merge_restrictive_policies,
+							   withCheckOptions,
+							   hasSubLinks,
+							   false);
+
+		/* Enforce the WITH CHECK clauses of the UPDATE policies */
+		add_with_check_options(rel, rt_index,
+							   WCO_RLS_UPDATE_CHECK,
+							   merge_permissive_policies,
+							   merge_restrictive_policies,
+							   withCheckOptions,
+							   hasSubLinks,
+							   false);
+	}
+
 	table_close(rel, NoLock);
 
 	/*
@@ -444,6 +532,14 @@ get_policies_for_relation(Relation relation, CmdType cmd, Oid user_id,
 					if (policy->polcmd == ACL_DELETE_CHR)
 						cmd_matches = true;
 					break;
+				case CMD_MERGE:
+
+					/*
+					 * We do not support a separate policy for MERGE command.
+					 * Instead it derives from the policies defined for other
+					 * commands.
+					 */
+					break;
 				default:
 					elog(ERROR, "unrecognized policy command type %d",
 						 (int) cmd);