1 files changed, 62 insertions, 39 deletions

diff --git a/src/backend/optimizer/prep/preptlist.c b/src/backend/optimizer/prep/preptlist.c
index afc733f1835..331e573f58e 100644
--- a/src/backend/optimizer/prep/preptlist.c
+++ b/src/backend/optimizer/prep/preptlist.c
@@ -4,20 +4,22 @@
  *	  Routines to preprocess the parse tree target list
  *
  * For INSERT and UPDATE queries, the targetlist must contain an entry for
- * each attribute of the target relation in the correct order.  For all query
+ * each attribute of the target relation in the correct order.  For UPDATE and
+ * DELETE queries, it must also contain junk tlist entries needed to allow the
+ * executor to identify the rows to be updated or deleted.  For all query
  * types, we may need to add junk tlist entries for Vars used in the RETURNING
  * list and row ID information needed for SELECT FOR UPDATE locking and/or
  * EvalPlanQual checking.
  *
- * The rewriter's rewriteTargetListIU and rewriteTargetListUD routines
- * also do preprocessing of the targetlist.  The division of labor between
- * here and there is partially historical, but it's not entirely arbitrary.
- * In particular, consider an UPDATE across an inheritance tree.  What the
- * rewriter does need be done only once (because it depends only on the
- * properties of the parent relation).  What's done here has to be done over
- * again for each child relation, because it depends on the column list of
- * the child, which might have more columns and/or a different column order
- * than the parent.
+ * The query rewrite phase also does preprocessing of the targetlist (see
+ * rewriteTargetListIU).  The division of labor between here and there is
+ * partially historical, but it's not entirely arbitrary.  In particular,
+ * consider an UPDATE across an inheritance tree.  What rewriteTargetListIU
+ * does need be done only once (because it depends only on the properties of
+ * the parent relation).  What's done here has to be done over again for each
+ * child relation, because it depends on the properties of the child, which
+ * might be of a different relation type, or have more columns and/or a
+ * different column order than the parent.
  *
  * The fact that rewriteTargetListIU sorts non-resjunk tlist entries by column
  * position, which expand_targetlist depends on, violates the above comment
@@ -47,11 +49,12 @@
 #include "optimizer/var.h"
 #include "parser/parsetree.h"
 #include "parser/parse_coerce.h"
+#include "rewrite/rewriteHandler.h"
 #include "utils/rel.h"
 
 
 static List *expand_targetlist(List *tlist, int command_type,
-				  Index result_relation, List *range_table);
+				  Index result_relation, Relation rel);
 
 
 /*
@@ -59,36 +62,61 @@ static List *expand_targetlist(List *tlist, int command_type,
  *	  Driver for preprocessing the parse tree targetlist.
  *
  *	  Returns the new targetlist.
+ *
+ * As a side effect, if there's an ON CONFLICT UPDATE clause, its targetlist
+ * is also preprocessed (and updated in-place).
  */
 List *
