From 551c8b59925a09ea399c10d364fd52d978e2d456 Mon Sep 17 00:00:00 2001
From: Etsuro Fujita <efujita@postgresql.org>
Date: Thu, 8 Aug 2019 21:41:12 +0900
Subject: [PATCH v12 2/5] Include result relation index if any in ForeignScan

FDWs that can perform an UPDATE/DELETE remotely using the "direct
modify" set of APIs need in some cases to access the result relation
properties for which they can currently look at
EState.es_result_relation_info, which the core executor laboriously
makes sure is set correctly.  An upcoming patch will remove that
field from EState.  So this commit installs a new field resultRelIndex
in ForeignScan node which will be set by the core planner for an FDW
to peruse during a "direct modification" operation.

Amit Langote, Etsuro Fujita
---
 contrib/postgres_fdw/postgres_fdw.c     | 26 ++++++++++++++++++--------
 doc/src/sgml/fdwhandler.sgml            | 10 ++++++----
 src/backend/executor/nodeForeignscan.c  |  5 ++++-
 src/backend/nodes/copyfuncs.c           |  1 +
 src/backend/nodes/outfuncs.c            |  1 +
 src/backend/nodes/readfuncs.c           |  1 +
 src/backend/optimizer/plan/createplan.c | 13 +++++++++++++
 src/backend/optimizer/plan/setrefs.c    | 17 +++++++++++++----
 src/include/nodes/plannodes.h           |  8 ++++++++
 9 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index a31abce..13e256f 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -200,6 +200,9 @@ typedef struct PgFdwDirectModifyState
 	Relation	rel;			/* relcache entry for the foreign table */
 	AttInMetadata *attinmeta;	/* attribute datatype conversion metadata */
 
+	int			resultRelIndex;	/* index of ResultRelInfo for the foreign table
+								 * in EState.es_result_relations */
+
 	/* extracted fdw_private data */
 	char	   *query;			/* text of UPDATE/DELETE command */
 	bool		has_returning;	/* is there a RETURNING clause? */
@@ -446,11 +449,12 @@ static List *build_remote_returning(Index rtindex, Relation rel,
 									List *returningList);
 static void rebuild_fdw_scan_tlist(ForeignScan *fscan, List *tlist);
 static void execute_dml_stmt(ForeignScanState *node);
