update-partition-key.patch

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index a666391..f9da3bd 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -1737,7 +1737,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 95e1589..273120a 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -624,6 +624,7 @@ ExecDelete(ItemPointer tupleid,
 		   TupleTableSlot *planSlot,
 		   EPQState *epqstate,
 		   EState *estate,
+		   bool   *already_deleted,
 		   bool canSetTag)
 {
 	ResultRelInfo *resultRelInfo;
@@ -632,6 +633,9 @@ ExecDelete(ItemPointer tupleid,
 	HeapUpdateFailureData hufd;
 	TupleTableSlot *slot = NULL;
 
+	if (already_deleted)
+		*already_deleted = false;
+
 	/*
 	 * get information on the (current) result relation
 	 */
@@ -775,6 +779,8 @@ ldelete:;
 					}
 				}
 				/* tuple already deleted; nothing to do */
+				if (already_deleted)
+					*already_deleted = true;
 				return NULL;
 
 			default:
@@ -877,7 +883,8 @@ ldelete:;
  * ----------------------------------------------------------------
  */
 static TupleTableSlot *
-ExecUpdate(ItemPointer tupleid,
+ExecUpdate(ModifyTableState *mtstate,
+		   ItemPointer tupleid,
 		   HeapTuple oldtuple,
 		   TupleTableSlot *slot,
 		   TupleTableSlot *planSlot,
@@ -986,6 +993,27 @@ lreplace:;
 			ExecWithCheckOptions(WCO_RLS_UPDATE_CHECK,
 								 resultRelInfo, slot, estate);
 
+		if (resultRelInfo->ri_PartitionCheck &&
+			!ExecPartitionCheck(resultRelInfo, slot, estate))
+		{
+			bool	already_deleted;
+
+			ExecDelete(tupleid, oldtuple, planSlot, epqstate, estate,
+					   &already_deleted, canSetTag);
+
+			if (already_deleted)
+				return NULL;
+			else
+			{
+				/*
+				 * 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
@@ -1312,7 +1340,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);
@@ -1582,12 +1610,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");
@@ -1727,7 +1755,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 
 	/* Build state for INSERT tuple routing */
 	rel = mtstate->resultRelInfo->ri_RelationDesc;
-	if (operation == CMD_INSERT &&
+	if ((operation == CMD_INSERT || operation == CMD_UPDATE) &&
 		rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 	{
 		PartitionDispatch *partition_dispatch_info;
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 02dbe7b..e9a2e07 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -224,6 +224,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 a1e9255..7f27f51 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 d7721ed..92603e9 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;