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

---
 contrib/postgres_fdw/postgres_fdw.c    | 22 ++++++++++++++++------
 doc/src/sgml/fdwhandler.sgml           | 11 ++++++++++-
 src/backend/executor/nodeForeignscan.c |  7 +++----
 src/backend/executor/nodeModifyTable.c | 15 ++++++++++++---
 src/include/foreign/fdwapi.h           |  1 +
 5 files changed, 42 insertions(+), 14 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..88d9f5da10 100644
--- a/src/backend/executor/nodeForeignscan.c
+++ b/src/backend/executor/nodeForeignscan.c
@@ -221,11 +221,10 @@ 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.  Direct modification scans
+	 * are initialized elsewhere; see ExecInitModifyTable().
 	 */
-	if (node->operation != CMD_SELECT)
-		fdwroutine->BeginDirectModify(scanstate, eflags);
-	else
+	if (node->operation == CMD_SELECT)
 		fdwroutine->BeginForeignScan(scanstate, eflags);
 
 	return scanstate;
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 9e0c8794c4..55ce709adb 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -2370,9 +2370,18 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 								   table_slot_callbacks(resultRelInfo->ri_RelationDesc));
 
 		/* Also let FDWs init themselves for foreign-table result rels */
-		if (!resultRelInfo->ri_usesFdwDirectModify &&
-			resultRelInfo->ri_FdwRoutine != NULL &&
-			resultRelInfo->ri_FdwRoutine->BeginForeignModify != NULL)
+		if (resultRelInfo->ri_usesFdwDirectModify)
+		{
+			ForeignScanState *scanstate;
+
+			Assert(IsA(mtstate->mt_plans[i], ForeignScanState));
+			scanstate = (ForeignScanState *) mtstate->mt_plans[i];
+			resultRelInfo->ri_FdwRoutine->BeginDirectModify(scanstate,
+															resultRelInfo,
+															eflags);
+		}
+		else if (resultRelInfo->ri_FdwRoutine != NULL &&
+				 resultRelInfo->ri_FdwRoutine->BeginForeignModify != NULL)
 		{
 			List	   *fdw_private = (List *) list_nth(node->fdwPrivLists, i);
 
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);
-- 
2.11.0