-static TupleTableSlot *get_returning_data(ForeignScanState *node);
+static TupleTableSlot *get_returning_data(ForeignScanState *node, ResultRelInfo *resultRelInfo);
 static void init_returning_filter(PgFdwDirectModifyState *dmstate,
 								  List *fdw_scan_tlist,
 								  Index rtindex);
 static TupleTableSlot *apply_returning_filter(PgFdwDirectModifyState *dmstate,
+											  ResultRelInfo *relInfo,
 											  TupleTableSlot *slot,
 											  EState *estate);
 static void prepare_query_params(PlanState *node,
@@ -2331,6 +2335,7 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 {
 	ForeignScan *fsplan = (ForeignScan *) node->ss.ps.plan;
 	EState	   *estate = node->ss.ps.state;
+	List	   *resultRelations = estate->es_plannedstmt->resultRelations;
 	PgFdwDirectModifyState *dmstate;
 	Index		rtindex;
 	RangeTblEntry *rte;
@@ -2355,7 +2360,9 @@ 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;
+	Assert(fsplan->resultRelIndex >= 0);
+	dmstate->resultRelIndex = fsplan->resultRelIndex;
+	rtindex = list_nth_int(resultRelations, fsplan->resultRelIndex);
 	rte = exec_rt_fetch(rtindex, estate);
 	userid = rte->checkAsUser ? rte->checkAsUser : GetUserId();
 
@@ -2450,7 +2457,10 @@ postgresIterateDirectModify(ForeignScanState *node)
 {
 	PgFdwDirectModifyState *dmstate = (PgFdwDirectModifyState *) node->fdw_state;
 	EState	   *estate = node->ss.ps.state;
-	ResultRelInfo *resultRelInfo = estate->es_result_relation_info;
+	ResultRelInfo *resultRelInfo = &estate->es_result_relations[dmstate->resultRelIndex];
+
+	/* The executor must have initialized the ResultRelInfo for us. */
+	Assert(resultRelInfo != NULL);
 
 	/*
 	 * If this is the first call after Begin, execute the statement.
@@ -2482,7 +2492,7 @@ postgresIterateDirectModify(ForeignScanState *node)
 	/*
 	 * Get the next RETURNING tuple.
 	 */
-	return get_returning_data(node);
+	return get_returning_data(node, resultRelInfo);
 }
 
 /*
@@ -4082,11 +4092,10 @@ execute_dml_stmt(ForeignScanState *node)
  * Get the result of a RETURNING clause.
  */
 static TupleTableSlot *
-get_returning_data(ForeignScanState *node)
+get_returning_data(ForeignScanState *node, ResultRelInfo *resultRelInfo)
 {
 	PgFdwDirectModifyState *dmstate = (PgFdwDirectModifyState *) node->fdw_state;
 	EState	   *estate = node->ss.ps.state;
-	ResultRelInfo *resultRelInfo = estate->es_result_relation_info;
 	TupleTableSlot *slot = node->ss.ss_ScanTupleSlot;
 	TupleTableSlot *resultSlot;
 
@@ -4141,7 +4150,8 @@ get_returning_data(ForeignScanState *node)
 		if (dmstate->rel)
 			resultSlot = slot;
 		else
-			resultSlot = apply_returning_filter(dmstate, slot, estate);
+			resultSlot = apply_returning_filter(dmstate, resultRelInfo, slot,
+												estate);
 	}
 	dmstate->next_tuple++;
 
@@ -4230,10 +4240,10 @@ init_returning_filter(PgFdwDirectModifyState *dmstate,
  */
 static TupleTableSlot *
 apply_returning_filter(PgFdwDirectModifyState *dmstate,
+					   ResultRelInfo *relInfo,
 					   TupleTableSlot *slot,
 					   EState *estate)
 {
-	ResultRelInfo *relInfo = estate->es_result_relation_info;
 	TupleDesc	resultTupType = RelationGetDescr(dmstate->resultRel);
 	TupleTableSlot *resultSlot;
 	Datum	   *values;
diff --git a/doc/src/sgml/fdwhandler.sgml b/doc/src/sgml/fdwhandler.sgml
index 72fa127..1f6de9b 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -893,8 +893,9 @@ BeginDirectModify(ForeignScanState *node,
      its <structfield>fdw_state</structfield> field is still NULL.  Information about
      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>).
+     <structname>ForeignScan</structname> plan node, which contains an integer field
+     giving the table's index in the query's list of result relations along with any
+     FDW-private information provided by <function>PlanDirectModify</function>.
      <literal>eflags</literal> contains flag bits describing the executor's
      operating mode for this plan node.
     </para>
@@ -926,8 +927,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>
+     obtained using the information 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..19433b3 100644
--- a/src/backend/executor/nodeForeignscan.c
+++ b/src/backend/executor/nodeForeignscan.c
@@ -221,10 +221,13 @@ 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)
+	{
+		Assert(node->resultRelIndex >= 0);
 		fdwroutine->BeginDirectModify(scanstate, eflags);
+	}
 	else
 		fdwroutine->BeginForeignScan(scanstate, eflags);
 
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 0409a40..de57744 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 f038648..fe10e5d 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 3d7a4e3..b157848 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -5541,6 +5541,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 make_modifytable(), if needed */
+	node->resultRelIndex = -1;
 
 	return node;
 }
@@ -6900,7 +6902,18 @@ make_modifytable(PlannerInfo *root,
 			!has_stored_generated_columns(subroot, rti))
 			direct_modify = fdwroutine->PlanDirectModify(subroot, node, rti, i);
 		if (direct_modify)
+		{
+			ForeignScan   *fscan = (ForeignScan *) list_nth(node->plans, i);
+
+			/*
+			 * For result relations that will be modified directly, the FDW
+			 * needs to know where to find them.
+			 */
+			Assert(IsA(fscan, ForeignScan));
+			fscan->resultRelIndex = i;
+
 			direct_modify_plans = bms_add_member(direct_modify_plans, i);
+		}
 
 		if (!direct_modify &&
 			fdwroutine != NULL &&
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index e647f2d..577a911 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -75,6 +75,7 @@ typedef struct
 {
 	PlannerInfo	   *root;
 	int				rtoffset;	/* offset in root->glob->finalrtable */
+	int				rroffset;	/* offset in root->glob->resultRelations */
 } set_plan_refs_context;
 
 /*
@@ -125,7 +126,7 @@ static bool trivial_subqueryscan(SubqueryScan *plan);
 static Plan *clean_up_removed_plan_level(Plan *parent, Plan *child);
 static void set_foreignscan_references(PlannerInfo *root,
 									   ForeignScan *fscan,
-									   int rtoffset);
+									   int rtoffset, int rroffset);
 static void set_customscan_references(CustomScan *cscan,
 									  set_plan_refs_context *context);
 static Plan *set_append_references(Append *aplan,
@@ -254,6 +255,7 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 {
 	PlannerGlobal *glob = root->glob;
 	int			rtoffset = list_length(glob->finalrtable);
+	int			rroffset = list_length(glob->resultRelations);
 	ListCell   *lc;
 	set_plan_refs_context	context;
 
@@ -308,6 +310,7 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 	/* Now fix the Plan tree */
 	context.root = root;
 	context.rtoffset = rtoffset;
+	context.rroffset = rroffset;
 	return set_plan_refs(plan, &context);
 }
 
@@ -506,6 +509,7 @@ set_plan_refs(Plan *plan, set_plan_refs_context *context)
 {
 	PlannerInfo *root = context->root;
 	int			rtoffset = context->rtoffset;
+	int			rroffset = context->rroffset;
 	ListCell   *l;
 
 	if (plan == NULL)
@@ -719,7 +723,8 @@ set_plan_refs(Plan *plan, set_plan_refs_context *context)
 			}
 			break;
 		case T_ForeignScan:
-			set_foreignscan_references(root, (ForeignScan *) plan, rtoffset);
+			set_foreignscan_references(root, (ForeignScan *) plan, rtoffset,
+									   rroffset);
 			break;
 		case T_CustomScan:
 			set_customscan_references((CustomScan *) plan, context);
@@ -985,7 +990,7 @@ set_plan_refs(Plan *plan, set_plan_refs_context *context)
 				 * resultRelIndex to reflect their starting position in the
 				 * global list.
 				 */
-				splan->resultRelIndex = list_length(root->glob->resultRelations);
+				splan->resultRelIndex = rroffset;
 				root->glob->resultRelations =
 					list_concat(root->glob->resultRelations,
 								splan->resultRelations);
@@ -1250,7 +1255,7 @@ clean_up_removed_plan_level(Plan *parent, Plan *child)
 static void
 set_foreignscan_references(PlannerInfo *root,
 						   ForeignScan *fscan,
-						   int rtoffset)
+						   int rtoffset, int rroffset)
 {
 	/* Adjust scanrelid if it's valid */
 	if (fscan->scan.scanrelid > 0)
@@ -1319,6 +1324,10 @@ set_foreignscan_references(PlannerInfo *root,
 	}
 
 	fscan->fs_relids = offset_relid_set(fscan->fs_relids, rtoffset);
+
+	/* Adjust resultRelIndex if needed */
+	if (fscan->resultRelIndex >= 0)
+		fscan->resultRelIndex += rroffset;
 }
 
 /*
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 83e0107..b0a3924 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -607,6 +607,11 @@ typedef struct WorkTableScan
  * When the plan node represents a foreign join, scan.scanrelid is zero and
  * fs_relids must be consulted to identify the join relation.  (fs_relids
  * is valid for simple scans as well, but will always match scan.scanrelid.)
+ *
+ * If an FDW's PlanDirectModify() callback decides to repurpose a ForeignScan
+ * node to store the information about an UPDATE or DELETE operation to
+ * to perform on a given foreign table result relation, resultRelIndex is set
+ * to identify the result relation.
  * ----------------
  */
 typedef struct ForeignScan
@@ -620,6 +625,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 UPDATE/DELETE; -1 for
+								 * other query types */
 } ForeignScan;
 
 /* ----------------
-- 
1.8.3.1

