0001-Initialize-WITH-CHECK-OPTIONS-and-RETURNING-expressi.patch

From c6456053de33fd43f9389ab056a0e67670be2f43 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Fri, 9 Apr 2021 15:23:56 +0900
Subject: [PATCH] Initialize WITH CHECK OPTIONS and RETURNING expressions on
 demand

---
 src/backend/executor/execMain.c               |  19 +-
 src/backend/executor/execPartition.c          |  33 +--
 src/backend/executor/nodeModifyTable.c        | 210 ++++++++++--------
 src/include/executor/executor.h               |   3 +-
 src/test/regress/expected/updatable_views.out |  18 +-
 5 files changed, 152 insertions(+), 131 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index b2e2df8773..fca60736f0 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -1950,7 +1950,8 @@ ExecConstraints(ResultRelInfo *resultRelInfo,
  */
 void
 ExecWithCheckOptions(WCOKind kind, ResultRelInfo *resultRelInfo,
-					 TupleTableSlot *slot, EState *estate)
+					 TupleTableSlot *slot, EState *estate,
+					 PlanState *parent)
 {
 	Relation	rel = resultRelInfo->ri_RelationDesc;
 	TupleDesc	tupdesc = RelationGetDescr(rel);
@@ -1958,6 +1959,22 @@ ExecWithCheckOptions(WCOKind kind, ResultRelInfo *resultRelInfo,
 	ListCell   *l1,
 			   *l2;
 
+	/* Initialize the quals if not already done. */
+	if (resultRelInfo->ri_WithCheckOptionExprs == NIL)
+	{
+		List   *wcoExprs = NIL;
+
+		foreach(l1, resultRelInfo->ri_WithCheckOptions)
+		{
+			WithCheckOption *wco = (WithCheckOption *) lfirst(l1);
+			ExprState  *wcoExpr = ExecInitQual((List *) wco->qual,
+											   parent);
+
+			wcoExprs = lappend(wcoExprs, wcoExpr);
+		}
+		resultRelInfo->ri_WithCheckOptionExprs = wcoExprs;
+	}
+
 	/*
 	 * We will use the EState's per-tuple context for evaluating constraint
 	 * expressions (creating it if it's not already there).
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 99780ebb96..7dccc55006 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -543,8 +543,6 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 	if (node && node->withCheckOptionLists != NIL)
 	{
 		List	   *wcoList;
-		List	   *wcoExprs = NIL;
-		ListCell   *ll;
 
 		/*
 		 * In the case of INSERT on a partitioned table, there is only one
@@ -582,18 +580,11 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 								RelationGetForm(partrel)->reltype,
 								&found_whole_row);
 		/* We ignore the value of found_whole_row. */
-
-		foreach(ll, wcoList)
-		{
-			WithCheckOption *wco = castNode(WithCheckOption, lfirst(ll));
-			ExprState  *wcoExpr = ExecInitQual(castNode(List, wco->qual),
-											   &mtstate->ps);
-
-			wcoExprs = lappend(wcoExprs, wcoExpr);
-		}
-
 		leaf_part_rri->ri_WithCheckOptions = wcoList;
-		leaf_part_rri->ri_WithCheckOptionExprs = wcoExprs;
+		/*
+		 * ri_WithCheckOptionExprs is built the first time it needs to be
+		 * used (see ExecWithCheckOptions()).
+		 */
 	}
 
 	/*
@@ -605,8 +596,6 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 	 */
 	if (node && node->returningLists != NIL)
 	{
-		TupleTableSlot *slot;
-		ExprContext *econtext;
 		List	   *returningList;
 
 		/* See the comment above for WCO lists. */
@@ -641,20 +630,10 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 		/* We ignore the value of found_whole_row. */
 
 		leaf_part_rri->ri_returningList = returningList;
-
 		/*
-		 * Initialize the projection itself.
-		 *
-		 * Use the slot and the expression context that would have been set up
-		 * in ExecInitModifyTable() for projection's output.
+		 * ri_projectReturning is built the first time it needs to be
+		 * used (see ExecProcessReturning()).
 		 */
-		Assert(mtstate->ps.ps_ResultTupleSlot != NULL);
-		slot = mtstate->ps.ps_ResultTupleSlot;
-		Assert(mtstate->ps.ps_ExprContext != NULL);
-		econtext = mtstate->ps.ps_ExprContext;
-		leaf_part_rri->ri_projectReturning =
-			ExecBuildProjectionInfo(returningList, econtext, slot,
-									&mtstate->ps, RelationGetDescr(partrel));
 	}
 
 	/* Set up information needed for routing tuples to the partition. */
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index c5a2a9a054..dba9f684a5 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -154,12 +154,37 @@ ExecCheckPlanOutput(Relation resultRel, List *targetList)
 				 errdetail("Query has too few columns.")));
 }
 