-preprocess_targetlist(PlannerInfo *root, List *tlist)
+preprocess_targetlist(PlannerInfo *root)
 {
 	Query	   *parse = root->parse;
 	int			result_relation = parse->resultRelation;
 	List	   *range_table = parse->rtable;
 	CmdType		command_type = parse->commandType;
+	RangeTblEntry *target_rte = NULL;
+	Relation	target_relation = NULL;
+	List	   *tlist;
 	ListCell   *lc;
 
 	/*
-	 * Sanity check: if there is a result relation, it'd better be a real
-	 * relation not a subquery.  Else parser or rewriter messed up.
+	 * If there is a result relation, open it so we can look for missing
+	 * columns and so on.  We assume that previous code already acquired at
+	 * least AccessShareLock on the relation, so we need no lock here.
 	 */
 	if (result_relation)
 	{
-		RangeTblEntry *rte = rt_fetch(result_relation, range_table);
+		target_rte = rt_fetch(result_relation, range_table);
+
+		/*
+		 * Sanity check: it'd better be a real relation not, say, a subquery.
+		 * Else parser or rewriter messed up.
+		 */
+		if (target_rte->rtekind != RTE_RELATION)
+			elog(ERROR, "result relation must be a regular relation");
 
-		if (rte->subquery != NULL || rte->relid == InvalidOid)
-			elog(ERROR, "subquery cannot be result relation");
+		target_relation = heap_open(target_rte->relid, NoLock);
 	}
+	else
+		Assert(command_type == CMD_SELECT);
+
+	/*
+	 * For UPDATE/DELETE, add any junk column(s) needed to allow the executor
+	 * to identify the rows to be updated or deleted.  Note that this step
+	 * scribbles on parse->targetList, which is not very desirable, but we
+	 * keep it that way to avoid changing APIs used by FDWs.
+	 */
+	if (command_type == CMD_UPDATE || command_type == CMD_DELETE)
+		rewriteTargetListUD(parse, target_rte, target_relation);
 
 	/*
 	 * for heap_form_tuple to work, the targetlist must match the exact order
 	 * of the attributes. We also need to fill in any missing attributes. -ay
 	 * 10/94
 	 */
+	tlist = parse->targetList;
 	if (command_type == CMD_INSERT || command_type == CMD_UPDATE)
 		tlist = expand_targetlist(tlist, command_type,
-								  result_relation, range_table);
+								  result_relation, target_relation);
 
 	/*
 	 * Add necessary junk columns for rowmarked rels.  These values are needed
@@ -193,19 +221,21 @@ preprocess_targetlist(PlannerInfo *root, List *tlist)
 		list_free(vars);
 	}
 
-	return tlist;
-}
+	/*
+	 * If there's an ON CONFLICT UPDATE clause, preprocess its targetlist too
+	 * while we have the relation open.
+	 */
+	if (parse->onConflict)
+		parse->onConflict->onConflictSet =
+			expand_targetlist(parse->onConflict->onConflictSet,
+							  CMD_UPDATE,
+							  result_relation,
+							  target_relation);
 
-/*
- * preprocess_onconflict_targetlist
- *	  Process ON CONFLICT SET targetlist.
- *
- *	  Returns the new targetlist.
- */
-List *
-preprocess_onconflict_targetlist(List *tlist, int result_relation, List *range_table)
-{
-	return expand_targetlist(tlist, CMD_UPDATE, result_relation, range_table);
+	if (target_relation)
+		heap_close(target_relation, NoLock);
+
+	return tlist;
 }
 
 
@@ -223,11 +253,10 @@ preprocess_onconflict_targetlist(List *tlist, int result_relation, List *range_t
  */
 static List *
 expand_targetlist(List *tlist, int command_type,
-				  Index result_relation, List *range_table)
+				  Index result_relation, Relation rel)
 {
 	List	   *new_tlist = NIL;
 	ListCell   *tlist_item;
-	Relation	rel;
 	int			attrno,
 				numattrs;
 
@@ -238,12 +267,8 @@ expand_targetlist(List *tlist, int command_type,
 	 * order; but we have to insert TLEs for any missing attributes.
 	 *
 	 * Scan the tuple description in the relation's relcache entry to make
-	 * sure we have all the user attributes in the right order.  We assume
-	 * that the rewriter already acquired at least AccessShareLock on the
-	 * relation, so we need no lock here.
+	 * sure we have all the user attributes in the right order.
 	 */
-	rel = heap_open(getrelid(result_relation, range_table), NoLock);
-
 	numattrs = RelationGetNumberOfAttributes(rel);
 
 	for (attrno = 1; attrno <= numattrs; attrno++)
@@ -386,8 +411,6 @@ expand_targetlist(List *tlist, int command_type,
 		tlist_item = lnext(tlist_item);
 	}
 
-	heap_close(rel, NoLock);
-
 	return new_tlist;
 }