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

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index a0fd6c15d2..61f1c4b9f2 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -254,7 +254,10 @@ ExecCheckTIDVisible(EState *estate,
  *		ExecInsert
  *
  *		For INSERT, we have to insert the tuple into the target relation
- *		and insert appropriate tuples into the index relations.
+ *		and insert appropriate tuples into the index relations.  When this
+ *		INSERT is a part of an UPDATE of partition-key, then the slot
+ *		containing the inserted tuple is passed back using output parameter
+ *		slot.
  *
  *		Returns RETURNING result if any, otherwise NULL.
  * ----------------------------------------------------------------
@@ -264,7 +267,9 @@ ExecInsert(ModifyTableState *mtstate,
 		   TupleTableSlot *slot,
 		   TupleTableSlot *planSlot,
 		   EState *estate,
-		   bool canSetTag)
+		   bool processReturning,
+		   bool canSetTag,
+		   TupleTableSlot **resultSlot)
 {
 	HeapTuple	tuple;
 	ResultRelInfo *resultRelInfo;
@@ -589,10 +594,14 @@ ExecInsert(ModifyTableState *mtstate,
 	if (resultRelInfo->ri_WithCheckOptions != NIL)
 		ExecWithCheckOptions(WCO_VIEW_CHECK, resultRelInfo, slot, estate);
 
-	/* Process RETURNING if present */
-	if (resultRelInfo->ri_projectReturning)
+	/* Process RETURNING if present and if requested */
+	if (processReturning && resultRelInfo->ri_projectReturning)
 		result = ExecProcessReturning(resultRelInfo, slot, planSlot);
 
+	/* If requested, pass back the inserted row */
+	if (resultSlot)
+		*resultSlot = slot;
+
 	return result;
 }
 
@@ -1067,7 +1076,9 @@ lreplace:;
 		if (partition_constraint_failed)
 		{
 			bool		tuple_deleted;
-			TupleTableSlot *ret_slot;
+			TupleTableSlot *orig_slot = slot;
+			TupleTableSlot *dest_slot = NULL;
+			TupleTableSlot *ret_slot = NULL;
 			TupleTableSlot *epqslot = NULL;
 			PartitionTupleRouting *proute = mtstate->mt_partition_tuple_routing;
 			int			map_index;
@@ -1095,7 +1106,7 @@ lreplace:;
 
 			/*
 			 * Row movement, part 1.  Delete the tuple, but skip RETURNING
-			 * processing. We want to return rows from INSERT.
+			 * processing. We want to return rows after INSERT.
 			 */
 			ExecDelete(mtstate, tupleid, oldtuple, planSlot, epqstate,
 					   estate, false, false /* canSetTag */ ,
@@ -1174,8 +1185,73 @@ lreplace:;
 			slot = ExecPrepareTupleRouting(mtstate, estate, proute,
 										   mtstate->rootResultRelInfo, slot);
 
-			ret_slot = ExecInsert(mtstate, slot, planSlot,
-								  estate, canSetTag);
+			/*
+			 * Row movement, part 2.  Insert the tuple, but skip RETURNING
+			 * processing, because the destination partition's projection may
+			 * not work correctly if some columns in the RETURNING list
+			 * reference planSlot, which is based on the source partition's
+			 * rowtype and junk attributes, and hence whose tuple descriptor
+			 * might differ from what the destination partition's projection
+			 * was told, because partitions may have different rowtypes and
+			 * junk attributes.
+			 */
+			ret_slot = ExecInsert(mtstate, slot, planSlot, estate, false,
+								  canSetTag, &dest_slot);
+			Assert(ret_slot == NULL);
+
+			/* Process RETURNING using the source relation's projection. */
+			if (resultRelInfo->ri_projectReturning && dest_slot)
+			{
+				/*
+				 * Convert the inserted tuple, if necessary.  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
+				 * rowtypes don't match, such cases would be rare in practice.
+				 */
+				if (dest_slot != orig_slot)
+				{
+					ResultRelInfo *destRelInfo = estate->es_result_relation_info;
+					Relation	destRelationDesc = destRelInfo->ri_RelationDesc;
+					TupleConversionMap *map;
+
+					map = convert_tuples_by_name(RelationGetDescr(destRelationDesc),
+												 RelationGetDescr(resultRelationDesc),
+												 gettext_noop("could not convert row type"));
+					if (map)
+					{
+						HeapTuple	old_tuple = ExecMaterializeSlot(dest_slot);
+						HeapTuple	new_tuple;
+
+						new_tuple = do_convert_tuple(old_tuple, map);
+
+						new_tuple->t_self = new_tuple->t_data->t_ctid =
+							old_tuple->t_self;
+						new_tuple->t_tableOid = old_tuple->t_tableOid;
+
+						HeapTupleHeaderSetXmin(new_tuple->t_data,
+											   HeapTupleHeaderGetRawXmin(old_tuple->t_data));
+						HeapTupleHeaderSetCmin(new_tuple->t_data,
+											   HeapTupleHeaderGetRawCommandId(old_tuple->t_data));
+						HeapTupleHeaderSetXmax(new_tuple->t_data,
+											   InvalidTransactionId);
+
+						if (RelationGetDescr(destRelationDesc)->tdhasoid)
+						{
+							Assert(RelationGetDescr(resultRelationDesc)->tdhasoid);
+							HeapTupleSetOid(new_tuple,
+											HeapTupleGetOid(old_tuple));
+						}
+
+						ExecStoreTuple(new_tuple, orig_slot,
+									   InvalidBuffer, true);
+					}
+					else
+						orig_slot = dest_slot;
+				}
+
+				ret_slot = ExecProcessReturning(resultRelInfo, orig_slot,
+												planSlot);
+			}
 
 			/* Revert ExecPrepareTupleRouting's node change. */
 			estate->es_result_relation_info = resultRelInfo;
@@ -2156,8 +2232,8 @@ ExecModifyTable(PlanState *pstate)
 				if (proute)
 					slot = ExecPrepareTupleRouting(node, estate, proute,
 												   resultRelInfo, slot);
-				slot = ExecInsert(node, slot, planSlot,
-								  estate, node->canSetTag);
+				slot = ExecInsert(node, slot, planSlot, estate,
+								  true, node->canSetTag, 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 2083345c8e..c35c0dd9f8 100644
--- a/src/test/regress/expected/update.out
+++ b/src/test/regress/expected/update.out
@@ -424,6 +424,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 8754ccb7b0..e6a61ace5b 100644
--- a/src/test/regress/sql/update.sql
+++ b/src/test/regress/sql/update.sql
@@ -223,6 +223,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;