+/* Initializes the RETURNING projection for given result relation. */
+static void
+ExecInitReturningProjection(ResultRelInfo *resultRelInfo,
+							ModifyTableState *mtstate)
+{
+	TupleTableSlot *slot;
+	ExprContext *econtext;
+
+	/* should not get called twice */
+	Assert(resultRelInfo->ri_projectReturning == NULL);
+
+	/* ExecInitModifyTable() should've initialized these. */
+	Assert(mtstate->ps.ps_ExprContext != NULL);
+	econtext = mtstate->ps.ps_ExprContext;
+	Assert(mtstate->ps.ps_ResultTupleSlot != NULL);
+	slot = mtstate->ps.ps_ResultTupleSlot;
+
+	Assert(resultRelInfo->ri_returningList != NIL);
+	resultRelInfo->ri_projectReturning =
+		ExecBuildProjectionInfo(resultRelInfo->ri_returningList,
+								econtext, slot, &mtstate->ps,
+								resultRelInfo->ri_RelationDesc->rd_att);
+}
+
 /*
  * ExecProcessReturning --- evaluate a RETURNING list
  *
  * resultRelInfo: current result rel
  * tupleSlot: slot holding tuple actually inserted/updated/deleted
  * planSlot: slot holding tuple returned by top subplan node
+ * mtstate: query's plan state
  *
  * Note: If tupleSlot is NULL, the FDW should have already provided econtext's
  * scan tuple.
@@ -169,10 +194,18 @@ ExecCheckPlanOutput(Relation resultRel, List *targetList)
 static TupleTableSlot *
 ExecProcessReturning(ResultRelInfo *resultRelInfo,
 					 TupleTableSlot *tupleSlot,
-					 TupleTableSlot *planSlot)
+					 TupleTableSlot *planSlot,
+					 ModifyTableState *mtstate)
 {
-	ProjectionInfo *projectReturning = resultRelInfo->ri_projectReturning;
-	ExprContext *econtext = projectReturning->pi_exprContext;
+	ProjectionInfo *projectReturning;
+	ExprContext *econtext;
+
+	/* Initialize the projection if not already done. */
+	if (resultRelInfo->ri_projectReturning == NULL)
+		ExecInitReturningProjection(resultRelInfo, mtstate);
+
+	projectReturning = resultRelInfo->ri_projectReturning;
+	econtext = projectReturning->pi_exprContext;
 
 	/* Make tuple and any needed join variables available to ExecProject */
 	if (tupleSlot)
@@ -767,7 +800,8 @@ ExecInsert(ModifyTableState *mtstate,
 		 * we are looking for at this point.
 		 */
 		if (resultRelInfo->ri_WithCheckOptions != NIL)
-			ExecWithCheckOptions(wco_kind, resultRelInfo, slot, estate);
+			ExecWithCheckOptions(wco_kind, resultRelInfo, slot, estate,
+								 &mtstate->ps);
 
 		/*
 		 * Check the constraints of the tuple.
@@ -962,11 +996,12 @@ ExecInsert(ModifyTableState *mtstate,
 	 * are looking for at this point.
 	 */
 	if (resultRelInfo->ri_WithCheckOptions != NIL)
-		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate);
+		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate,
+							 &mtstate->ps);
 
 	/* Process RETURNING if present */
