v4-0001-Fix-a-bug-with-RETURNING-when-UPDATE-moves-tuple.patch

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 20a4c474cc..6476ee8cf1 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -153,9 +153,6 @@ ExecCheckPlanOutput(Relation resultRel, List *targetList)
  * tupleSlot: slot holding tuple actually inserted/updated/deleted
  * planSlot: slot holding tuple returned by top subplan node
  *
- * Note: If tupleSlot is NULL, the FDW should have already provided econtext's
- * scan tuple.
- *
  * Returns a slot holding the result tuple
  */
 static TupleTableSlot *
@@ -166,17 +163,28 @@ ExecProcessReturning(ResultRelInfo *resultRelInfo,
 	ProjectionInfo *projectReturning = resultRelInfo->ri_projectReturning;
 	ExprContext *econtext = projectReturning->pi_exprContext;
 
-	/* Make tuple and any needed join variables available to ExecProject */
+	/*
+	 * Make tuple and any needed join variables available to ExecProject
+	 *
+	 * Note: If tupleSlot is NULL, it means this is called for modifying a
+	 * foreign table directly, in which case the FDW should have already
+	 * provided econtext's ecxt_scantuple.
+	 */
 	if (tupleSlot)
 		econtext->ecxt_scantuple = tupleSlot;
-	econtext->ecxt_outertuple = planSlot;
+	else
+	{
+		Assert(resultRelInfo->ri_usesFdwDirectModify);
+		Assert(econtext->ecxt_scantuple);
 
-	/*
-	 * RETURNING expressions might reference the tableoid column, so
-	 * reinitialize tts_tableOid before evaluating them.
-	 */
-	econtext->ecxt_scantuple->tts_tableOid =
-		RelationGetRelid(resultRelInfo->ri_RelationDesc);
+		/*
+		 * The FDW may already have initialized tts_tableOid, but in case it
+		 * didn't, do so now.
+		 */
+		econtext->ecxt_scantuple->tts_tableOid =
+			RelationGetRelid(resultRelInfo->ri_RelationDesc);
+	}
+	econtext->ecxt_outertuple = planSlot;
 
 	/* Compute the RETURNING expressions */
 	return ExecProject(projectReturning);
@@ -368,7 +376,7 @@ ExecComputeStoredGenerated(EState *estate, TupleTableSlot *slot, CmdType cmdtype
  *		For INSERT, we have to insert the tuple into the target relation
  *		and insert appropriate tuples into the index relations.
  *
- *		Returns RETURNING result if any, otherwise NULL.
+ *		Returns RETURNING result if any and requested, otherwise NULL.
  * ----------------------------------------------------------------
  */
 static TupleTableSlot *
@@ -376,7 +384,9 @@ ExecInsert(ModifyTableState *mtstate,
 		   TupleTableSlot *slot,
 		   TupleTableSlot *planSlot,
 		   EState *estate,
-		   bool canSetTag)
+		   bool canSetTag,
+		   bool processReturning,
+		   TupleTableSlot **resultSlot)
 {
 	ResultRelInfo *resultRelInfo;
 	Relation	resultRelationDesc;
@@ -677,9 +687,13 @@ ExecInsert(ModifyTableState *mtstate,
 		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate);
 
 	/* Process RETURNING if present */
-	if (resultRelInfo->ri_projectReturning)
+	if (processReturning && resultRelInfo->ri_projectReturning)
 		result = ExecProcessReturning(resultRelInfo, slot, planSlot);
 
+	/* If requested, pass back the inserted row */
+	if (resultSlot)
+		*resultSlot = slot;
+
 	return result;
 }
 
@@ -1202,11 +1216,14 @@ lreplace:;
 		if (partition_constraint_failed)
 		{
 			bool		tuple_deleted;
-			TupleTableSlot *ret_slot;
+			TupleTableSlot *orig_slot = slot;
+			TupleTableSlot *res_slot = NULL;
+			TupleTableSlot *ret_slot = NULL;
 			TupleTableSlot *epqslot = NULL;
 			PartitionTupleRouting *proute = mtstate->mt_partition_tuple_routing;
 			int			map_index;
 			TupleConversionMap *tupconv_map;
+			ResultRelInfo *destRel;
 
 			/*
 			 * Disallow an INSERT ON CONFLICT DO UPDATE that causes the
@@ -1307,9 +1324,18 @@ lreplace:;
 			Assert(mtstate->rootResultRelInfo != NULL);
 			slot = ExecPrepareTupleRouting(mtstate, estate, proute,
 										   mtstate->rootResultRelInfo, slot);
+			destRel = estate->es_result_relation_info;
 
-			ret_slot = ExecInsert(mtstate, slot, planSlot,
-								  estate, canSetTag);
+			/*
+			 * Tell ExecInsert() to not process RETURNING, because using the
+			 * the destination partition's RETURNING projection may not work
+			 * correctly if some columns in the RETURNING list reference
+			 * planSlot, which is based on the source partition's tuple
+			 * descriptor.
+			 */
+			ret_slot = ExecInsert(mtstate, slot, planSlot, estate, canSetTag,
+								  false, &res_slot);
+			Assert(ret_slot == NULL);
 
 			/* Revert ExecPrepareTupleRouting's node change. */
 			estate->es_result_relation_info = resultRelInfo;
@@ -1319,6 +1345,38 @@ lreplace:;
 				mtstate->mt_transition_capture->tcs_map = saved_tcs_map;
 			}
 
+			/* Process RETURNING using the source relation's projection. */
+			if (resultRelInfo->ri_projectReturning && res_slot)
+			{
+				/*
+				 * Switch back to the original slot.  While switching the slot,
+				 * also check if we need to convert the tuple, because the
+				 * tuples descriptors of the source and the destination
+				 * partitions may not match.  While it's not great that we have
+				 * to check and build the map from scratch for every tuple that
+				 * is moved between partitions whose tuple descriptors don't
+				 * match, such cases should be rare in practice.
+				 */
+				if (res_slot != orig_slot)
+				{
+					AttrMap *map;
+
+					map = build_attrmap_by_name_if_req(RelationGetDescr(destRel->ri_RelationDesc),
+													   RelationGetDescr(resultRelInfo->ri_RelationDesc));
+					if (map)
+						orig_slot = execute_attr_map_slot(map,
+														  res_slot,
+														  orig_slot);
+
+					/*
+					 * Override tts_tableOid with the OID of the destination
+					 * partition.
+					 */
+					orig_slot->tts_tableOid = RelationGetRelid(destRel->ri_RelationDesc);
+				}
+				return ExecProcessReturning(resultRelInfo, orig_slot, planSlot);
+			}
+
 			return ret_slot;
 		}
 
