v4-0001-Revise-BeginDirectModify-API-to-pass-ResultRelInf.patch

From f7231c6d04da106321f8b43b1a55dc10b8cb2b82 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Wed, 31 Jul 2019 16:38:43 +0900
Subject: [PATCH v4 1/4] Revise BeginDirectModify API to pass ResultRelInfo
 directly

For ExecInitForeignScan() to efficiently get the ResultRelInfo to
pass to BeginDirectModify(), add a field to ForeignScan that gives
the index of a given result relation in the query's list of result
relations.
---
 contrib/postgres_fdw/postgres_fdw.c     | 22 ++++++++++++++++------
 doc/src/sgml/fdwhandler.sgml            | 11 ++++++++++-
 src/backend/executor/nodeForeignscan.c  | 13 ++++++++++---
 src/backend/nodes/copyfuncs.c           |  1 +
 src/backend/nodes/outfuncs.c            |  1 +
 src/backend/optimizer/plan/createplan.c |  2 ++
 src/backend/optimizer/plan/setrefs.c    | 28 ++++++++++++++++++++++------
 src/include/foreign/fdwapi.h            |  1 +
 src/include/nodes/plannodes.h           |  4 ++++
 9 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 033aeb2556..1b60ff88b1 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -205,6 +205,9 @@ typedef struct PgFdwDirectModifyState
 	List	   *retrieved_attrs;	/* attr numbers retrieved by RETURNING */
 	bool		set_processed;	/* do we set the command es_processed? */
 
+	/* Information about the relation being modified */
+	ResultRelInfo *resultRelInfo;
+
 	/* for remote query execution */
 	PGconn	   *conn;			/* connection for the update */
 	int			numParams;		/* number of parameters passed to query */
@@ -360,7 +363,9 @@ static bool postgresPlanDirectModify(PlannerInfo *root,
 									 ModifyTable *plan,
 									 Index resultRelation,
 									 int subplan_index);
