update-partition-key_v3.patch

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index c28cf9c..2878ff3 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -1778,7 +1778,7 @@ ExecRelCheck(ResultRelInfo *resultRelInfo,
  *
  * Note: This is called *iff* resultRelInfo is the main target table.
  */
-static bool
+bool
 ExecPartitionCheck(ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
 				   EState *estate)
 {
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 29c6a6e..287af13 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -62,7 +62,7 @@ static bool ExecOnConflictUpdate(ModifyTableState *mtstate,
 					 EState *estate,
 					 bool canSetTag,
 					 TupleTableSlot **returning);
-
+static void ExecInitPartitionReturningProjection(ModifyTableState *mtstate, Relation root_rel);
 /*
  * Verify that the tuples to be produced by INSERT or UPDATE match the
  * target relation's rowtype
@@ -625,6 +625,7 @@ ExecDelete(ItemPointer tupleid,
 		   TupleTableSlot *planSlot,
 		   EPQState *epqstate,
 		   EState *estate,
+		   bool   *concurrently_deleted,
 		   bool canSetTag)
 {
 	ResultRelInfo *resultRelInfo;
@@ -633,6 +634,9 @@ ExecDelete(ItemPointer tupleid,
 	HeapUpdateFailureData hufd;
 	TupleTableSlot *slot = NULL;
 
+	if (concurrently_deleted)
+		*concurrently_deleted = false;
+
 	/*
 	 * get information on the (current) result relation
 	 */
@@ -776,6 +780,8 @@ ldelete:;
 					}
 				}
 				/* tuple already deleted; nothing to do */
+				if (concurrently_deleted)
+					*concurrently_deleted = true;
 				return NULL;
 
 			default:
@@ -878,7 +884,8 @@ ldelete:;
  * ----------------------------------------------------------------
  */
 static TupleTableSlot *
