From db6849c6ef3c3329911bcba5d75a88d86164ca7b Mon Sep 17 00:00:00 2001
From: Etsuro Fujita <efujita@postgresql.org>
Date: Thu, 8 Aug 2019 21:41:12 +0900
Subject: [PATCH v11 1/4] Remove dependency on estate->es_result_relation_info
 from FDW APIs.

FDW APIs for executing a foreign table direct modification assumed that
the FDW would obtain the target foreign table's ResultRelInfo from
estate->es_result_relation_info of the passed-in ForeignScanState node,
but the upcoming patch(es) to refactor partitioning-related code in
nodeModifyTable.c will remove the es_result_relation_info variable.
Revise BeginDirectModify()'s API to pass the ResultRelInfo explicitly, to
remove the dependency on that variable from the FDW APIs.  For
ExecInitForeignScan() to efficiently get the ResultRelInfo to pass to
BeginDirectModify(), add a field to ForeignScan that gives the index of
the target foreign table in the list of the query result relations.

Patch by Amit Langote, following a proposal by Andres Freund, reviewed by
Andres Freund and me

Discussion: https://postgr.es/m/20190718010911.l6xcdv6birtxiei4@alap3.anarazel.de
---
 contrib/postgres_fdw/postgres_fdw.c     | 25 +++++++++++++++++++------
 doc/src/sgml/fdwhandler.sgml            |  8 ++++++--
 src/backend/executor/nodeForeignscan.c  | 14 ++++++++++----
 src/backend/nodes/copyfuncs.c           |  1 +
 src/backend/nodes/outfuncs.c            |  1 +
 src/backend/nodes/readfuncs.c           |  1 +
 src/backend/optimizer/plan/createplan.c |  2 ++
 src/backend/optimizer/plan/setrefs.c    | 15 +++++++++++++++
 src/include/foreign/fdwapi.h            |  1 +
 src/include/nodes/plannodes.h           |  3 +++
 10 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 9fc53ca..3af5b24 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -218,6 +218,7 @@ typedef struct PgFdwDirectModifyState
 	int			num_tuples;		/* # of result tuples */
 	int			next_tuple;		/* index of next one to return */
 	Relation	resultRel;		/* relcache entry for the target relation */
+	ResultRelInfo *resultRelInfo;	/* ResultRelInfo for the target relation */
 	AttrNumber *attnoMap;		/* array of attnums of input user columns */
 	AttrNumber	ctidAttno;		/* attnum of input ctid column */
 	AttrNumber	oidAttno;		/* attnum of input oid column */
@@ -361,7 +362,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,
@@ -2319,6 +2322,11 @@ postgresPlanDirectModify(PlannerInfo *root,
 			rebuild_fdw_scan_tlist(fscan, returningList);
 	}
 
+	/*
+	 * Set the index of the subplan result rel.
+	 */
+	fscan->resultRelIndex = subplan_index;
+
 	table_close(rel, NoLock);
 	return true;
 }
@@ -2328,7 +2336,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;
@@ -2356,7 +2366,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();
 
@@ -2389,6 +2399,9 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 		dmstate->rel = NULL;
 	}
 
+	/* Save the ResultRelInfo for the target relation. */
+	dmstate->resultRelInfo = rinfo;
+
 	/* Initialize state variable */
 	dmstate->num_tuples = -1;	/* -1 means not set yet */
 
@@ -2451,7 +2464,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.
@@ -4087,7 +4100,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;
 
