fix-tuple-routing-for-foreign-partitions-2.patch

*** a/contrib/postgres_fdw/expected/postgres_fdw.out
--- b/contrib/postgres_fdw/expected/postgres_fdw.out
***************
*** 7454,7459 **** select tableoid::regclass, * FROM itrtest;
--- 7454,7501 ----
   remp1    | 1 | foo
  (1 row)
  
+ delete from itrtest;
+ drop index loct1_idx;
+ -- Test that remote triggers work with insert tuple routing
+ create function br_insert_trigfunc() returns trigger as $$
+ begin
+ 	new.b := new.b || ' triggered !';
+ 	return new;
+ end
+ $$ language plpgsql;
+ create trigger loct1_br_insert_trigger before insert on loct1
+ 	for each row execute procedure br_insert_trigfunc();
+ create trigger loct2_br_insert_trigger before insert on loct2
+ 	for each row execute procedure br_insert_trigfunc();
+ -- The new values are concatenated with ' triggered !'
+ insert into itrtest values (1, 'foo') returning *;
+  a |        b        
+ ---+-----------------
+  1 | foo triggered !
+ (1 row)
+ 
+ insert into itrtest values (2, 'qux') returning *;
+  a |        b        
+ ---+-----------------
+  2 | qux triggered !
+ (1 row)
+ 
+ insert into itrtest values (1, 'test1'), (2, 'test2') returning *;
+  a |         b         
+ ---+-------------------
+  1 | test1 triggered !
+  2 | test2 triggered !
+ (2 rows)
+ 
+ with result as (insert into itrtest values (1, 'test1'), (2, 'test2') returning *) select * from result;
+  a |         b         
+ ---+-------------------
+  1 | test1 triggered !
+  2 | test2 triggered !
+ (2 rows)
+ 
+ drop trigger loct1_br_insert_trigger on loct1;
+ drop trigger loct2_br_insert_trigger on loct2;
  drop table itrtest;
  drop table loct1;
  drop table loct2;
***************
*** 7518,7523 **** select tableoid::regclass, * FROM locp;
--- 7560,7615 ----
  
  -- The executor should not let unexercised FDWs shut down
  update utrtest set a = 1 where b = 'foo';
+ -- Test that remote triggers work with update tuple routing
+ create trigger loct_br_insert_trigger before insert on loct
+ 	for each row execute procedure br_insert_trigfunc();
+ delete from utrtest;
+ insert into utrtest values (2, 'qux');
+ -- Check case where the foreign partition is a subplan target rel
+ explain (verbose, costs off)
+ update utrtest set a = 1 where a = 1 or a = 2 returning *;
+                                           QUERY PLAN                                          
+ ----------------------------------------------------------------------------------------------
+  Update on public.utrtest
+    Output: remp.a, remp.b
+    Foreign Update on public.remp
+    Update on public.locp
+    ->  Foreign Update on public.remp
+          Remote SQL: UPDATE public.loct SET a = 1 WHERE (((a = 1) OR (a = 2))) RETURNING a, b
+    ->  Seq Scan on public.locp
+          Output: 1, locp.b, locp.ctid
+          Filter: ((locp.a = 1) OR (locp.a = 2))
+ (9 rows)
+ 
+ -- The new values are concatenated with ' triggered !'
+ update utrtest set a = 1 where a = 1 or a = 2 returning *;
+  a |        b        
+ ---+-----------------
+  1 | qux triggered !
+ (1 row)
+ 
+ delete from utrtest;
+ insert into utrtest values (2, 'qux');
+ -- Check case where the foreign partition isn't a subplan target rel
+ explain (verbose, costs off)
+ update utrtest set a = 1 where a = 2 returning *;
+               QUERY PLAN              
+ --------------------------------------
+  Update on public.utrtest
+    Output: locp.a, locp.b
+    Update on public.locp
+    ->  Seq Scan on public.locp
+          Output: 1, locp.b, locp.ctid
+          Filter: (locp.a = 2)
+ (6 rows)
+ 
+ -- The new values are concatenated with ' triggered !'
+ update utrtest set a = 1 where a = 2 returning *;
+  a |        b        
+ ---+-----------------
+  1 | qux triggered !
+ (1 row)
+ 
  drop table utrtest;
  drop table loct;
  -- Test copy tuple routing
