From 98f3be1315ca534c844287695d94efa990be11f1 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Tue, 27 Dec 2016 16:56:58 +0900
Subject: [PATCH 7/7] Fix RETURNING to work correctly after tuple-routing

In ExecInsert(), do not switch back to the original resultRelInfo until
after we finish ExecProcessReturning(), so that RETURNING projection
is done considering the partition the tuple was routed to.  For the
projection to work correctly, we must initialize the same for each
leaf partition during ModifyTableState initialization.

With this commit, map_partition_varattnos() now accepts one more argument
viz. target_varno.  Previously, it assumed varno = 1 for its input
expressions, which limited its applicability.  It was enought so far
since its usage was limited to partition constraints. To use it with
expressions such as an INSERT's returning list, we must be prepared for
varnos != 1 as in the change above.

Reported by: n/a
Patch by: Amit Langote
Reports: n/a
---
 src/backend/catalog/partition.c        | 13 +++++-----
 src/backend/commands/tablecmds.c       |  1 +
 src/backend/executor/execMain.c        |  4 +--
 src/backend/executor/nodeModifyTable.c | 46 +++++++++++++++++++++++++++-------
 src/include/catalog/partition.h        |  3 ++-
 src/test/regress/expected/insert.out   | 28 ++++++++++++++++++++-
 src/test/regress/sql/insert.sql        | 14 +++++++++++
 7 files changed, 90 insertions(+), 19 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 23eaaf062f..bbc010105e 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -881,7 +881,8 @@ get_qual_from_partbound(Relation rel, Relation parent, Node *bound)
  * different from the parent's.
  */
 List *
-map_partition_varattnos(List *expr, Relation partrel, Relation parent)
+map_partition_varattnos(List *expr, int target_varno,
+						Relation partrel, Relation parent)
 {
 	TupleDesc	tupdesc = RelationGetDescr(parent);
 	AttrNumber	attno;
@@ -906,7 +907,7 @@ map_partition_varattnos(List *expr, Relation partrel, Relation parent)
 	}
 
 	expr = (List *) map_variable_attnos((Node *) expr,
-										1, 0,
+										target_varno, 0,
 										part_attnos,
 										tupdesc->natts,
 										&found_whole_row);
@@ -1517,8 +1518,8 @@ generate_partition_qual(Relation rel)
 		result = list_concat(generate_partition_qual(parent),
 							 copyObject(rel->rd_partcheck));
 
-		/* Mark Vars with correct attnos */
-		result = map_partition_varattnos(result, rel, parent);
+		/* Mark Vars with correct attnos (dummy varno = 1) */
+		result = map_partition_varattnos(result, 1, rel, parent);
 
 		/* Keep the parent locked until commit */
 		heap_close(parent, NoLock);
@@ -1544,8 +1545,8 @@ generate_partition_qual(Relation rel)
 	else
 		result = my_qual;
 
