v4-0002-Include-result-relation-index-if-any-in-ForeignSc.patch

From 25fcde25709cac15fefba950115f99b825f8d6b3 Mon Sep 17 00:00:00 2001
From: Etsuro Fujita <efujita@postgresql.org>
Date: Thu, 8 Aug 2019 21:41:12 +0900
Subject: [PATCH v4 2/3] 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.  However that means the executor must
ensure that es_result_relation_info points to the correct result
relation at all times, especially during inherited updates.  This
requirement gets in the way of number of projects related to changing
how ModifyTable operates.  For example, an upcoming patch will change
things such that there will be one source plan for all result
relations whereas currently there is one per result relation, an
arrangement which makes it convenient to switch the result relation
when the source plan changes.

This commit installs a new field 'resultRelIndex' in ForeignScan node
which must be set by an FDW if the node will be used to carry out an
UPDATE/DELETE operation on a given foreign table, which is the case
if the FDW manages to push that operations to the remote side.  This
commit also modifies postgres_fdw to implement that.

Amit Langote, Etsuro Fujita
---
 contrib/postgres_fdw/postgres_fdw.c     | 31 +++++++++++++++++++++++--------
 doc/src/sgml/fdwhandler.sgml            | 14 +++++++++-----
 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 |  2 ++
 src/backend/optimizer/plan/setrefs.c    | 15 +++++++++++++++
 src/include/nodes/plannodes.h           |  8 ++++++++
 9 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 5583fff..963c2ac 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,
@@ -2300,6 +2304,11 @@ postgresPlanDirectModify(PlannerInfo *root,
 	fscan->operation = operation;
 
 	/*
+	 * Store the index of the result relation we will be operating on.
+	 */
+	fscan->resultRelIndex = subplan_index;
+
+	/*
 	 * Update the fdw_exprs list that will be available to the executor.
 	 */
 	fscan->fdw_exprs = params_list;
@@ -2339,6 +2348,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;
@@ -2363,7 +2373,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();
 
@@ -2458,7 +2470,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.
@@ -2490,7 +2505,7 @@ postgresIterateDirectModify(ForeignScanState *node)
 	/*
 	 * Get the next RETURNING tuple.
 	 */
-	return get_returning_data(node);
+	return get_returning_data(node, resultRelInfo);
 }
 
 /*
@@ -4090,11 +4105,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;
 
@@ -4149,7 +4163,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++;
 
@@ -4238,10 +4253,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 16a08e7..7c6c48a 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -868,7 +868,9 @@ PlanDirectModify(PlannerInfo *root,
      be set to the <literal>CmdType</literal> enumeration appropriately; that is,
      <literal>CMD_UPDATE</literal> for <command>UPDATE</command>,
      <literal>CMD_INSERT</literal> for <command>INSERT</command>, and
-     <literal>CMD_DELETE</literal> for <command>DELETE</command>.
+     <literal>CMD_DELETE</literal> for <command>DELETE</command>.  Also,
+     its <structfield>resultRelIndex</structfield> field must be set to
+     <literal>subplan_index</literal>.
     </para>
 
     <para>
@@ -896,8 +898,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>
@@ -929,8 +932,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 fa046f6..4a32d5f 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -762,6 +762,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 c09bb41..5f372ac 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -699,6 +699,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 92bb7ad..872cf10 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -2018,6 +2018,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 11c7eb5..352b5cf 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -5576,6 +5576,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;
 }
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index a10fd2d..5fa41d4 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -972,6 +972,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,
@@ -1327,6 +1334,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/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 360034d..fba746e 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -608,6 +608,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, it must also set the
+ * resultRelIndex field to identify the result relation.
  * ----------------
  */
 typedef struct ForeignScan
@@ -621,6 +626,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