Add support for MERGE SQL command

MERGE performs actions that modify rows in the target table using a
source table or query. MERGE provides a single SQL statement that can
conditionally INSERT/UPDATE/DELETE rows -- a task that would otherwise
require multiple PL statements.  For example,

MERGE INTO target AS t
USING source AS s
ON t.tid = s.sid
WHEN MATCHED AND t.balance > s.delta THEN
  UPDATE SET balance = t.balance - s.delta
WHEN MATCHED THEN
  DELETE
WHEN NOT MATCHED AND s.delta > 0 THEN
  INSERT VALUES (s.sid, s.delta)
WHEN NOT MATCHED THEN
  DO NOTHING;

MERGE works with regular tables, partitioned tables and inheritance
hierarchies, including column and row security enforcement, as well as
support for row and statement triggers and transition tables therein.

MERGE is optimized for OLTP and is parameterizable, though also useful
for large scale ETL/ELT. MERGE is not intended to be used in preference
to existing single SQL commands for INSERT, UPDATE or DELETE since there
is some overhead.  MERGE can be used from PL/pgSQL.

MERGE does not support targetting updatable views or foreign tables, and
RETURNING clauses are not allowed either.  These limitations are likely
fixable with sufficient effort.  Rewrite rules are also not supported,
but it's not clear that we'd want to support them.

Author: Pavan Deolasee <pavan.deolasee@gmail.com>
Author: Álvaro Herrera <alvherre@alvh.no-ip.org>
Author: Amit Langote <amitlangote09@gmail.com>
Author: Simon Riggs <simon.riggs@enterprisedb.com>
Reviewed-by: Peter Eisentraut <peter.eisentraut@enterprisedb.com>
Reviewed-by: Andres Freund <andres@anarazel.de> (earlier versions)
Reviewed-by: Peter Geoghegan <pg@bowt.ie> (earlier versions)
Reviewed-by: Robert Haas <robertmhaas@gmail.com> (earlier versions)
Reviewed-by: Japin Li <japinli@hotmail.com>
Reviewed-by: Justin Pryzby <pryzby@telsasoft.com>
Reviewed-by: Tomas Vondra <tomas.vondra@enterprisedb.com>
Reviewed-by: Zhihong Yu <zyu@yugabyte.com>
Discussion: https://postgr.es/m/CANP8+jKitBSrB7oTgT9CY2i1ObfOt36z0XMraQc+Xrz8QB0nXA@mail.gmail.com
Discussion: https://postgr.es/m/CAH2-WzkJdBuxj9PO=2QaO9-3h3xGbQPZ34kJH=HukRekwM-GZg@mail.gmail.com
Discussion: https://postgr.es/m/20201231134736.GA25392@alvherre.pgsql

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);