From 20b1c4c6c7206e364f169cc30b3bdc1b0114eb1b Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 12 Mar 2025 21:10:31 +0900
Subject: [PATCH v1] Ensure first result relation is included even if pruned

This started as an investigation into crashes and misbehavior in MERGE
statements when all target partitions were pruned. These issues became
possible after commit cbc127917e, which introduced tracking of
unpruned relids to avoid processing pruned relations.

Some executor code paths rely on ModifyTableState.resultRelInfo[0] being
present and initialized, even if no result relations remain after
pruning. For example, ExecMerge() and ExecMergeNotMatched() use the first
resultRelInfo to determine the appropriate action. Similarly,
ExecInitPartitionInfo() assumes at least one result relation exists.

To preserve these assumptions, ExecDoInitialPruning() now ensures the
first result relation is locked and marked as unpruned, and
ExecInitModifyTable() includes it in the initialized result relation list
even if it was pruned.

This avoids crashes and undefined behavior while allowing pruning of
other result relations to proceed as before.

Bug: #18830
Reported-by: Robins Tharakan <tharakan@gmail.com>
Diagnozed-by: Tender Wang <tndrwang@gmail.com>
Diagnozed-by: Dean Rasheed <dean.a.rasheed@gmail.com>
Suggested-by: Dean Rasheed <dean.a.rasheed@gmail.com>
Discussion: https://postgr.es/m/18830-1f31ea1dc930d444%40postgresql.org
---
 src/backend/executor/execPartition.c   | 23 +++++++++++++++++++++++
 src/backend/executor/nodeModifyTable.c | 12 +++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 5cd5e2eeb80..c93c2afad03 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1819,6 +1819,7 @@ adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap)
 void
 ExecDoInitialPruning(EState *estate)
 {
+	PlannedStmt *stmt = estate->es_plannedstmt;
 	ListCell   *lc;
 	List	   *locked_relids = NIL;
 
@@ -1867,6 +1868,28 @@ ExecDoInitialPruning(EState *estate)
 												validsubplans);
 	}
 
+	/*
+	 * Lock the first result relation even if it was pruned. Some executor
+	 * paths (e.g., in nodeModifyTable.c and execPartition.c) expect this
+	 * relation to be locked regardless of pruning. Also add it to
+	 * es_unpruned_relids so that ExecInitModifyTable() processes it.
+	 */
+	if (stmt->resultRelations)
+	{
+		Index	firstResultRel = linitial_int(stmt->resultRelations);
+
+		if (!bms_is_member(firstResultRel, estate->es_unpruned_relids))
+		{
+			RangeTblEntry *rte = exec_rt_fetch(firstResultRel, estate);
+
+			Assert(rte->rtekind == RTE_RELATION && rte->rellockmode != NoLock);
+			LockRelationOid(rte->relid, rte->rellockmode);
+			locked_relids = lappend_int(locked_relids, firstResultRel);
+			estate->es_unpruned_relids = bms_add_member(estate->es_unpruned_relids,
+														firstResultRel);
+		}
+	}
+
 	/*
 	 * Release the useless locks if the plan won't be executed.  This is the
 	 * same as what CheckCachedPlan() in plancache.c does.
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index a9142796743..bcc80a79046 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -4495,7 +4495,16 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	{
 		Index		rti = lfirst_int(l);
 
-		if (bms_is_member(rti, estate->es_unpruned_relids))
+		/*
+		 * The first result relation might have been pruned, but we still need
+		 * to include it to support executor code paths that access
+		 * ModifyTableState.resultRelInfo[0], e.g., for evaluating result
+		 * relation expressions.
+		 *
+		 * ExecDoInitialPruning() already locked it and added it to
+		 * es_unpruned_relids.
+		 */
+		if (i == 0 || bms_is_member(rti, estate->es_unpruned_relids))
 		{
 			resultRelations = lappend_int(resultRelations, rti);
 			if (node->withCheckOptionLists)
@@ -4536,6 +4545,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 		i++;
 	}
 	nrels = list_length(resultRelations);
+	Assert (nrels > 0);
 
 	/*
 	 * create state structure
-- 
2.43.0