@@ -2244,7 +2302,8 @@ ExecModifyTable(PlanState *pstate)
 					slot = ExecPrepareTupleRouting(node, estate, proute,
 												   resultRelInfo, slot);
 				slot = ExecInsert(node, slot, planSlot,
-								  estate, node->canSetTag);
+								  estate, node->canSetTag,
+								  true, NULL);
 				/* Revert ExecPrepareTupleRouting's state change. */
 				if (proute)
 					estate->es_result_relation_info = resultRelInfo;
diff --git a/src/test/regress/expected/update.out b/src/test/regress/expected/update.out
index bf939d79f6..807005c40a 100644
--- a/src/test/regress/expected/update.out
+++ b/src/test/regress/expected/update.out
@@ -445,6 +445,46 @@ UPDATE range_parted set c = 95 WHERE a = 'b' and b > 10 and c > 100 returning (r
  part_c_1_100   | b | 17 |  95 | 19 | 
 (6 rows)
 
+-- The following tests computing RETURNING when the source and the destination
+-- partitions of a UPDATE row movement operation have different tuple
+-- descriptors, which has been shown to be problematic in the cases where the
+-- RETURNING targetlist contains non-target relation attributes that are
+-- computed by referring to the source partition plan's output tuple.  Also,
+-- a trigger on the destination relation may change the tuple, which must be
+-- reflected in the RETURNING output, so we test that too.
+CREATE TABLE part_c_1_c_20 (LIKE range_parted);
+ALTER TABLE part_c_1_c_20 DROP a, DROP b, ADD a text, ADD b bigint;
+ALTER TABLE range_parted ATTACH PARTITION part_c_1_c_20 FOR VALUES FROM ('c', 1) TO ('c', 20);
+CREATE FUNCTION trigfunc () RETURNS TRIGGER LANGUAGE plpgsql as $$ BEGIN NEW.e := 'in trigfunc()'; RETURN NEW; END; $$;
+CREATE TRIGGER trig BEFORE INSERT ON part_c_1_c_20 FOR EACH ROW EXECUTE FUNCTION trigfunc();
+UPDATE range_parted r set a = 'c' FROM (VALUES ('a', 1), ('a', 10), ('b', 12)) s(x, y) WHERE s.x = r.a AND s.y = r.b RETURNING tableoid::regclass, *;
+   tableoid    | a | b  |  c  | d |       e       | x | y  
+---------------+---+----+-----+---+---------------+---+----
+ part_c_1_c_20 | c |  1 |   1 | 1 | in trigfunc() | a |  1
+ part_c_1_c_20 | c | 10 | 200 | 1 | in trigfunc() | a | 10
+ part_c_1_c_20 | c | 12 |  96 | 1 | in trigfunc() | b | 12
+(3 rows)
+
+DROP TRIGGER trig ON part_c_1_c_20;
+DROP FUNCTION trigfunc;
+:init_range_parted;
+CREATE FUNCTION trigfunc () RETURNS TRIGGER LANGUAGE plpgsql as $$ BEGIN RETURN NULL; END; $$;
+CREATE TRIGGER trig BEFORE INSERT ON part_c_1_c_20 FOR EACH ROW EXECUTE FUNCTION trigfunc();
+UPDATE range_parted r set a = 'c' FROM (VALUES ('a', 1), ('a', 10), ('b', 12)) s(x, y) WHERE s.x = r.a AND s.y = r.b RETURNING tableoid::regclass, *;
+ tableoid | a | b | c | d | e | x | y 
+----------+---+---+---+---+---+---+---
+(0 rows)
+
+:show_data;
+   partname   | a | b  |  c  | d  | e 
+--------------+---+----+-----+----+---
+ part_c_1_100 | b | 13 |  97 |  2 | 
+ part_d_15_20 | b | 15 | 105 | 16 | 
+ part_d_15_20 | b | 17 | 105 | 19 | 
+(3 rows)
+
+DROP TABLE part_c_1_c_20;
+DROP FUNCTION trigfunc;
 -- Transition tables with update row movement
 :init_range_parted;
 CREATE FUNCTION trans_updatetrigfunc() RETURNS trigger LANGUAGE plpgsql AS
diff --git a/src/test/regress/sql/update.sql b/src/test/regress/sql/update.sql
index 8c558a7bc7..dc7274a4bb 100644
--- a/src/test/regress/sql/update.sql
+++ b/src/test/regress/sql/update.sql
@@ -236,6 +236,31 @@ DROP VIEW upview;
 UPDATE range_parted set c = 95 WHERE a = 'b' and b > 10 and c > 100 returning (range_parted), *;
 :show_data;
 
+-- The following tests computing RETURNING when the source and the destination
+-- partitions of a UPDATE row movement operation have different tuple
+-- descriptors, which has been shown to be problematic in the cases where the
+-- RETURNING targetlist contains non-target relation attributes that are
+-- computed by referring to the source partition plan's output tuple.  Also,
+-- a trigger on the destination relation may change the tuple, which must be
+-- reflected in the RETURNING output, so we test that too.
+CREATE TABLE part_c_1_c_20 (LIKE range_parted);
+ALTER TABLE part_c_1_c_20 DROP a, DROP b, ADD a text, ADD b bigint;
+ALTER TABLE range_parted ATTACH PARTITION part_c_1_c_20 FOR VALUES FROM ('c', 1) TO ('c', 20);
+CREATE FUNCTION trigfunc () RETURNS TRIGGER LANGUAGE plpgsql as $$ BEGIN NEW.e := 'in trigfunc()'; RETURN NEW; END; $$;
+CREATE TRIGGER trig BEFORE INSERT ON part_c_1_c_20 FOR EACH ROW EXECUTE FUNCTION trigfunc();
+UPDATE range_parted r set a = 'c' FROM (VALUES ('a', 1), ('a', 10), ('b', 12)) s(x, y) WHERE s.x = r.a AND s.y = r.b RETURNING tableoid::regclass, *;
+
+DROP TRIGGER trig ON part_c_1_c_20;
+DROP FUNCTION trigfunc;
+
+:init_range_parted;
+CREATE FUNCTION trigfunc () RETURNS TRIGGER LANGUAGE plpgsql as $$ BEGIN RETURN NULL; END; $$;
+CREATE TRIGGER trig BEFORE INSERT ON part_c_1_c_20 FOR EACH ROW EXECUTE FUNCTION trigfunc();
+UPDATE range_parted r set a = 'c' FROM (VALUES ('a', 1), ('a', 10), ('b', 12)) s(x, y) WHERE s.x = r.a AND s.y = r.b RETURNING tableoid::regclass, *;
+:show_data;
+
+DROP TABLE part_c_1_c_20;
+DROP FUNCTION trigfunc;
 
 -- Transition tables with update row movement
 :init_range_parted;