@@ -4234,7 +4247,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 7479303..d2cfc6c 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -881,6 +881,7 @@ PlanDirectModify(PlannerInfo *root,
 <programlisting>
 void
 BeginDirectModify(ForeignScanState *node,
+                  ResultRelInfo *rinfo,
                   int eflags);
 </programlisting>
 
@@ -894,6 +895,8 @@ BeginDirectModify(ForeignScanState *node,
      <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>).
+     In addition, the <structname>ResultRelInfo</structname> struct also
+     contains information about the target foreign table.
      <literal>eflags</literal> contains flag bits describing the executor's
      operating mode for this plan node.
     </para>
@@ -925,8 +928,9 @@ IterateDirectModify(ForeignScanState *node);
      tuple table slot (the node's <structfield>ScanTupleSlot</structfield> should be
      used for this purpose).  The data that was actually inserted, updated
      or deleted must be stored in the
-     <literal>es_result_relation_info-&gt;ri_projectReturning-&gt;pi_exprContext-&gt;ecxt_scantuple</literal>
-     of the node's <structname>EState</structname>.
+     <literal>ri_projectReturning-&gt;pi_exprContext-&gt;ecxt_scantuple</literal>
+     of the target foreign table's <structname>ResultRelInfo</structname>
+     passed to <function>BeginDirectModify</function>.
      Return NULL if no more rows are available.
      Note that this is called in a short-lived memory context that will be
      reset between invocations.  Create a memory context in
diff --git a/src/backend/executor/nodeForeignscan.c b/src/backend/executor/nodeForeignscan.c
index 513471a..3021252 100644
--- a/src/backend/executor/nodeForeignscan.c
+++ b/src/backend/executor/nodeForeignscan.c
@@ -221,12 +221,18 @@ ExecInitForeignScan(ForeignScan *node, EState *estate, int eflags)
 			ExecInitNode(outerPlan(node), estate, eflags);
 
 	/*
-	 * Tell the FDW to initialize the scan.
+	 * Tell the FDW to initialize the scan or the direct modification.
 	 */
-	if (node->operation != CMD_SELECT)
-		fdwroutine->BeginDirectModify(scanstate, eflags);
-	else
+	if (node->operation == CMD_SELECT)
 		fdwroutine->BeginForeignScan(scanstate, eflags);
+	else
+	{
+		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 89c409d..2afb195 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -761,6 +761,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 e2f1775..15fd85a 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -698,6 +698,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/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 42050ab..4024a80 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -2017,6 +2017,7 @@ _readForeignScan(void)
 	READ_NODE_FIELD(fdw_recheck_quals);
 	READ_BITMAPSET_FIELD(fs_relids);
 	READ_BOOL_FIELD(fsSystemCol);
+	READ_INT_FIELD(resultRelIndex);
 
 	READ_DONE();
 }
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index eb9543f..362bc44 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -5559,6 +5559,8 @@ make_foreignscan(List *qptlist,
 	node->fs_relids = NULL;
 	/* fsSystemCol will be filled in by create_foreignscan_plan */
 	node->fsSystemCol = false;
+	/* resultRelIndex will be set by PlanDirectModify/setrefs.c, if needed */
+	node->resultRelIndex = -1;
 
 	return node;
 }
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index baefe0e..f3d1a12 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -904,6 +904,13 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 					rc->rti += rtoffset;
 					rc->prti += rtoffset;
 				}
+				/*
+				 * Caution: Do not change the relative ordering of this loop
+				 * and the statement below that adds the result relations to
+				 * root->glob->resultRelations, because we need to use the
+				 * current value of list_length(root->glob->resultRelations)
+				 * in some plans.
+				 */
 				foreach(l, splan->plans)
 				{
 					lfirst(l) = set_plan_refs(root,
@@ -1243,6 +1250,14 @@ set_foreignscan_references(PlannerInfo *root,
 	}
 
 	fscan->fs_relids = offset_relid_set(fscan->fs_relids, rtoffset);
+
+	/*
+	 * Adjust resultRelIndex if it's valid (note that we are called before
+	 * adding the RT indexes of ModifyTable result relations to the global
+	 * list)
+	 */
+	if (fscan->resultRelIndex >= 0)
+		fscan->resultRelIndex += list_length(root->glob->resultRelations);
 }
 
 /*
diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h
index 95556df..5c60592 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 83e0107..7314d2f 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -620,6 +620,9 @@ 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;	/* index of foreign table in the list of query
+								 * result relations for INSERT/UPDATE/DELETE;
+								 * -1 for SELECT */
 } ForeignScan;
 
 /* ----------------
-- 
1.8.3.1