-	if (resultRelInfo->ri_projectReturning)
-		result = ExecProcessReturning(resultRelInfo, slot, planSlot);
+	if (resultRelInfo->ri_returningList)
+		result = ExecProcessReturning(resultRelInfo, slot, planSlot, mtstate);
 
 	return result;
 }
@@ -1022,7 +1057,8 @@ ExecBatchInsert(ModifyTableState *mtstate,
 		 * comment in ExecInsert.
 		 */
 		if (resultRelInfo->ri_WithCheckOptions != NIL)
-			ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate);
+			ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate,
+								 &mtstate->ps);
 	}
 
 	if (canSetTag && numInserted > 0)
@@ -1345,7 +1381,7 @@ ldelete:;
 						 ar_delete_trig_tcs);
 
 	/* Process RETURNING if present and if requested */
-	if (processReturning && resultRelInfo->ri_projectReturning)
+	if (processReturning && resultRelInfo->ri_returningList)
 	{
 		/*
 		 * We have to put the target tuple into a slot, which means first we
@@ -1373,7 +1409,7 @@ ldelete:;
 			}
 		}
 
-		rslot = ExecProcessReturning(resultRelInfo, slot, planSlot);
+		rslot = ExecProcessReturning(resultRelInfo, slot, planSlot, mtstate);
 
 		/*
 		 * Before releasing the target tuple again, make sure rslot has a
@@ -1707,7 +1743,8 @@ lreplace:;
 			 * kind we are looking for at this point.
 			 */
 			ExecWithCheckOptions(WCO_RLS_UPDATE_CHECK,
-								 resultRelInfo, slot, estate);
+								 resultRelInfo, slot, estate,
+								 &mtstate->ps);
 		}
 
 		/*
@@ -1934,11 +1971,12 @@ lreplace:;
 	 * are looking for at this point.
 	 */
 	if (resultRelInfo->ri_WithCheckOptions != NIL)
-		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate);
+		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate,
+							 &mtstate->ps);
 
 	/* Process RETURNING if present */
-	if (resultRelInfo->ri_projectReturning)
-		return ExecProcessReturning(resultRelInfo, slot, planSlot);
+	if (resultRelInfo->ri_returningList)
+		return ExecProcessReturning(resultRelInfo, slot, planSlot, mtstate);
 
 	return NULL;
 }
@@ -2132,7 +2170,8 @@ ExecOnConflictUpdate(ModifyTableState *mtstate,
 		 */
 		ExecWithCheckOptions(WCO_RLS_CONFLICT_CHECK, resultRelInfo,
 							 existing,
-							 mtstate->ps.state);
+							 mtstate->ps.state,
+							 &mtstate->ps);
 	}
 
 	/* Project the new tuple version */
@@ -2439,7 +2478,7 @@ ExecModifyTable(PlanState *pstate)
 			 * ExecProcessReturning by IterateDirectModify, so no need to
 			 * provide it here.
 			 */
-			slot = ExecProcessReturning(resultRelInfo, NULL, planSlot);
+			slot = ExecProcessReturning(resultRelInfo, NULL, planSlot, node);
 
 			return slot;
 		}
@@ -2748,6 +2787,33 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	if (!(eflags & EXEC_FLAG_EXPLAIN_ONLY))
 		ExecSetupTransitionCaptureState(mtstate, estate);
 
