From b5a268bdb661f2fc4ecf327418d6cd7a03c98d69 Mon Sep 17 00:00:00 2001
From: Etsuro Fujita <efujita@postgresql.org>
Date: Thu, 8 Aug 2019 21:41:12 +0900
Subject: [PATCH v13 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 resultRelation
in ForeignScan node which will be set by the core planner for an FDW
to peruse during a "direct modification" operation; it gives the
range table index of the target foreign table.

Amit Langote, Etsuro Fujita
---
 contrib/postgres_fdw/postgres_fdw.c     | 23 +++++++++++++++--------
 doc/src/sgml/fdwhandler.sgml            | 13 +++++++------
 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    |  4 ++++
 src/include/nodes/plannodes.h           |  8 ++++++++
 9 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index a31abce..8eaec6c 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -200,6 +200,8 @@ typedef struct PgFdwDirectModifyState
 	Relation	rel;			/* relcache entry for the foreign table */
 	AttInMetadata *attinmeta;	/* attribute datatype conversion metadata */
 
+	Index			resultRelation;	/* Target foreign table's RT index */
+
 	/* extracted fdw_private data */
 	char	   *query;			/* text of UPDATE/DELETE command */
 	bool		has_returning;	/* is there a RETURNING clause? */
@@ -446,11 +448,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,
@@ -2355,7 +2358,8 @@ 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->resultRelation > 0);
+	dmstate->resultRelation = rtindex = fsplan->resultRelation;
 	rte = exec_rt_fetch(rtindex, estate);
 	userid = rte->checkAsUser ? rte->checkAsUser : GetUserId();
 
@@ -2450,7 +2454,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->resultRelation - 1];
+
+	/* 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 +2489,7 @@ postgresIterateDirectModify(ForeignScanState *node)
 	/*
 	 * Get the next RETURNING tuple.
 	 */
-	return get_returning_data(node);
+	return get_returning_data(node, resultRelInfo);
 }
 
 /*
@@ -4082,11 +4089,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 +4147,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 +4237,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..e995514 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -893,10 +893,10 @@ 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>).
-     <literal>eflags</literal> contains flag bits describing the executor's
-     operating mode for this plan node.
+     <structname>ForeignScan</structname> plan node, which contains the table's range
+     table index 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>
 
     <para>
@@ -926,8 +926,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..515860d 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->resultRelation > 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 7974fa0..70476d8 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -759,6 +759,7 @@ _copyForeignScan(const ForeignScan *from)
 	COPY_NODE_FIELD(fdw_recheck_quals);
 	COPY_BITMAPSET_FIELD(fs_relids);
 	COPY_SCALAR_FIELD(fsSystemCol);
+	COPY_SCALAR_FIELD(resultRelation);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index f35d2a0..2d6efb5 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -696,6 +696,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(resultRelation);
 }
 
 static void
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 73b8d9d..1a9d01f 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -2015,6 +2015,7 @@ _readForeignScan(void)
 	READ_NODE_FIELD(fdw_recheck_quals);
 	READ_BITMAPSET_FIELD(fs_relids);
 	READ_BOOL_FIELD(fsSystemCol);
+	READ_INT_FIELD(resultRelation);
 
 	READ_DONE();
 }
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 881eaf4..fd13ac9 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;
+	/* resultRelation will be set by make_modifytable(), if needed */
+	node->resultRelation = 0;
 
 	return node;
 }
@@ -6898,7 +6900,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->resultRelation = rti;
+
 			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 66372d3..6db6e72 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -1315,6 +1315,10 @@ set_foreignscan_references(PlannerInfo *root,
 	}
 
 	fscan->fs_relids = offset_relid_set(fscan->fs_relids, rtoffset);
+
+	/* Adjust resultRelation if needed */
+	if (fscan->resultRelation > 0)
+		fscan->resultRelation += rtoffset;
 }
 
 /*
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 0666546..cadfdb9 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -605,6 +605,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, resultRelation is set
+ * to identify the relation.
  * ----------------
  */
 typedef struct ForeignScan
@@ -618,6 +623,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 */
+	Index		resultRelation;	/* Target foreign table's RT index; valid for
+								 * result relations of UPDATE/DELETE; 0 for
+								 * other query types */
 } ForeignScan;
 
 /* ----------------
-- 
1.8.3.1