-static void postgresBeginDirectModify(ForeignScanState *node, int eflags);
+static void postgresBeginDirectModify(ForeignScanState *node,
+						  ResultRelInfo *rinfo,
+						  int eflags);
 static TupleTableSlot *postgresIterateDirectModify(ForeignScanState *node);
 static void postgresEndDirectModify(ForeignScanState *node);
 static void postgresExplainForeignScan(ForeignScanState *node,
@@ -2340,7 +2345,9 @@ postgresPlanDirectModify(PlannerInfo *root,
  *		Prepare a direct foreign table modification
  */
 static void
-postgresBeginDirectModify(ForeignScanState *node, int eflags)
+postgresBeginDirectModify(ForeignScanState *node,
+						  ResultRelInfo *rinfo,
+						  int eflags)
 {
 	ForeignScan *fsplan = (ForeignScan *) node->ss.ps.plan;
 	EState	   *estate = node->ss.ps.state;
@@ -2368,7 +2375,7 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 	 * Identify which user to do the remote access as.  This should match what
 	 * ExecCheckRTEPerms() does.
 	 */
-	rtindex = estate->es_result_relation_info->ri_RangeTableIndex;
+	rtindex = rinfo->ri_RangeTableIndex;
 	rte = exec_rt_fetch(rtindex, estate);
 	userid = rte->checkAsUser ? rte->checkAsUser : GetUserId();
 
@@ -2414,6 +2421,9 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 	dmstate->set_processed = intVal(list_nth(fsplan->fdw_private,
 											 FdwDirectModifyPrivateSetProcessed));
 
+	/* Save the ResultRelInfo of the relation being modified. */
+	dmstate->resultRelInfo = rinfo;
+
 	/* Create context for per-tuple temp workspace. */
 	dmstate->temp_cxt = AllocSetContextCreate(estate->es_query_cxt,
 											  "postgres_fdw temporary data",
@@ -2463,7 +2473,7 @@ postgresIterateDirectModify(ForeignScanState *node)
 {
 	PgFdwDirectModifyState *dmstate = (PgFdwDirectModifyState *) node->fdw_state;
 	EState	   *estate = node->ss.ps.state;
-	ResultRelInfo *resultRelInfo = estate->es_result_relation_info;
+	ResultRelInfo *resultRelInfo = dmstate->resultRelInfo;
 
 	/*
 	 * If this is the first call after Begin, execute the statement.
@@ -4033,7 +4043,7 @@ get_returning_data(ForeignScanState *node)
 {
 	PgFdwDirectModifyState *dmstate = (PgFdwDirectModifyState *) node->fdw_state;
 	EState	   *estate = node->ss.ps.state;
-	ResultRelInfo *resultRelInfo = estate->es_result_relation_info;
+	ResultRelInfo *resultRelInfo = dmstate->resultRelInfo;
 	TupleTableSlot *slot = node->ss.ss_ScanTupleSlot;
 	TupleTableSlot *resultSlot;
 
@@ -4180,7 +4190,7 @@ apply_returning_filter(PgFdwDirectModifyState *dmstate,
 					   TupleTableSlot *slot,
 					   EState *estate)
 {
-	ResultRelInfo *relInfo = estate->es_result_relation_info;
+	ResultRelInfo *relInfo = dmstate->resultRelInfo;
 	TupleDesc	resultTupType = RelationGetDescr(dmstate->resultRel);
 	TupleTableSlot *resultSlot;
 	Datum	   *values;
diff --git a/doc/src/sgml/fdwhandler.sgml b/doc/src/sgml/fdwhandler.sgml
index 27b94fb611..04c2eccd1c 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -871,6 +871,7 @@ PlanDirectModify(PlannerInfo *root,
 <programlisting>
 void
 BeginDirectModify(ForeignScanState *node,
+                  ResultRelInfo *rinfo,
                   int eflags);
 </programlisting>
 
@@ -883,7 +884,9 @@ BeginDirectModify(ForeignScanState *node,
      the table to modify is accessible through the
      <structname>ForeignScanState</structname> node (in particular, from the underlying
      <structname>ForeignScan</structname> plan node, which contains any FDW-private
-     information provided by <function>PlanDirectModify</function>).
+     information provided by <function>PlanDirectModify</function>).  In
+     addition, <literal>rinfo</literal> also contains information describing
+     the target foreign table.
      <literal>eflags</literal> contains flag bits describing the executor's
      operating mode for this plan node.
     </para>
@@ -895,6 +898,12 @@ BeginDirectModify(ForeignScanState *node,
      for <function>ExplainDirectModify</function> and <function>EndDirectModify</function>.
     </para>
 
+    <note>
+     Also note that it's a good idea to store the <literal>rinfo</literal>
+     in the <structfield>fdw_state</structfield> for
+     <function>IterateDirectModify</function> to use.
+    </node>
+
     <para>
      If the <function>BeginDirectModify</function> pointer is set to
      <literal>NULL</literal>, no attempts to execute a direct modification on the
diff --git a/src/backend/executor/nodeForeignscan.c b/src/backend/executor/nodeForeignscan.c
index 52af1dac5c..9824c16e09 100644
--- a/src/backend/executor/nodeForeignscan.c
+++ b/src/backend/executor/nodeForeignscan.c
@@ -223,10 +223,17 @@ ExecInitForeignScan(ForeignScan *node, EState *estate, int eflags)
 	/*
 	 * Tell the FDW to initialize the scan.
 	 */
-	if (node->operation != CMD_SELECT)
-		fdwroutine->BeginDirectModify(scanstate, eflags);
-	else
+	if (node->operation == CMD_SELECT)
 		fdwroutine->BeginForeignScan(scanstate, eflags);
+	else
+	{
+		/* Perform initializations for a direct modification. */
+		ResultRelInfo *resultRelInfo;
+
+		Assert(node->resultRelIndex >= 0);
+		resultRelInfo = &estate->es_result_relations[node->resultRelIndex];
+		fdwroutine->BeginDirectModify(scanstate, resultRelInfo, eflags);
+	}
 
 	return scanstate;
 }
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index a2617c7cfd..e981298a75 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -758,6 +758,7 @@ _copyForeignScan(const ForeignScan *from)
 	COPY_NODE_FIELD(fdw_recheck_quals);
 	COPY_BITMAPSET_FIELD(fs_relids);
 	COPY_SCALAR_FIELD(fsSystemCol);
+	COPY_SCALAR_FIELD(resultRelIndex);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index e6ce8e2110..80eedc4a24 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -695,6 +695,7 @@ _outForeignScan(StringInfo str, const ForeignScan *node)
 	WRITE_NODE_FIELD(fdw_recheck_quals);
 	WRITE_BITMAPSET_FIELD(fs_relids);
 	WRITE_BOOL_FIELD(fsSystemCol);
+	WRITE_INT_FIELD(resultRelIndex);
 }
 
 static void
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index f2325694c5..ff78167f79 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -5455,6 +5455,8 @@ make_foreignscan(List *qptlist,
 	node->fs_relids = NULL;
 	/* fsSystemCol will be filled in by create_foreignscan_plan */
 	node->fsSystemCol = false;
+	/* this might be filled to a >= 0 value by set_plan_refs() */
+	node->resultRelIndex = -1;
 
 	return node;
 }
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index 329ebd5f28..b233eb7dce 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -781,6 +781,7 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 		case T_ModifyTable:
 			{
 				ModifyTable *splan = (ModifyTable *) plan;
+				int		resultRelIndex;
 
 				Assert(splan->plan.targetlist == NIL);
 				Assert(splan->plan.qual == NIL);
@@ -877,12 +878,6 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 					rc->rti += rtoffset;
 					rc->prti += rtoffset;
 				}
-				foreach(l, splan->plans)
-				{
-					lfirst(l) = set_plan_refs(root,
-											  (Plan *) lfirst(l),
-											  rtoffset);
-				}
 
 				/*
 				 * Append this ModifyTable node's final result relation RT
@@ -908,6 +903,27 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 						lappend_int(root->glob->rootResultRelations,
 									splan->rootRelation);
 				}
+
+				resultRelIndex = splan->resultRelIndex;
+				foreach(l, splan->plans)
+				{
+					lfirst(l) = set_plan_refs(root,
+											  (Plan *) lfirst(l),
+											  rtoffset);
+
+					/*
+					 * For foreign table result relations, save their index
+					 * in the global list of result relations into the
+					 * corresponding ForeignScan nodes.
+					 */
+					if (IsA(lfirst(l), ForeignScan))
+					{
+						ForeignScan *fscan = (ForeignScan *) lfirst(l);
+
+						fscan->resultRelIndex = resultRelIndex;
+					}
+					resultRelIndex++;
+				}
 			}
 			break;
 		case T_Append:
diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h
index 822686033e..adf39bc618 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -112,6 +112,7 @@ typedef bool (*PlanDirectModify_function) (PlannerInfo *root,
 										   int subplan_index);
 
 typedef void (*BeginDirectModify_function) (ForeignScanState *node,
+											ResultRelInfo *rinfo,
 											int eflags);
 
 typedef TupleTableSlot *(*IterateDirectModify_function) (ForeignScanState *node);
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 8e6594e355..047dd73dd1 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -616,6 +616,10 @@ typedef struct ForeignScan
 	List	   *fdw_recheck_quals;	/* original quals not in scan.plan.qual */
 	Bitmapset  *fs_relids;		/* RTIs generated by this scan */
 	bool		fsSystemCol;	/* true if any "system column" is needed */
+	int			resultRelIndex;	/* For non-SELECT operations, this contains
+								 * the offset of result relation in a
+								 * query-global list of result relations; -1
+								 * otherwise */
 } ForeignScan;
 
 /* ----------------
-- 
2.11.0