From bc3ab4f67da7da6f740bf204671876c88b4f25d4 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Thu, 11 Jun 2020 09:54:08 +0900
Subject: [PATCH v2] Fix a bug with RETURNING when UPDATE "moves" tuple

Currently, the RETURNING list is projected by ExecInsert() that
is running on the destination partition and hence using that
partition's version of the RETURNING projection.  If the RETURNING
list contains some non-target relation's columns, they refer to the
corresponding columns in the plan's output tuple ('planSlot'), whose
descriptor is based on the *source* result relation's descriptor,
which can lead to attribute mismatch errors when projecting with
the destination relation's RETURNING projection.

Fix is to perform the projection in ExecUpdate using the source
relation's RETURNING projection after returning from ExecInsert().
---
 src/backend/executor/nodeModifyTable.c | 48 +++++++++++++++++++++++++++++-----
 src/test/regress/expected/update.out   | 22 ++++++++++++++++
 src/test/regress/sql/update.sql        | 15 +++++++++++
 3 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 20a4c47..97b1c99 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -368,7 +368,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 +376,8 @@ ExecInsert(ModifyTableState *mtstate,
 		   TupleTableSlot *slot,
 		   TupleTableSlot *planSlot,
 		   EState *estate,
-		   bool canSetTag)
+		   bool canSetTag,
+		   bool processReturning)
 {
 	ResultRelInfo *resultRelInfo;
 	Relation	resultRelationDesc;
@@ -677,7 +678,7 @@ 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);
 
 	return result;
@@ -1202,11 +1203,13 @@ lreplace:;
 		if (partition_constraint_failed)
 		{
 			bool		tuple_deleted;
+			TupleTableSlot *orig_slot = slot;
 			TupleTableSlot *ret_slot;
 			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 +1310,17 @@ 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);
 
 			/* Revert ExecPrepareTupleRouting's node change. */
 			estate->es_result_relation_info = resultRelInfo;
@@ -1319,6 +1330,30 @@ lreplace:;
 				mtstate->mt_transition_capture->tcs_map = saved_tcs_map;
 			}
 
+			/* Process RETURNING using the source relation's projection. */
+			if (resultRelInfo->ri_projectReturning)
+			{
+				AttrMap *map;
+
+				/*
+				 * 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 (slot != orig_slot)
+				{
+					map = build_attrmap_by_name_if_req(RelationGetDescr(destRel->ri_RelationDesc),
+													   RelationGetDescr(resultRelInfo->ri_RelationDesc));
+					if (map)
+						slot = execute_attr_map_slot(map, slot, orig_slot);
+				}
+				return ExecProcessReturning(resultRelInfo, slot, planSlot);
+			}
+
 			return ret_slot;
 		}
 
@@ -2244,7 +2279,8 @@ ExecModifyTable(PlanState *pstate)
 					slot = ExecPrepareTupleRouting(node, estate, proute,
 												   resultRelInfo, slot);
 				slot = ExecInsert(node, slot, planSlot,
-								  estate, node->canSetTag);
+								  estate, node->canSetTag,
+								  true);
 				/* 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 bf939d7..e272ff6 100644
--- a/src/test/regress/expected/update.out
+++ b/src/test/regress/expected/update.out
@@ -445,6 +445,28 @@ 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_a_1_a_10  | c |  1 |   1 | 1 | in trigfunc() | a |  1
+ part_a_10_a_20 | c | 10 | 200 | 1 | in trigfunc() | a | 10
+ part_c_1_100   | c | 12 |  96 | 1 | in trigfunc() | b | 12
+(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 8c558a7..ae38983 100644
--- a/src/test/regress/sql/update.sql
+++ b/src/test/regress/sql/update.sql
@@ -236,6 +236,21 @@ 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 TABLE part_c_1_c_20;
+DROP FUNCTION trigfunc;
 
 -- Transition tables with update row movement
 :init_range_parted;
-- 
1.8.3.1