*** a/contrib/postgres_fdw/postgres_fdw.c
--- b/contrib/postgres_fdw/postgres_fdw.c
***************
*** 1974,1982 **** postgresBeginForeignInsert(ModifyTableState *mtstate,
  						   ResultRelInfo *resultRelInfo)
  {
  	PgFdwModifyState *fmstate;
! 	Plan	   *plan = mtstate->ps.plan;
  	Relation	rel = resultRelInfo->ri_RelationDesc;
- 	RangeTblEntry *rte;
  	Query	   *query;
  	PlannerInfo *root;
  	TupleDesc	tupdesc = RelationGetDescr(rel);
--- 1974,1983 ----
  						   ResultRelInfo *resultRelInfo)
  {
  	PgFdwModifyState *fmstate;
! 	ModifyTable *plan = (ModifyTable *) mtstate->ps.plan;
! 	EState	   *estate = mtstate->ps.state;
! 	Index		resultRelation = resultRelInfo->ri_RangeTableIndex;
  	Relation	rel = resultRelInfo->ri_RelationDesc;
  	Query	   *query;
  	PlannerInfo *root;
  	TupleDesc	tupdesc = RelationGetDescr(rel);
***************
*** 1985,2002 **** postgresBeginForeignInsert(ModifyTableState *mtstate,
  	List	   *targetAttrs = NIL;
  	List	   *retrieved_attrs = NIL;
  	bool		doNothing = false;
  
  	initStringInfo(&sql);
  
  	/* Set up largely-dummy planner state. */
- 	rte = makeNode(RangeTblEntry);
- 	rte->rtekind = RTE_RELATION;
- 	rte->relid = RelationGetRelid(rel);
- 	rte->relkind = RELKIND_FOREIGN_TABLE;
  	query = makeNode(Query);
  	query->commandType = CMD_INSERT;
! 	query->resultRelation = 1;
! 	query->rtable = list_make1(rte);
  	root = makeNode(PlannerInfo);
  	root->parse = query;
  
--- 1986,2047 ----
  	List	   *targetAttrs = NIL;
  	List	   *retrieved_attrs = NIL;
  	bool		doNothing = false;
+ 	bool		replace_rte = false;
+ 	ListCell   *lc = NULL;
+ 	RangeTblEntry *old_rte = NULL;
  
  	initStringInfo(&sql);
  
+ 	/*
+ 	 * If the foreign table is a partition, temporarily replace the parent's
+ 	 * RTE in the range table with a new target RTE describing the foreign
+ 	 * table for use by deparseInsertSql() and create_foreign_modify() below.
+ 	 */
+ 	if (resultRelInfo->ri_PartitionRoot)
+ 	{
+ 		replace_rte = true;
+ 
+ 		/*
+ 		 * If this is invoked by an UPDATE, we need special handling: if the
+ 		 * partition is one of the UPDATE subplan target rels, use the target
+ 		 * rel's RTE without replacement; else replace the RTE for the first
+ 		 * subplan target rel with the new target RTE as we have created
+ 		 * expressions in the partition's ResultRelInfo such as RETURNING by
+ 		 * translating the attnos of the corresponding expressions in the
+ 		 * ResultRelInfo for the first subplan target rel before we get here.
+ 		 */
+ 		if (plan && plan->operation == CMD_UPDATE)
+ 		{
+ 			if (resultRelation != plan->nominalRelation)
+ 			{
+ 				/* partition that is a subplan target rel */
+ 				replace_rte = false;
+ 			}
+ 			else
+ 			{
+ 				/* partition that isn't a subplan target rel */ 
+ 				resultRelation = mtstate->resultRelInfo[0].ri_RangeTableIndex;
+ 			}
+ 		}
+ 
+ 		if (replace_rte)
+ 		{
+ 			RangeTblEntry *new_rte;
+ 
+ 			lc = list_nth_cell(estate->es_range_table, resultRelation - 1);
+ 			old_rte = lfirst(lc);
+ 			new_rte = copyObject(old_rte);
+ 			new_rte->relid = RelationGetRelid(rel);
+ 			new_rte->relkind = RELKIND_FOREIGN_TABLE;
+ 			lfirst(lc) = new_rte;
+ 		}
+ 	}
+ 
  	/* Set up largely-dummy planner state. */
  	query = makeNode(Query);
  	query->commandType = CMD_INSERT;
! 	query->resultRelation = resultRelation;
! 	query->rtable = estate->es_range_table;
  	root = makeNode(PlannerInfo);
  	root->parse = query;
  
***************
*** 2012,2018 **** postgresBeginForeignInsert(ModifyTableState *mtstate,
  	/* Check if we add the ON CONFLICT clause to the remote query. */
  	if (plan)
  	{
! 		OnConflictAction onConflictAction = ((ModifyTable *) plan)->onConflictAction;
  
  		/* We only support DO NOTHING without an inference specification. */
  		if (onConflictAction == ONCONFLICT_NOTHING)
--- 2057,2063 ----
  	/* Check if we add the ON CONFLICT clause to the remote query. */
  	if (plan)
  	{
! 		OnConflictAction onConflictAction = plan->onConflictAction;
  
  		/* We only support DO NOTHING without an inference specification. */
  		if (onConflictAction == ONCONFLICT_NOTHING)
***************
*** 2023,2033 **** postgresBeginForeignInsert(ModifyTableState *mtstate,
  	}
  
  	/* Construct the SQL command string. */
! 	deparseInsertSql(&sql, root, 1, rel, targetAttrs, doNothing,
  					 resultRelInfo->ri_returningList, &retrieved_attrs);
  
  	/* Construct an execution state. */
! 	fmstate = create_foreign_modify(mtstate->ps.state,
  									resultRelInfo,
  									CMD_INSERT,
  									NULL,
--- 2068,2078 ----
  	}
  
  	/* Construct the SQL command string. */
! 	deparseInsertSql(&sql, root, resultRelation, rel, targetAttrs, doNothing,
  					 resultRelInfo->ri_returningList, &retrieved_attrs);
  
  	/* Construct an execution state. */
! 	fmstate = create_foreign_modify(estate,
  									resultRelInfo,
  									CMD_INSERT,
  									NULL,
***************
*** 2036,2041 **** postgresBeginForeignInsert(ModifyTableState *mtstate,
--- 2081,2090 ----
  									retrieved_attrs != NIL,
  									retrieved_attrs);
  
+ 	/* Restore the saved RTE. */
+ 	if (replace_rte)
+ 		lfirst(lc) = old_rte;
+ 
  	resultRelInfo->ri_FdwState = fmstate;
  }
  
*** a/contrib/postgres_fdw/sql/postgres_fdw.sql
--- b/contrib/postgres_fdw/sql/postgres_fdw.sql
***************
*** 1804,1809 **** insert into itrtest values (1, 'bar') on conflict (a) do update set b = excluded
--- 1804,1834 ----
  
  select tableoid::regclass, * FROM itrtest;
  
+ delete from itrtest;
+ 
+ drop index loct1_idx;
+ 
+ -- Test that remote triggers work with insert tuple routing
+ create function br_insert_trigfunc() returns trigger as $$
+ begin
+ 	new.b := new.b || ' triggered !';
+ 	return new;
+ end
+ $$ language plpgsql;
+ create trigger loct1_br_insert_trigger before insert on loct1
+ 	for each row execute procedure br_insert_trigfunc();
+ create trigger loct2_br_insert_trigger before insert on loct2
+ 	for each row execute procedure br_insert_trigfunc();
+ 
+ -- The new values are concatenated with ' triggered !'
+ insert into itrtest values (1, 'foo') returning *;
+ insert into itrtest values (2, 'qux') returning *;
+ insert into itrtest values (1, 'test1'), (2, 'test2') returning *;
+ with result as (insert into itrtest values (1, 'test1'), (2, 'test2') returning *) select * from result;
+ 
+ drop trigger loct1_br_insert_trigger on loct1;
+ drop trigger loct2_br_insert_trigger on loct2;
+ 
  drop table itrtest;
  drop table loct1;
  drop table loct2;
***************
*** 1836,1841 **** select tableoid::regclass, * FROM locp;
--- 1861,1888 ----
  -- The executor should not let unexercised FDWs shut down
  update utrtest set a = 1 where b = 'foo';
  
+ -- Test that remote triggers work with update tuple routing
+ create trigger loct_br_insert_trigger before insert on loct
+ 	for each row execute procedure br_insert_trigfunc();
+ 
+ delete from utrtest;
+ insert into utrtest values (2, 'qux');
+ 
+ -- Check case where the foreign partition is a subplan target rel
+ explain (verbose, costs off)
+ update utrtest set a = 1 where a = 1 or a = 2 returning *;
+ -- The new values are concatenated with ' triggered !'
+ update utrtest set a = 1 where a = 1 or a = 2 returning *;
+ 
+ delete from utrtest;
+ insert into utrtest values (2, 'qux');
+ 
+ -- Check case where the foreign partition isn't a subplan target rel
+ explain (verbose, costs off)
+ update utrtest set a = 1 where a = 2 returning *;
+ -- The new values are concatenated with ' triggered !'
+ update utrtest set a = 1 where a = 2 returning *;
+ 
  drop table utrtest;
  drop table loct;
  
*** a/src/backend/executor/execPartition.c
--- b/src/backend/executor/execPartition.c
***************
*** 330,335 **** ExecInitPartitionInfo(ModifyTableState *mtstate,
--- 330,342 ----
  					  estate->es_instrument);
  
  	/*
+ 	 * Verify result relation is a valid target for an INSERT.  An UPDATE of a
+ 	 * partition-key becomes a DELETE+INSERT operation, so this check is still
+ 	 * required when the operation is CMD_UPDATE.
+ 	 */
+ 	CheckValidResultRel(leaf_part_rri, CMD_INSERT);
+ 
+ 	/*
  	 * Since we've just initialized this ResultRelInfo, it's not in any list
  	 * attached to the estate as yet.  Add it, so that it can be found later.
  	 *
***************
*** 341,349 **** ExecInitPartitionInfo(ModifyTableState *mtstate,
  		lappend(estate->es_tuple_routing_result_relations,
  				leaf_part_rri);
  
- 	/* Set up information needed for routing tuples to this partition. */
- 	ExecInitRoutingInfo(mtstate, estate, proute, leaf_part_rri, partidx);
- 
  	/*
  	 * Open partition indices.  The user may have asked to check for conflicts
  	 * within this leaf partition and do "nothing" instead of throwing an
--- 348,353 ----
***************
*** 466,471 **** ExecInitPartitionInfo(ModifyTableState *mtstate,
--- 470,478 ----
  									&mtstate->ps, RelationGetDescr(partrel));
  	}
  
+ 	/* Set up information needed for routing tuples to the partition. */
+ 	ExecInitRoutingInfo(mtstate, estate, proute, leaf_part_rri, partidx);
+ 
  	/*
  	 * If there is an ON CONFLICT clause, initialize state for it.
  	 */
***************
*** 612,619 **** ExecInitPartitionInfo(ModifyTableState *mtstate,
  
  /*
   * ExecInitRoutingInfo
!  *		Set up information needed for routing tuples to a leaf partition if
!  *		routable; else abort the operation
   */
  void
  ExecInitRoutingInfo(ModifyTableState *mtstate,
--- 619,625 ----
  
  /*
   * ExecInitRoutingInfo
!  *		Set up information needed for routing tuples to a leaf partition
   */
  void
  ExecInitRoutingInfo(ModifyTableState *mtstate,
***************
*** 624,632 **** ExecInitRoutingInfo(ModifyTableState *mtstate,
  {
  	MemoryContext oldContext;
  
- 	/* Verify the partition is a valid target for INSERT */
- 	CheckValidResultRel(partRelInfo, CMD_INSERT);
- 
  	/*
  	 * Switch into per-query memory context.
  	 */
--- 630,635 ----
*** a/src/backend/executor/nodeModifyTable.c
--- b/src/backend/executor/nodeModifyTable.c
***************
*** 1700,1719 **** ExecPrepareTupleRouting(ModifyTableState *mtstate,
  										partidx);
  
  	/*
! 	 * Set up information needed for routing tuples to the partition if we
! 	 * didn't yet (ExecInitRoutingInfo would abort the operation if the
! 	 * partition isn't routable).
  	 *
  	 * Note: an UPDATE of a partition key invokes an INSERT that moves the
! 	 * tuple to a new partition.  This setup would be needed for a subplan
  	 * partition of such an UPDATE that is chosen as the partition to route
! 	 * the tuple to.  The reason we do this setup here rather than in
  	 * ExecSetupPartitionTupleRouting is to avoid aborting such an UPDATE
  	 * unnecessarily due to non-routable subplan partitions that may not be
  	 * chosen for update tuple movement after all.
  	 */
  	if (!partrel->ri_PartitionReadyForRouting)
  		ExecInitRoutingInfo(mtstate, estate, proute, partrel, partidx);
  
  	/*
  	 * Make it look like we are inserting into the partition.
--- 1700,1723 ----
  										partidx);
  
  	/*
! 	 * Check whether the partition is routable if we didn't yet
  	 *
  	 * Note: an UPDATE of a partition key invokes an INSERT that moves the
! 	 * tuple to a new partition.  This check would be applied to a subplan
  	 * partition of such an UPDATE that is chosen as the partition to route
! 	 * the tuple to.  The reason we do this check here rather than in
  	 * ExecSetupPartitionTupleRouting is to avoid aborting such an UPDATE
  	 * unnecessarily due to non-routable subplan partitions that may not be
  	 * chosen for update tuple movement after all.
  	 */
  	if (!partrel->ri_PartitionReadyForRouting)
+ 	{
+ 		/* Verify the partition is a valid target for INSERT. */
+ 		CheckValidResultRel(partrel, CMD_INSERT);
+ 
+ 		/* Set up information needed for routing tuples to the partition. */
  		ExecInitRoutingInfo(mtstate, estate, proute, partrel, partidx);
+ 	}
  
  	/*
  	 * Make it look like we are inserting into the partition.