-ExecUpdate(ItemPointer tupleid,
+ExecUpdate(ModifyTableState *mtstate,
+		   ItemPointer tupleid,
 		   HeapTuple oldtuple,
 		   TupleTableSlot *slot,
 		   TupleTableSlot *planSlot,
@@ -987,6 +994,69 @@ lreplace:;
 			ExecWithCheckOptions(WCO_RLS_UPDATE_CHECK,
 								 resultRelInfo, slot, estate);
 
+		if (resultRelInfo->ri_PartitionCheck &&
+			!ExecPartitionCheck(resultRelInfo, slot, estate))
+		{
+			bool	is_partitioned_table = true;
+
+			if (!mtstate->mt_partition_dispatch_info)
+			{
+				ModifyTable *node = (ModifyTable *) mtstate->ps.plan;
+				Relation root_rel;
+
+				/* root table RT index is at the head of partitioned_rels */
+				if (node->partitioned_rels)
+				{
+					Index	root_rti;
+					Oid		root_oid;
+
+					root_rti = linitial_int(node->partitioned_rels);
+					root_oid = getrelid(root_rti, estate->es_range_table);
+					root_rel = heap_open(root_oid, NoLock);	/* locked by InitPlan */
+				}
+				else
+					root_rel = mtstate->resultRelInfo->ri_RelationDesc;
+
+				is_partitioned_table =
+					root_rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE;
+
+				if (is_partitioned_table)
+					ExecSetupPartitionTupleRouting(
+										root_rel,
+										&mtstate->mt_partition_dispatch_info,
+										&mtstate->mt_partitions,
+										&mtstate->mt_partition_tupconv_maps,
+										&mtstate->mt_partition_tuple_slot,
+										&mtstate->mt_num_dispatch,
+										&mtstate->mt_num_partitions);
+
+				/* Build a projection for each leaf partition rel. */
+				ExecInitPartitionReturningProjection(mtstate, root_rel);
+			}
+
+			/*
+			 * If it's not a partitioned table after all, let it fall through
+			 * the usual error handling.
+			 */
+			if (is_partitioned_table)
+			{
+				bool	concurrently_deleted;
+
+				ExecDelete(tupleid, oldtuple, planSlot, epqstate, estate,
+						   &concurrently_deleted, canSetTag);
+
+				if (concurrently_deleted)
+					return NULL;
+
+				/*
+				 * Don't update estate.es_processed updated again. ExecDelete()
+				 * has already done it above. So use canSetTag=false.
+				 */
+				return ExecInsert(mtstate, slot, planSlot, NULL,
+									  ONCONFLICT_NONE, estate, false);
+			}
+		}
+
 		/*
 		 * Check the constraints of the tuple.  Note that we pass the same
 		 * slot for the orig_slot argument, because unlike ExecInsert(), no
@@ -1313,7 +1383,7 @@ ExecOnConflictUpdate(ModifyTableState *mtstate,
 	 */
 
 	/* Execute UPDATE with projection */
-	*returning = ExecUpdate(&tuple.t_self, NULL,
+	*returning = ExecUpdate(mtstate, &tuple.t_self, NULL,
 							mtstate->mt_conflproj, planSlot,
 							&mtstate->mt_epqstate, mtstate->ps.state,
 							canSetTag);
@@ -1583,12 +1653,12 @@ ExecModifyTable(ModifyTableState *node)
 								  estate, node->canSetTag);
 				break;
 			case CMD_UPDATE:
-				slot = ExecUpdate(tupleid, oldtuple, slot, planSlot,
+				slot = ExecUpdate(node, tupleid, oldtuple, slot, planSlot,
 								&node->mt_epqstate, estate, node->canSetTag);
 				break;
 			case CMD_DELETE:
 				slot = ExecDelete(tupleid, oldtuple, planSlot,
-								&node->mt_epqstate, estate, node->canSetTag);
+								&node->mt_epqstate, estate, NULL, node->canSetTag);
 				break;
 			default:
 				elog(ERROR, "unknown operation");
@@ -1837,7 +1907,6 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	{
 		TupleTableSlot *slot;
 		ExprContext *econtext;
-		List	   *returningList;
 
 		/*
 		 * Initialize result tuple slot and assign its rowtype using the first
@@ -1872,30 +1941,10 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 		}
 
 		/*
-		 * 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.
+		 * Build a projection for each leaf partition rel. 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++;
-		}
+		ExecInitPartitionReturningProjection(mtstate, rel);
 	}
 	else
 	{
@@ -2124,6 +2173,56 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 }
 
 /* ----------------------------------------------------------------
+ *		ExecInitPartitionReturningProjection
+ *
+ * Initialize stuff required to handle RETURNING for leaf partitions.
+ * We don't build the returningList for each partition within the planner, but
+ * simple translation of the varattnos for each partition suffices.  This
+ * actually is helpful only for INSERT case; UPDATE/DELETE are handled
+ * differently.
+ * ----------------------------------------------------------------
+ */
+static void
+ExecInitPartitionReturningProjection(ModifyTableState *mtstate, Relation root_rel)
+{
+	ResultRelInfo  *resultRelInfo = mtstate->mt_partitions;
+	ModifyTable	   *node = (ModifyTable *) mtstate->ps.plan;
+	TupleTableSlot *returning_slot = mtstate->ps.ps_ResultTupleSlot;
+	List		   *returningList;
+	int				i;
+
+	/*
+	 * If there is no returning clause, or if we have already initialized the
+	 * returning projection info, there is nothing to be done.
+	 */
+	if (node->returningLists == NIL ||
+		(resultRelInfo && resultRelInfo->ri_projectReturning != NULL) ||
+		mtstate->mt_num_partitions == 0)
+		return;
+
+	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, root_rel);
+		rliststate = (List *) ExecInitExpr((Expr *) rlist, &mtstate->ps);
+		resultRelInfo->ri_projectReturning =
+			ExecBuildProjectionInfo(rliststate,
+									mtstate->ps.ps_ExprContext,
+									returning_slot,
+									resultRelInfo->ri_RelationDesc->rd_att);
+		resultRelInfo++;
+	}
+}
+
+
+/* ----------------------------------------------------------------
  *		ExecEndModifyTable
  *
  *		Shuts down the plan.
@@ -2154,10 +2253,19 @@ ExecEndModifyTable(ModifyTableState *node)
 	 * Close all the partitioned tables, leaf partitions, and their indices
 	 *
 	 * Remember node->mt_partition_dispatch_info[0] corresponds to the root
-	 * partitioned table, which we must not try to close, because it is the
-	 * main target table of the query that will be closed by ExecEndPlan().
-	 * Also, tupslot is NULL for the root partitioned table.
+	 * partitioned table, which should not be closed if it is the main target
+	 * table of the query, which will be closed by ExecEndPlan(). Also, tupslot
+	 * is NULL for the root partitioned table.
 	 */
+	if (node->mt_num_dispatch > 0)
+	{
+		Relation	root_partition;
+
+		root_partition = node->mt_partition_dispatch_info[0]->reldesc;
+		if (root_partition != node->resultRelInfo->ri_RelationDesc)
+			heap_close(root_partition, NoLock);
+	}
+
 	for (i = 1; i < node->mt_num_dispatch; i++)
 	{
 		PartitionDispatch pd = node->mt_partition_dispatch_info[i];
@@ -2165,6 +2273,7 @@ ExecEndModifyTable(ModifyTableState *node)
 		heap_close(pd->reldesc, NoLock);
 		ExecDropSingleTupleTableSlot(pd->tupslot);
 	}
+
 	for (i = 0; i < node->mt_num_partitions; i++)
 	{
 		ResultRelInfo *resultRelInfo = node->mt_partitions + i;
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index a5c75e7..1fc7cb2 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -225,6 +225,9 @@ extern int ExecFindPartition(ResultRelInfo *resultRelInfo,
 				  PartitionDispatch *pd,
 				  TupleTableSlot *slot,
 				  EState *estate);
+extern bool ExecPartitionCheck(ResultRelInfo *resultRelInfo,
+							TupleTableSlot *slot,
+							EState *estate);
 
 #define EvalPlanQualSetSlot(epqstate, slot)  ((epqstate)->origslot = (slot))
 extern void EvalPlanQualFetchRowMarks(EPQState *epqstate);
diff --git a/src/test/regress/expected/update.out b/src/test/regress/expected/update.out
index 9366f04..99c8046 100644
--- a/src/test/regress/expected/update.out
+++ b/src/test/regress/expected/update.out
@@ -209,13 +209,12 @@ create table part_b_1_b_10 partition of range_parted for values from ('b', 1) to
 create table part_b_10_b_20 partition of range_parted for values from ('b', 10) to ('b', 20);
 insert into part_a_1_a_10 values ('a', 1);
 insert into part_b_10_b_20 values ('b', 10);
--- fail
+-- fail (row movement happens only within the partition subtree)
 update part_a_1_a_10 set a = 'b' where a = 'a';
 ERROR:  new row for relation "part_a_1_a_10" violates partition constraint
 DETAIL:  Failing row contains (b, 1).
+-- ok (row movement)
 update range_parted set b = b - 1 where b = 10;
-ERROR:  new row for relation "part_b_10_b_20" violates partition constraint
-DETAIL:  Failing row contains (b, 9).
 -- ok
 update range_parted set b = b + 1 where b = 10;
 -- cleanup
diff --git a/src/test/regress/sql/update.sql b/src/test/regress/sql/update.sql
index 6637119..7667793 100644
--- a/src/test/regress/sql/update.sql
+++ b/src/test/regress/sql/update.sql
@@ -119,8 +119,9 @@ create table part_b_10_b_20 partition of range_parted for values from ('b', 10)
 insert into part_a_1_a_10 values ('a', 1);
 insert into part_b_10_b_20 values ('b', 10);
 
--- fail
+-- fail (row movement happens only within the partition subtree)
 update part_a_1_a_10 set a = 'b' where a = 'a';
+-- ok (row movement)
 update range_parted set b = b - 1 where b = 10;
 -- ok
 update range_parted set b = b + 1 where b = 10;