-	/* Mark Vars with correct attnos */
-	result = map_partition_varattnos(result, rel, parent);
+	/* Mark Vars with correct attnos (dummy varno = 1) */
+	result = map_partition_varattnos(result, 1, rel, parent);
 
 	/* Save a copy of *only* this rel's partition qual in the relcache */
 	oldcxt = MemoryContextSwitchTo(CacheMemoryContext);
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 2e51124eb3..caf6b36255 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -13349,6 +13349,7 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 			constr = linitial(partConstraint);
 			my_constr = make_ands_implicit((Expr *) constr);
 			tab->partition_constraint = map_partition_varattnos(my_constr,
+																1,
 																part_rel,
 																rel);
 			/* keep our lock until commit */
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index dfa8bdff92..45940706d7 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -1279,10 +1279,10 @@ InitResultRelInfo(ResultRelInfo *resultRelInfo,
 		/*
 		 * This is not our own partition constraint, but rather an ancestor's.
 		 * So any Vars in it bear the ancestor's attribute numbers.  We must
-		 * switch them to our own.
+		 * switch them to our own. (dummy varno = 1)
 		 */
 		if (partition_check != NIL)
-			partition_check = map_partition_varattnos(partition_check,
+			partition_check = map_partition_varattnos(partition_check, 1,
 													  resultRelationDesc,
 													  partition_root);
 	}
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 825a15f42d..95262114f1 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -262,7 +262,8 @@ ExecInsert(ModifyTableState *mtstate,
 	Relation	resultRelationDesc;
 	Oid			newId;
 	List	   *recheckIndexes = NIL;
-	TupleTableSlot *oldslot = slot;
+	TupleTableSlot *oldslot = slot,
+				   *result = NULL;
 
 	/*
 	 * get the heap tuple out of the tuple table slot, making sure we have a
@@ -574,12 +575,6 @@ ExecInsert(ModifyTableState *mtstate,
 
 	list_free(recheckIndexes);
 
-	if (saved_resultRelInfo)
-	{
-		resultRelInfo = saved_resultRelInfo;
-		estate->es_result_relation_info = resultRelInfo;
-	}
-
 	/*
 	 * Check any WITH CHECK OPTION constraints from parent views.  We are
 	 * required to do this after testing all constraints and uniqueness
@@ -597,9 +592,15 @@ ExecInsert(ModifyTableState *mtstate,
 
 	/* Process RETURNING if present */
 	if (resultRelInfo->ri_projectReturning)
-		return ExecProcessReturning(resultRelInfo, slot, planSlot);
+		result = ExecProcessReturning(resultRelInfo, slot, planSlot);
 
-	return NULL;
+	if (saved_resultRelInfo)
+	{
+		resultRelInfo = saved_resultRelInfo;
+		estate->es_result_relation_info = resultRelInfo;
+	}
+
+	return result;
 }
 
 /* ----------------------------------------------------------------
@@ -1786,6 +1787,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	{
 		TupleTableSlot *slot;
 		ExprContext *econtext;
+		List		*returningList;
 
 		/*
 		 * Initialize result tuple slot and assign its rowtype using the first
@@ -1818,6 +1820,32 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 									 resultRelInfo->ri_RelationDesc->rd_att);
 			resultRelInfo++;
 		}
+
+		/*
+		 * Build a projection for each leaf partition rel.  Note that we
+		 * didn't build the returningList for each partition within the
+		 * planner, but simple translation of the varattnos for each
+		 * partition will suffice.  This only occurs for the INSERT case;
+		 * UPDATE/DELETE are handled above.
+		 */
+		resultRelInfo = mtstate->mt_partitions;
+		returningList = linitial(node->returningLists);
+		for (i = 0; i < mtstate->mt_num_partitions; i++)
+		{
+			Relation	partrel = resultRelInfo->ri_RelationDesc;
+			List	   *rlist,
+					   *rliststate;
+
+			/* varno = node->nominalRelation */
+			rlist = map_partition_varattnos(returningList,
+											node->nominalRelation,
+											partrel, rel);
+			rliststate = (List *) ExecInitExpr((Expr *) rlist, &mtstate->ps);
+			resultRelInfo->ri_projectReturning =
+				ExecBuildProjectionInfo(rliststate, econtext, slot,
+									 resultRelInfo->ri_RelationDesc->rd_att);
+			resultRelInfo++;
+		}
 	}
 	else
 	{
diff --git a/src/include/catalog/partition.h b/src/include/catalog/partition.h
index 78220d6ac6..ba334328f3 100644
--- a/src/include/catalog/partition.h
+++ b/src/include/catalog/partition.h
@@ -77,7 +77,8 @@ extern bool partition_bounds_equal(PartitionKey key,
 extern void check_new_partition_bound(char *relname, Relation parent, Node *bound);
 extern Oid get_partition_parent(Oid relid);
 extern List *get_qual_from_partbound(Relation rel, Relation parent, Node *bound);
-extern List *map_partition_varattnos(List *expr, Relation partrel, Relation parent);
+extern List *map_partition_varattnos(List *expr, int target_varno,
+						Relation partrel, Relation parent);
 extern List *RelationGetPartitionQual(Relation rel);
 
 /* For tuple routing */
diff --git a/src/test/regress/expected/insert.out b/src/test/regress/expected/insert.out
index ed0513d2ff..c8c47ba829 100644
--- a/src/test/regress/expected/insert.out
+++ b/src/test/regress/expected/insert.out
@@ -381,8 +381,34 @@ DETAIL:  Failing row contains (1, 2).
 insert into p1 (a, b) values (2, 3);
 ERROR:  new row for relation "p11" violates partition constraint
 DETAIL:  Failing row contains (3, 2).
+-- check that RETURNING works correctly with tuple-routing
+alter table p drop constraint check_b;
+create table p12 partition of p1 for values from (5) to (10);
+create table p2 (b int not null, a int not null);
+alter table p attach partition p2 for values from (1, 10) to (1, 20);
+create table p3 partition of p for values from (1, 20) to (1, 30);
+create table p4 (like p);
+alter table p4 drop a;
+alter table p4 add a int not null;
+alter table p attach partition p4 for values from (1, 30) to (1, 40);
+with ins (a, b, c) as
+  (insert into p (b, a) select s.a, 1 from generate_series(2, 39) s(a) returning tableoid::regclass, *)
+  select a, b, min(c), max(c) from ins group by a, b order by 1;
+  a  | b | min | max 
+-----+---+-----+-----
+ p11 | 1 |   2 |   4
+ p12 | 1 |   5 |   9
+ p2  | 1 |  10 |  19
+ p3  | 1 |  20 |  29
+ p4  | 1 |  30 |  39
+(5 rows)
+
 -- cleanup
 drop table p cascade;
-NOTICE:  drop cascades to 2 other objects
+NOTICE:  drop cascades to 6 other objects
 DETAIL:  drop cascades to table p1
 drop cascades to table p11
+drop cascades to table p12
+drop cascades to table p2
+drop cascades to table p3
+drop cascades to table p4
diff --git a/src/test/regress/sql/insert.sql b/src/test/regress/sql/insert.sql
index dca89b9286..2bf8cc1659 100644
--- a/src/test/regress/sql/insert.sql
+++ b/src/test/regress/sql/insert.sql
@@ -223,5 +223,19 @@ insert into p values (1, 2);
 -- selected by tuple-routing
 insert into p1 (a, b) values (2, 3);
 
+-- check that RETURNING works correctly with tuple-routing
+alter table p drop constraint check_b;
+create table p12 partition of p1 for values from (5) to (10);
+create table p2 (b int not null, a int not null);
+alter table p attach partition p2 for values from (1, 10) to (1, 20);
+create table p3 partition of p for values from (1, 20) to (1, 30);
+create table p4 (like p);
+alter table p4 drop a;
+alter table p4 add a int not null;
+alter table p attach partition p4 for values from (1, 30) to (1, 40);
+with ins (a, b, c) as
+  (insert into p (b, a) select s.a, 1 from generate_series(2, 39) s(a) returning tableoid::regclass, *)
+  select a, b, min(c), max(c) from ins group by a, b order by 1;
+
 -- cleanup
 drop table p cascade;
-- 
2.11.0