+	/*
+	 * Initialize result tuple slot and assign its rowtype using the first
+	 * RETURNING list.  We assume the rest will look the same.
+	 */
+	if (node->returningLists)
+	{
+		mtstate->ps.plan->targetlist = (List *) linitial(node->returningLists);
+
+		/* Set up a slot for the output of the RETURNING projection(s) */
+		ExecInitResultTupleSlotTL(&mtstate->ps, &TTSOpsVirtual);
+
+		/* Need an econtext too */
+		if (mtstate->ps.ps_ExprContext == NULL)
+			ExecAssignExprContext(estate, &mtstate->ps);
+	}
+	else
+	{
+		/*
+		 * We still must construct a dummy result tuple type, because InitPlan
+		 * expects one (maybe should change that?).
+		 */
+		mtstate->ps.plan->targetlist = NIL;
+		ExecInitResultTypeTL(&mtstate->ps);
+
+		mtstate->ps.ps_ExprContext = NULL;
+	}
+
 	/*
 	 * Open all the result relations and initialize the ResultRelInfo structs.
 	 * (But root relation was initialized above, if it's part of the array.)
@@ -2856,6 +2922,44 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 					elog(ERROR, "could not find junk wholerow column");
 			}
 		}
+
+		/*
+		 * Initialize any WITH CHECK OPTION constraints if needed.
+		 */
+		if (node->withCheckOptionLists)
+		{
+			List   *wcoList = (List *) list_nth(node->withCheckOptionLists, i);
+
+			resultRelInfo->ri_WithCheckOptions = wcoList;
+			/*
+			 * ri_WithCheckOptionExprs is built the first time it needs to be
+			 * used (see ExecWithCheckOptions()).
+			 */
+		}
+
+		/*
+		 * Initialize RETURNING projections if needed.
+		 */
+		if (node->returningLists)
+		{
+			List  *rlist = (List *) list_nth(node->returningLists, i);
+
+			resultRelInfo->ri_returningList = rlist;
+
+			/*
+			 * ri_projectReturning is built the first time it needs to be used
+			 * (see ExecProcessReturning()), unless the relation is to be
+			 * "directly modified", in that case, IterateDirectModify()
+			 * expects ri_projectReturning to be valid.  Because the child
+			 * relations (their ResultRelInfos) that are directly modified are
+			 * accessed from their corresponding ForeignScanState nodes,
+			 * they can get switched without the control returning to the top-
+			 * level ModifyTable, so there's no better way than initializing
+			 * the projection for all such result relations here.
+			 */
+			if (resultRelInfo->ri_usesFdwDirectModify)
+				ExecInitReturningProjection(resultRelInfo, mtstate);
+		}
 	}
 
 	/*
@@ -2883,80 +2987,6 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 		mtstate->mt_partition_tuple_routing =
 			ExecSetupPartitionTupleRouting(estate, rel);
 
-	/*
-	 * Initialize any WITH CHECK OPTION constraints if needed.
-	 */
-	resultRelInfo = mtstate->resultRelInfo;
-	foreach(l, node->withCheckOptionLists)
-	{
-		List	   *wcoList = (List *) lfirst(l);
-		List	   *wcoExprs = NIL;
-		ListCell   *ll;
-
-		foreach(ll, wcoList)
-		{
-			WithCheckOption *wco = (WithCheckOption *) lfirst(ll);
-			ExprState  *wcoExpr = ExecInitQual((List *) wco->qual,
-											   &mtstate->ps);
-
-			wcoExprs = lappend(wcoExprs, wcoExpr);
-		}
-
-		resultRelInfo->ri_WithCheckOptions = wcoList;
-		resultRelInfo->ri_WithCheckOptionExprs = wcoExprs;
-		resultRelInfo++;
-	}
-
-	/*
-	 * Initialize RETURNING projections if needed.
-	 */
-	if (node->returningLists)
-	{
-		TupleTableSlot *slot;
-		ExprContext *econtext;
-
-		/*
-		 * Initialize result tuple slot and assign its rowtype using the first
-		 * RETURNING list.  We assume the rest will look the same.
-		 */
-		mtstate->ps.plan->targetlist = (List *) linitial(node->returningLists);
-
-		/* Set up a slot for the output of the RETURNING projection(s) */
-		ExecInitResultTupleSlotTL(&mtstate->ps, &TTSOpsVirtual);
-		slot = mtstate->ps.ps_ResultTupleSlot;
-
-		/* Need an econtext too */
-		if (mtstate->ps.ps_ExprContext == NULL)
-			ExecAssignExprContext(estate, &mtstate->ps);
-		econtext = mtstate->ps.ps_ExprContext;
-
-		/*
-		 * Build a projection for each result rel.
-		 */
-		resultRelInfo = mtstate->resultRelInfo;
-		foreach(l, node->returningLists)
-		{
-			List	   *rlist = (List *) lfirst(l);
-
-			resultRelInfo->ri_returningList = rlist;
-			resultRelInfo->ri_projectReturning =
-				ExecBuildProjectionInfo(rlist, econtext, slot, &mtstate->ps,
-										resultRelInfo->ri_RelationDesc->rd_att);
-			resultRelInfo++;
-		}
-	}
-	else
-	{
-		/*
-		 * We still must construct a dummy result tuple type, because InitPlan
-		 * expects one (maybe should change that?).
-		 */
-		mtstate->ps.plan->targetlist = NIL;
-		ExecInitResultTypeTL(&mtstate->ps);
-
-		mtstate->ps.ps_ExprContext = NULL;
-	}
-
 	/* Set the list of arbiter indexes if needed for ON CONFLICT */
 	resultRelInfo = mtstate->resultRelInfo;
 	if (node->onConflictAction != ONCONFLICT_NONE)
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 6eae134c08..622cb26c88 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -211,7 +211,8 @@ extern bool ExecPartitionCheck(ResultRelInfo *resultRelInfo,
 extern void ExecPartitionCheckEmitError(ResultRelInfo *resultRelInfo,
 										TupleTableSlot *slot, EState *estate);
 extern void ExecWithCheckOptions(WCOKind kind, ResultRelInfo *resultRelInfo,
-								 TupleTableSlot *slot, EState *estate);
+								 TupleTableSlot *slot, EState *estate,
+								 PlanState *parent);
 extern LockTupleMode ExecUpdateLockMode(EState *estate, ResultRelInfo *relinfo);
 extern ExecRowMark *ExecFindRowMark(EState *estate, Index rti, bool missing_ok);
 extern ExecAuxRowMark *ExecBuildAuxRowMark(ExecRowMark *erm, List *targetlist);
diff --git a/src/test/regress/expected/updatable_views.out b/src/test/regress/expected/updatable_views.out
index cdff914b93..dd309283be 100644
--- a/src/test/regress/expected/updatable_views.out
+++ b/src/test/regress/expected/updatable_views.out
@@ -1892,28 +1892,22 @@ UPDATE rw_view1 SET a = a + 5; -- should fail
 ERROR:  new row violates check option for view "rw_view1"
 DETAIL:  Failing row contains (15).
 EXPLAIN (costs off) INSERT INTO rw_view1 VALUES (5);
-                       QUERY PLAN                        
----------------------------------------------------------
+      QUERY PLAN      
+----------------------
  Insert on base_tbl b
    ->  Result
-   SubPlan 1
-     ->  Index Only Scan using ref_tbl_pkey on ref_tbl r
-           Index Cond: (a = b.a)
-(5 rows)
+(2 rows)
 
 EXPLAIN (costs off) UPDATE rw_view1 SET a = a + 5;
-                        QUERY PLAN                         
------------------------------------------------------------
+               QUERY PLAN                
+-----------------------------------------
  Update on base_tbl b
    ->  Hash Join
          Hash Cond: (b.a = r.a)
          ->  Seq Scan on base_tbl b
          ->  Hash
                ->  Seq Scan on ref_tbl r
-   SubPlan 1
-     ->  Index Only Scan using ref_tbl_pkey on ref_tbl r_1
-           Index Cond: (a = b.a)
-(9 rows)
+(6 rows)
 
 DROP TABLE base_tbl, ref_tbl CASCADE;
 NOTICE:  drop cascades to view rw_view1
-- 
2.24.1