From f586635ab49f3027546a7bda4c4f6017b946f333 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Sat, 4 Apr 2026 20:43:14 +0900
Subject: [PATCH v11 4/4] Use pruning-aware locking for single-statement cached
 plans

For single-statement reused generic plans, perform initial partition
pruning before acquiring execution locks, then lock only the
surviving partitions.

Add ExecutorPrepAndLock() which encapsulates the pruning-aware lock
sequence: lock unprunable relations, call ExecutorPrep() to run
initial pruning, then lock survivors.  Plan validity is checked
after each step; ExecutorPrepCleanup() handles the case where the
plan is invalidated between prep and execution.

Extend PortalLockCachedPlan() to use the pruning-aware path for
eligible plans (single-statement reused generic, non-utility).
All other cases continue using the conservative lock-all path
from the previous commit.

Track firstResultRels in PlannerGlobal and PlannedStmt so they
are locked even if pruned, preserving ExecInitModifyTable()
assumptions about the first result relation being available.

Multi-statement CachedPlans (from rule rewriting) always use
conservative locking, since PortalRunMulti() executes statements
sequentially with CCI between them and later statements' pruning
expressions may depend on earlier ones' effects.  In principle,
this could be relaxed if the planner can prove that no pruning
expression reads state modified by an earlier statement, but that
is left for a future patch.

Regression tests are included to verify:

- Only surviving partitions are locked when pruning is enabled, and
  all partitions are locked when it is disabled (pg_locks inspection).
- Multiple ModifyTable nodes (via writable CTEs) handle the case where
  all target partitions are pruned, exercising firstResultRels.
- Plan invalidation during pruning-aware lock setup (DDL triggered by
  a pruning expression) discards the prep state and replans cleanly.
- Multi-statement CachedPlans (from rule rewriting) fall back to
  locking all partitions, avoiding stale pruning results.

Note for extension authors: code that accesses partition relations
through EState must check that the RT index is a member of
es_unpruned_relids before opening the relation.  Previously this
was an optimization; it is now a correctness requirement, because
pruned partitions may not be locked.
---
 src/backend/commands/explain.c                |  45 +++--
 src/backend/commands/prepare.c                |  30 ++-
 src/backend/executor/execMain.c               | 142 ++++++++++++++
 src/backend/executor/nodeModifyTable.c        |   5 +-
 src/backend/optimizer/plan/planner.c          |   1 +
 src/backend/optimizer/plan/setrefs.c          |  18 ++
 src/backend/tcop/pquery.c                     |  54 ++++-
 src/backend/utils/cache/plancache.c           |  16 ++
 src/include/commands/explain.h                |   3 +-
 src/include/executor/executor.h               |   4 +
 src/include/nodes/pathnodes.h                 |   3 +
 src/include/nodes/plannodes.h                 |  10 +
 src/include/utils/plancache.h                 |   2 +
 src/test/regress/expected/partition_prune.out | 184 ++++++++++++++++++
 src/test/regress/expected/plancache.out       |  63 ++++++
 src/test/regress/sql/partition_prune.sql      | 116 +++++++++++
 src/test/regress/sql/plancache.sql            |  52 +++++
 17 files changed, 720 insertions(+), 28 deletions(-)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index e4b70166b0e..60cd912ace1 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -374,7 +374,8 @@ standard_ExplainOneQuery(Query *query, int cursorOptions,
 	/* run it (if needed) and produce output */
 	ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
 				   &planduration, (es->buffers ? &bufusage : NULL),
-				   es->memory ? &mem_counters : NULL);
+				   es->memory ? &mem_counters : NULL,
+				   NULL);
 }
 
 /*
@@ -498,7 +499,8 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 			   const char *queryString, ParamListInfo params,
 			   QueryEnvironment *queryEnv, const instr_time *planduration,
 			   const BufferUsage *bufusage,
-			   const MemoryContextCounters *mem_counters)
+			   const MemoryContextCounters *mem_counters,
+			   QueryDesc *prep_qd)
 {
 	DestReceiver *dest;
 	QueryDesc  *queryDesc;
@@ -527,13 +529,6 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	 */
 	INSTR_TIME_SET_CURRENT(starttime);
 
-	/*
-	 * Use a snapshot with an updated command ID to ensure this query sees
-	 * results of any previously executed queries.
-	 */
-	PushCopiedSnapshot(GetActiveSnapshot());
-	UpdateActiveSnapshotCommandId();
-
 	/*
 	 * We discard the output if we have no use for it.  If we're explaining
 	 * CREATE TABLE AS, we'd better use the appropriate tuple receiver, while
@@ -549,10 +544,34 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	else
 		dest = None_Receiver;
 
-	/* Create a QueryDesc for the query */
-	queryDesc = CreateQueryDesc(plannedstmt, queryString,
-								GetActiveSnapshot(), InvalidSnapshot,
-								dest, params, queryEnv, instrument_option);
+	/*
+	 * Create a QueryDesc for the query, or use the one provided by the
+	 * caller.  When reusing a prep QueryDesc, its snapshot was set at
+	 * creation time; we push it as active for ExecutorStart and override the
+	 * destination and instrument options, which were not known when the
+	 * caller created it.
+	 */
+	if (prep_qd)
+	{
+		PushActiveSnapshot(GetActiveSnapshot());
+		queryDesc = prep_qd;
+		Assert(queryDesc->dest == None_Receiver);
+		queryDesc->dest = dest;
+		queryDesc->instrument_options = instrument_option;
+	}
+	else
+	{
+		/*
+		 * Use a snapshot with an updated command ID to ensure this query sees
+		 * results of any previously executed queries.
+		 */
+		PushCopiedSnapshot(GetActiveSnapshot());
+		UpdateActiveSnapshotCommandId();
+		queryDesc = CreateQueryDesc(plannedstmt, queryString,
+									GetActiveSnapshot(), InvalidSnapshot,
+									dest, params, queryEnv,
+									instrument_option);
+	}
 
 	/* Select execution options */
 	if (es->analyze)
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 03d7a98fc58..3bbbc052149 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -588,6 +588,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	MemoryContextCounters mem_counters;
 	MemoryContext planner_ctx = NULL;
 	MemoryContext saved_ctx = NULL;
+	QueryDesc  *prep_qd = NULL;
 
 	if (es->memory)
 	{
@@ -640,8 +641,31 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 							  pstate->p_queryEnv);
 		plan_list = cplan->stmt_list;
 
-		if (AcquireExecutorLocks(cplan))
+		if (!CachedPlanCanPrep(cplan, entry->plansource))
+		{
+			if (AcquireExecutorLocks(cplan))
+				break;
+			ReleaseCachedPlan(cplan, CurrentResourceOwner);
+			continue;
+		}
+
+		prep_qd = CreateQueryDesc(linitial_node(PlannedStmt, plan_list),
+								  query_string,
+								  GetActiveSnapshot(),
+								  InvalidSnapshot,
+								  None_Receiver,	/* ExplainOnePlan will fix */
+								  paramLI,
+								  pstate->p_queryEnv,
+								  0 /* ExplainOnePlan will fix */ );
+		if (ExecutorPrepAndLock(prep_qd,
+								CurrentResourceOwner,
+								es->generic ? EXEC_FLAG_EXPLAIN_GENERIC : 0,
+								&cplan->is_valid))
 			break;
+
+		/* Try again. */
+		ExecutorPrepCleanup(prep_qd);
+		FreeQueryDesc(prep_qd);
 		ReleaseCachedPlan(cplan, CurrentResourceOwner);
 	}
 
@@ -664,6 +688,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	plan_list = cplan->stmt_list;
 
 	/* Explain each query */
+	Assert(prep_qd == NULL || list_length(plan_list) == 1);
 	foreach(p, plan_list)
 	{
 		PlannedStmt *pstmt = lfirst_node(PlannedStmt, p);
@@ -671,7 +696,8 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 		if (pstmt->commandType != CMD_UTILITY)
 			ExplainOnePlan(pstmt, into, es, query_string, paramLI, pstate->p_queryEnv,
 						   &planduration, (es->buffers ? &bufusage : NULL),
-						   es->memory ? &mem_counters : NULL);
+						   es->memory ? &mem_counters : NULL,
+						   prep_qd);
 		else
 			ExplainOneUtility(pstmt->utilityStmt, into, es, pstate, paramLI);
 
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 735c80e08a9..7333c0f66d5 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -324,6 +324,124 @@ standard_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	MemoryContextSwitchTo(oldcontext);
 }
 
+/*
+ * LockRangeTableRelids
+ * 		Acquire or release locks on the specified relids, which reference
+ * 		entries in the provided range table.
+ *
+ * Helper for AcquireExecutorLocksUnpruned().
+ */
+static void
+LockRangeTableRelids(List *rtable, Bitmapset *relids, bool acquire)
+{
+	int			rtindex = -1;
+
+	while ((rtindex = bms_next_member(relids, rtindex)) >= 0)
+	{
+		RangeTblEntry *rte = list_nth_node(RangeTblEntry, rtable, rtindex - 1);
+
+		Assert(rte->rtekind == RTE_RELATION ||
+			   (rte->rtekind == RTE_SUBQUERY && OidIsValid(rte->relid)));
+
+		/*
+		 * Acquire the appropriate type of lock on each relation OID. Note
+		 * that we don't actually try to open the rel, and hence will not fail
+		 * if it's been dropped entirely --- we'll just transiently acquire a
+		 * non-conflicting lock.
+		 */
+		if (acquire)
+			LockRelationOid(rte->relid, rte->rellockmode);
+		else
+			UnlockRelationOid(rte->relid, rte->rellockmode);
+	}
+}
+
+/*
+ * AcquireExecutorLocksPrepared
+ *
+ * Acquire or release execution locks using pruning results already computed
+ * by ExecutorPrep() and stored in queryDesc->estate.
+ *
+ * This is intended for single-statement reused generic-plan paths that
+ * choose pruning-aware locking instead of the conservative
+ * AcquireExecutorLocks() path.
+ */
+static void
+AcquireExecutorLocksPrepared(QueryDesc *queryDesc, bool acquire)
+{
+	PlannedStmt *plannedstmt = queryDesc->plannedstmt;
+	EState	   *estate = queryDesc->estate;
+	Bitmapset  *lock_relids;
+	ListCell   *lc;
+
+	Assert(queryDesc != NULL);
+	Assert(estate != NULL);
+	Assert(plannedstmt != NULL);
+	Assert(plannedstmt->commandType != CMD_UTILITY);
+
+	lock_relids = bms_difference(estate->es_unpruned_relids,
+								 plannedstmt->unprunableRelids);
+
+	/*
+	 * Keep the first result relation of each ModifyTable locked even if
+	 * pruning removed all target partitions.  ExecInitModifyTable() relies on
+	 * one such relation remaining available.
+	 */
+	foreach(lc, plannedstmt->firstResultRels)
+	{
+		Index		rti = lfirst_int(lc);
+
+		lock_relids = bms_add_member(lock_relids, rti);
+	}
+
+	LockRangeTableRelids(plannedstmt->rtable, lock_relids, acquire);
+
+	bms_free(lock_relids);
+
+}
+
+/*
+ * ExecutorPrepAndLock
+ *		Perform pruning-aware locking for a single PlannedStmt.
+ *
+ * Locks unprunable relations first, then runs ExecutorPrep() to
+ * determine which partitions survive initial pruning, then locks
+ * only those survivors.  Checks *is_valid after each locking step
+ * to detect plan invalidation (e.g., from concurrent DDL or DDL
+ * triggered by a pruning expression).
+ *
+ * Returns true if the plan is still valid and all needed locks are
+ * held.  Returns false if the plan was invalidated at any point, in
+ * which case all acquired locks have been released and the caller
+ * should discard the QueryDesc and retry with a fresh plan.
+ */
+bool
+ExecutorPrepAndLock(QueryDesc *queryDesc, ResourceOwner owner,
+					int eflags, bool *is_valid)
+{
+	PlannedStmt *pstmt = queryDesc->plannedstmt;
+
+	/* Lock unprunable rels before pruning can access them. */
+	LockRangeTableRelids(pstmt->rtable, pstmt->unprunableRelids, true);
+	if (!*is_valid)
+	{
+		LockRangeTableRelids(pstmt->rtable, pstmt->unprunableRelids, false);
+		return false;
+	}
+
+	/* Run pruning and lock survivors. */
+	ExecutorPrep(queryDesc, owner, eflags);
+	AcquireExecutorLocksPrepared(queryDesc, true);
+	if (!*is_valid)
+	{
+		AcquireExecutorLocksPrepared(queryDesc, false);
+		LockRangeTableRelids(pstmt->rtable, pstmt->unprunableRelids, false);
+		return false;
+	}
+
+	return true;
+}
+
 /*
  * ExecutorPrep
  *
@@ -382,6 +500,30 @@ ExecutorPrep(QueryDesc *queryDesc, ResourceOwner owner, int eflags)
 	CurrentResourceOwner = oldowner;
 }
 
+/*
+ * ExecutorPrepCleanup
+ *		Clean up an EState that was created by ExecutorPrep() but never
+ *		passed to ExecutorStart().  This happens when the plan is
+ *		invalidated between prep and execution, and the caller must
+ *		discard the prepped state before retrying with a fresh plan.
+ *
+ * Unlike ExecutorEnd(), this does not expect a fully initialized
+ * plan state tree -- only the range table relations and the
+ * EState itself need to be freed.
+ */
+void
+ExecutorPrepCleanup(QueryDesc *queryDesc)
+{
+	EState	   *estate = queryDesc->estate;
+
+	if (estate == NULL)
+		return;
+
+	ExecCloseRangeTableRelations(estate);
+	FreeExecutorState(estate);
+	queryDesc->estate = NULL;
+}
+
 /* ----------------------------------------------------------------
  *		ExecutorRun
  *
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index dfd7b33aa9b..8bc5c36e09d 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -5112,8 +5112,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	 * as a reference for building the ResultRelInfo of the target partition.
 	 * In either case, it doesn't matter which result relation is kept, so we
 	 * just keep the first one, if all others have been pruned.  See also,
-	 * ExecDoInitialPruning(), which ensures that this first result relation
-	 * has been locked.
+	 * AcquireExecutorLocksUnpruned(), which ensures that this first result
+	 * relation has been locked.
 	 */
 	i = 0;
 	foreach(l, node->resultRelations)
@@ -5127,6 +5127,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 			/* all result relations pruned; keep the first one */
 			keep_rel = true;
 			rti = linitial_int(node->resultRelations);
+			Assert(list_member_int(estate->es_plannedstmt->firstResultRels, rti));
 			i = 0;
 		}
 
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 4ec76ce31a9..ace1cbacc91 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -657,6 +657,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 											  glob->prunableRelids);
 	result->permInfos = glob->finalrteperminfos;
 	result->subrtinfos = glob->subrtinfos;
+	result->firstResultRels = glob->firstResultRels;
 	result->appendRelations = glob->appendRelations;
 	result->subplans = glob->subplans;
 	result->rewindPlanIDs = glob->rewindPlanIDs;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index ff0e875f2a2..6ee51f06920 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -384,6 +384,24 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 		}
 	}
 
+	/*
+	 * Record the first result relation if it belongs to the set of initially
+	 * prunable relations.  We use bms_next_member() to get the
+	 * lowest-numbered leaf result rel, which matches
+	 * linitial_int(ModifyTable.resultRelations) because partition expansion
+	 * preserves RT index order.  ExecInitModifyTable() asserts that the
+	 * recorded index matches what it actually needs.
+	 */
+	if (root->leaf_result_relids)
+	{
+		Index		firstResultRel = bms_next_member(root->leaf_result_relids, -1);
+
+		firstResultRel += rtoffset;
+		if (bms_is_member(firstResultRel, root->glob->prunableRelids))
+			root->glob->firstResultRels =
+				lappend_int(root->glob->firstResultRels, firstResultRel);
+	}
+
 	return result;
 }
 
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 1b22515d56e..af732821139 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -59,7 +59,9 @@ static uint64 DoPortalRunFetch(Portal portal,
 							   long count,
 							   DestReceiver *dest);
 static void DoPortalRewind(Portal portal);
-static bool PortalLockCachedPlan(Portal portal);
+static bool PortalLockCachedPlan(Portal portal, bool do_prep,
+								 ParamListInfo params,
+								 QueryDesc **queryDesc_p);
 
 
 /*
@@ -492,9 +494,14 @@ restart:
 				 * the destination to DestNone.
 				 *
 				 * If the portal is backed by a cached plan, acquire execution
-				 * locks via PortalLockCachedPlan().  If the plan is
-				 * invalidated during locking, it replans and may change the
-				 * portal strategy, requiring us to restart PortalStart().
+				 * locks via PortalLockCachedPlan().  For eligible plans
+				 * (single-statement reused generic), this performs
+				 * pruning-aware locking: it runs ExecutorPrep() on the
+				 * QueryDesc to determine which partitions survive initial
+				 * pruning, then locks only those.  If the plan is invalidated
+				 * during this process, it replans and rebuilds the QueryDesc.
+				 * If replanning changes the portal strategy, we must restart
+				 * PortalStart() to redispatch.
 				 */
 				queryDesc = CreateQueryDesc(linitial_node(PlannedStmt, portal->stmts),
 											portal->sourceText,
@@ -506,7 +513,7 @@ restart:
 											0);
 				if (portal->cplan)
 				{
-					if (PortalLockCachedPlan(portal))
+					if (PortalLockCachedPlan(portal, true, params, &queryDesc))
 					{
 						PopActiveSnapshot();
 						goto restart;
@@ -552,7 +559,7 @@ restart:
 			case PORTAL_ONE_MOD_WITH:
 				if (portal->cplan)
 				{
-					if (PortalLockCachedPlan(portal))
+					if (PortalLockCachedPlan(portal, false, NULL, NULL))
 						goto restart;
 				}
 
@@ -608,7 +615,7 @@ restart:
 				 */
 				if (portal->cplan)
 				{
-					if (PortalLockCachedPlan(portal))
+					if (PortalLockCachedPlan(portal, false, NULL, NULL))
 						goto restart;
 				}
 
@@ -1825,15 +1832,32 @@ EnsurePortalSnapshotExists(void)
  *		Acquire execution locks for a cached-plan-backed portal,
  *		retrying with a fresh plan if the current one is invalidated.
  *
+ * If do_prep is true and the plan is eligible (single-statement reused
+ * generic plan), performs pruning-aware locking via ExecutorPrep() and
+ * populates portal->queryDesc with the prepped QueryDesc.  Otherwise
+ * falls back to locking all relations in the plan.
+ *
  * Returns true if replanning changed portal->strategy, meaning the
- * caller must redispatch.  Returns false once locks are held.
+ * caller must redispatch.  Returns false once locks are held and the
+ * plan is valid for execution.
  */
 static bool
-PortalLockCachedPlan(Portal portal)
+PortalLockCachedPlan(Portal portal, bool do_prep,
+					 ParamListInfo params,
+					 QueryDesc **prep_qd)
 {
 	PortalStrategy start_strategy = portal->strategy;
 
-	if (AcquireExecutorLocks(portal->cplan))
+	if (do_prep && CachedPlanCanPrep(portal->cplan, portal->plansource))
+	{
+		Assert(prep_qd);
+		if (ExecutorPrepAndLock(*prep_qd, portal->resowner, 0,
+								&portal->cplan->is_valid))
+			return false;
+		ExecutorPrepCleanup(*prep_qd);
+		FreeQueryDesc(*prep_qd);
+	}
+	else if (AcquireExecutorLocks(portal->cplan))
 		return false;
 
 	/* Replan.  Locks will be taken freshly. */
@@ -1849,5 +1873,15 @@ PortalLockCachedPlan(Portal portal)
 	if (portal->strategy != start_strategy)
 		return true;
 
+	if (prep_qd)
+	{
+		Assert(list_length(portal->stmts) == 1);
+		*prep_qd = CreateQueryDesc(linitial_node(PlannedStmt, portal->stmts),
+								   portal->sourceText,
+								   GetActiveSnapshot(), InvalidSnapshot,
+								   None_Receiver, params,
+								   portal->queryEnv, 0);
+	}
+
 	return false;
 }
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index f7fe366859c..fca2f84081e 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -1977,6 +1977,22 @@ AcquireExecutorLocks(CachedPlan *cplan)
 	return true;
 }
 
+/*
+ * CachedPlanCanPrep
+ *		Check whether a cached plan is eligible for pruning-aware locking
+ *		via ExecutorPrepAndLock().
+ *
+ * Only single-statement reused generic plans with a non-utility command
+ * qualify.
+ */
+bool
+CachedPlanCanPrep(CachedPlan *cplan, CachedPlanSource *plansource)
+{
+	return (cplan == plansource->gplan &&
+			list_length(cplan->stmt_list) == 1 &&
+			linitial_node(PlannedStmt, cplan->stmt_list)->commandType != CMD_UTILITY);
+}
+
 /*
  * AcquirePlannerLocks: acquire locks needed for planning of a querytree list;
  * or release them if acquire is false.
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 472e141bba3..3a03355e6b6 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -69,7 +69,8 @@ extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into,
 						   ParamListInfo params, QueryEnvironment *queryEnv,
 						   const instr_time *planduration,
 						   const BufferUsage *bufusage,
-						   const MemoryContextCounters *mem_counters);
+						   const MemoryContextCounters *mem_counters,
+						   QueryDesc *prep_qd);
 
 extern void ExplainPrintPlan(ExplainState *es, QueryDesc *queryDesc);
 extern void ExplainPrintTriggers(ExplainState *es,
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 491c4886506..fef5aadcdfa 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -21,6 +21,7 @@
 #include "nodes/lockoptions.h"
 #include "nodes/parsenodes.h"
 #include "utils/memutils.h"
+#include "utils/resowner.h"
 
 
 /*
@@ -235,6 +236,9 @@ ExecGetJunkAttribute(TupleTableSlot *slot, AttrNumber attno, bool *isNull)
  */
 extern void ExecutorStart(QueryDesc *queryDesc, int eflags);
 extern void standard_ExecutorStart(QueryDesc *queryDesc, int eflags);
+extern bool ExecutorPrepAndLock(QueryDesc *queryDesc, ResourceOwner owner,
+								int eflags, bool *is_valid);
+extern void ExecutorPrepCleanup(QueryDesc *queryDesc);
 extern void ExecutorRun(QueryDesc *queryDesc,
 						ScanDirection direction, uint64 count);
 extern void standard_ExecutorRun(QueryDesc *queryDesc,
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 693b879f76d..8753e05152b 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -217,6 +217,9 @@ typedef struct PlannerGlobal
 	/* "flat" list of integer RT indexes */
 	List	   *resultRelations;
 
+	/* "flat" list of integer RT indexes (one per ModifyTable node) */
+	List	   *firstResultRels;
+
 	/* "flat" list of AppendRelInfos */
 	List	   *appendRelations;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 14a1dfed2b9..7f6f7cda781 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -120,6 +120,16 @@ typedef struct PlannedStmt
 	/* RT indexes of relations targeted by INSERT/UPDATE/DELETE/MERGE */
 	Bitmapset  *resultRelationRelids;
 
+	/*
+	 * rtable indexes of first target relation in each ModifyTable node in the
+	 * plan for INSERT/UPDATE/DELETE/MERGE.  NIL if resultRelations is NIL.
+	 *
+	 * These are used by AcquireExecutorLocksUnpruned() to ensure that the
+	 * first result rel for each ModifyTable remains locked even if pruned;
+	 * see ExecInitModifyTable() for the executor side assumptions.
+	 */
+	List	   *firstResultRels;
+
 	/* list of AppendRelInfo nodes */
 	List	   *appendRelations;
 
diff --git a/src/include/utils/plancache.h b/src/include/utils/plancache.h
index e0fc403e717..2941d3a301b 100644
--- a/src/include/utils/plancache.h
+++ b/src/include/utils/plancache.h
@@ -254,4 +254,6 @@ extern bool CachedPlanIsSimplyValid(CachedPlanSource *plansource,
 extern CachedExpression *GetCachedExpression(Node *expr);
 extern void FreeCachedExpression(CachedExpression *cexpr);
 
+extern bool CachedPlanCanPrep(CachedPlan *cplan, CachedPlanSource *plansource);
+
 #endif							/* PLANCACHE_H */
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
index deacdd75807..61781389d2f 100644
--- a/src/test/regress/expected/partition_prune.out
+++ b/src/test/regress/expected/partition_prune.out
@@ -4824,3 +4824,187 @@ select min(a) over (partition by a order by a) from part_abc where a >= stable_o
 
 drop view part_abc_view;
 drop table part_abc;
+--
+-- Verify that pruning-aware locking skips pruned partitions
+-- when reusing a generic cached plan.
+--
+set plan_cache_mode to force_generic_plan;
+create table prunelock_p (a int) partition by list (a);
+create table prunelock_p1 partition of prunelock_p for values in (1);
+create table prunelock_p2 partition of prunelock_p for values in (2);
+create table prunelock_p3 partition of prunelock_p for values in (3);
+prepare prunelock_q (int) as select * from prunelock_p where a = $1;
+-- Force generic plan creation
+explain (costs off) execute prunelock_q(1);
+                  QUERY PLAN                  
+----------------------------------------------
+ Append
+   Subplans Removed: 2
+   ->  Seq Scan on prunelock_p1 prunelock_p_1
+         Filter: (a = $1)
+(4 rows)
+
+-- Execute and check which child partitions are locked
+begin;
+execute prunelock_q(1);
+ a 
+---
+(0 rows)
+
+select c.relname
+  from pg_locks l
+  join pg_class c on c.oid = l.relation
+ where l.pid = pg_backend_pid()
+   and c.relname like 'prunelock_p_'
+ order by c.relname;
+   relname    
+--------------
+ prunelock_p1
+(1 row)
+
+commit;
+deallocate prunelock_q;
+-- Turn pruning off
+set enable_partition_pruning to off;
+prepare prunelock_q (int) as select * from prunelock_p where a = $1;
+-- Force generic plan creation
+explain (costs off) execute prunelock_q(1);
+                  QUERY PLAN                  
+----------------------------------------------
+ Append
+   ->  Seq Scan on prunelock_p1 prunelock_p_1
+         Filter: (a = $1)
+   ->  Seq Scan on prunelock_p2 prunelock_p_2
+         Filter: (a = $1)
+   ->  Seq Scan on prunelock_p3 prunelock_p_3
+         Filter: (a = $1)
+(7 rows)
+
+-- Execute and check which child partitions are locked
+begin;
+execute prunelock_q(1);
+ a 
+---
+(0 rows)
+
+select c.relname
+  from pg_locks l
+  join pg_class c on c.oid = l.relation
+ where l.pid = pg_backend_pid()
+   and c.relname like 'prunelock_p_'
+ order by c.relname;
+   relname    
+--------------
+ prunelock_p1
+ prunelock_p2
+ prunelock_p3
+(3 rows)
+
+commit;
+deallocate prunelock_q;
+reset enable_partition_pruning;
+--
+-- Verify firstResultRels handling with multiple ModifyTable nodes
+-- (writable CTEs) targeting a partitioned table.  When a pruning
+-- parameter matches no partition, all result relations are pruned
+-- and the executor must still find a usable first result relation
+-- for each ModifyTable node.
+--
+prepare prunelock_mt_q (int, int) as
+  with upd1 as (update prunelock_p set a = a),
+       upd2 as (update prunelock_p set a = a where a = $2)
+  update prunelock_p set a = a where a = $1;
+-- Force generic plan creation
+explain (costs off) execute prunelock_mt_q(1, 2);
+                         QUERY PLAN                         
+------------------------------------------------------------
+ Update on prunelock_p
+   Update on prunelock_p1 prunelock_p_1
+   CTE upd1
+     ->  Update on prunelock_p prunelock_p_3
+           Update on prunelock_p1 prunelock_p_4
+           Update on prunelock_p2 prunelock_p_5
+           Update on prunelock_p3 prunelock_p_6
+           ->  Append
+                 ->  Seq Scan on prunelock_p1 prunelock_p_4
+                 ->  Seq Scan on prunelock_p2 prunelock_p_5
+                 ->  Seq Scan on prunelock_p3 prunelock_p_6
+   CTE upd2
+     ->  Update on prunelock_p prunelock_p_7
+           Update on prunelock_p2 prunelock_p_8
+           ->  Append
+                 Subplans Removed: 2
+                 ->  Seq Scan on prunelock_p2 prunelock_p_8
+                       Filter: (a = $2)
+   ->  Append
+         Subplans Removed: 2
+         ->  Seq Scan on prunelock_p1 prunelock_p_1
+               Filter: (a = $1)
+(22 rows)
+
+-- All partitions pruned: value 4 matches no partition, so each
+-- ModifyTable must still initialize correctly with no matching
+-- result relations.
+explain (costs off) execute prunelock_mt_q(4, 5);
+                         QUERY PLAN                         
+------------------------------------------------------------
+ Update on prunelock_p
+   CTE upd1
+     ->  Update on prunelock_p prunelock_p_2
+           Update on prunelock_p1 prunelock_p_3
+           Update on prunelock_p2 prunelock_p_4
+           Update on prunelock_p3 prunelock_p_5
+           ->  Append
+                 ->  Seq Scan on prunelock_p1 prunelock_p_3
+                 ->  Seq Scan on prunelock_p2 prunelock_p_4
+                 ->  Seq Scan on prunelock_p3 prunelock_p_5
+   CTE upd2
+     ->  Update on prunelock_p prunelock_p_6
+           ->  Append
+                 Subplans Removed: 3
+   ->  Append
+         Subplans Removed: 3
+(16 rows)
+
+deallocate prunelock_mt_q;
+drop table prunelock_p;
+--
+-- Verify that pruning-aware locking falls back to locking all
+-- partitions for multi-statement CachedPlans.  Rule rewriting can
+-- expand a single statement into multiple PlannedStmts, and later
+-- statements must not have their pruning evaluated before earlier
+-- ones have executed, since CCI between statements can change what
+-- pruning expressions see.
+--
+create table prune_config (val int);
+insert into prune_config values (1);
+create table multistmt_pt (a int, b int) partition by list (a);
+create table multistmt_pt_1 partition of multistmt_pt for values in (1);
+create table multistmt_pt_2 partition of multistmt_pt for values in (2);
+insert into multistmt_pt values (1, 0), (2, 0);
+create function get_prune_val() returns int as $$
+  select val from prune_config;
+$$ language sql stable;
+create rule config_upd_rule as on update to multistmt_pt
+  do also update prune_config set val = 2;
+set plan_cache_mode to force_generic_plan;
+prepare multi_q as update multistmt_pt set b = b + 1 where a = get_prune_val();
+-- first execute creates the generic plan
+execute multi_q;
+-- reset for the real test
+update prune_config set val = 1;
+update multistmt_pt set b = 0;
+-- second execute reuses the plan; pruning-aware locking kicks in
+execute multi_q;
+select * from multistmt_pt order by a;
+ a | b 
+---+---
+ 1 | 0
+ 2 | 1
+(2 rows)
+
+deallocate multi_q;
+drop rule config_upd_rule on multistmt_pt;
+drop function get_prune_val;
+drop table multistmt_pt, prune_config;
+reset plan_cache_mode;
diff --git a/src/test/regress/expected/plancache.out b/src/test/regress/expected/plancache.out
index 4e59188196c..3043dbfac2d 100644
--- a/src/test/regress/expected/plancache.out
+++ b/src/test/regress/expected/plancache.out
@@ -398,3 +398,66 @@ select name, generic_plans, custom_plans from pg_prepared_statements
 (1 row)
 
 drop table test_mode;
+-- This exercises the CachedPlanPrepCleanup() path, which must free
+-- the EState created by ExecutorPrep() when the plan is invalidated
+-- before execution begins.  The pruning expression uses a stable SQL
+-- function that calls a volatile plpgsql function.  That function
+-- performs DDL on a partition when a separate "signal" table says to
+-- do so.  The second EXECUTE should replan cleanly after the DDL.
+set plan_cache_mode to force_generic_plan;
+create table inval_during_pruning_p (a int) partition by list (a);
+create table inval_during_pruning_p1 partition of inval_during_pruning_p for values in (1);
+create table inval_during_pruning_p2 partition of inval_during_pruning_p for values in (2);
+insert into inval_during_pruning_p values (1), (2);
+create table inval_during_pruning_signal (create_idx bool not null);
+insert into inval_during_pruning_signal values (false);
+create or replace function invalidate_plancache_func() returns int
+as $$
+declare
+	create_index bool;
+begin
+	-- Perform DDL on a partition if asked to
+	select create_idx into create_index from inval_during_pruning_signal for update;
+	if create_index = true then
+		raise notice 'creating index on partition inval_during_pruning_p1';
+		create index on inval_during_pruning_p1 (a);
+		update inval_during_pruning_signal set create_idx = false;
+	end if;
+	-- value that pruning will match against partition bounds
+	return 1;
+end;
+$$ language plpgsql volatile;
+create or replace function stable_pruning_val() returns int as $$
+	select invalidate_plancache_func();
+$$ language sql stable;
+prepare inval_during_pruning_q as select * from inval_during_pruning_p where a = stable_pruning_val();
+-- Build a generic plan and run pruning once, but don't set the signal
+-- for invalidate_plancache_func() to perform the DDL.
+explain (verbose, costs off) execute inval_during_pruning_q;
+                                QUERY PLAN                                 
+---------------------------------------------------------------------------
+ Append
+   Subplans Removed: 1
+   ->  Seq Scan on public.inval_during_pruning_p1 inval_during_pruning_p_1
+         Output: inval_during_pruning_p_1.a
+         Filter: (inval_during_pruning_p_1.a = stable_pruning_val())
+(5 rows)
+
+-- Reuse the generic plan.  Make invalidate_plancache_func() perform DDL
+-- during this execution, which should force replanning without errors.
+update inval_during_pruning_signal set create_idx = true;
+explain (verbose, costs off) execute inval_during_pruning_q;
+NOTICE:  creating index on partition inval_during_pruning_p1
+                                QUERY PLAN                                 
+---------------------------------------------------------------------------
+ Append
+   Subplans Removed: 1
+   ->  Seq Scan on public.inval_during_pruning_p1 inval_during_pruning_p_1
+         Output: inval_during_pruning_p_1.a
+         Filter: (inval_during_pruning_p_1.a = stable_pruning_val())
+(5 rows)
+
+deallocate inval_during_pruning_q;
+drop table inval_during_pruning_p, inval_during_pruning_signal;
+drop function invalidate_plancache_func, stable_pruning_val;
+reset plan_cache_mode;
diff --git a/src/test/regress/sql/partition_prune.sql b/src/test/regress/sql/partition_prune.sql
index d93c0c03bab..692415a8d9f 100644
--- a/src/test/regress/sql/partition_prune.sql
+++ b/src/test/regress/sql/partition_prune.sql
@@ -1447,3 +1447,119 @@ select min(a) over (partition by a order by a) from part_abc where a >= stable_o
 
 drop view part_abc_view;
 drop table part_abc;
+
+--
+-- Verify that pruning-aware locking skips pruned partitions
+-- when reusing a generic cached plan.
+--
+set plan_cache_mode to force_generic_plan;
+
+create table prunelock_p (a int) partition by list (a);
+create table prunelock_p1 partition of prunelock_p for values in (1);
+create table prunelock_p2 partition of prunelock_p for values in (2);
+create table prunelock_p3 partition of prunelock_p for values in (3);
+
+prepare prunelock_q (int) as select * from prunelock_p where a = $1;
+
+-- Force generic plan creation
+explain (costs off) execute prunelock_q(1);
+
+-- Execute and check which child partitions are locked
+begin;
+execute prunelock_q(1);
+
+select c.relname
+  from pg_locks l
+  join pg_class c on c.oid = l.relation
+ where l.pid = pg_backend_pid()
+   and c.relname like 'prunelock_p_'
+ order by c.relname;
+commit;
+
+deallocate prunelock_q;
+
+-- Turn pruning off
+set enable_partition_pruning to off;
+
+prepare prunelock_q (int) as select * from prunelock_p where a = $1;
+
+-- Force generic plan creation
+explain (costs off) execute prunelock_q(1);
+
+-- Execute and check which child partitions are locked
+begin;
+execute prunelock_q(1);
+
+select c.relname
+  from pg_locks l
+  join pg_class c on c.oid = l.relation
+ where l.pid = pg_backend_pid()
+   and c.relname like 'prunelock_p_'
+ order by c.relname;
+commit;
+
+deallocate prunelock_q;
+reset enable_partition_pruning;
+
+--
+-- Verify firstResultRels handling with multiple ModifyTable nodes
+-- (writable CTEs) targeting a partitioned table.  When a pruning
+-- parameter matches no partition, all result relations are pruned
+-- and the executor must still find a usable first result relation
+-- for each ModifyTable node.
+--
+prepare prunelock_mt_q (int, int) as
+  with upd1 as (update prunelock_p set a = a),
+       upd2 as (update prunelock_p set a = a where a = $2)
+  update prunelock_p set a = a where a = $1;
+
+-- Force generic plan creation
+explain (costs off) execute prunelock_mt_q(1, 2);
+
+-- All partitions pruned: value 4 matches no partition, so each
+-- ModifyTable must still initialize correctly with no matching
+-- result relations.
+explain (costs off) execute prunelock_mt_q(4, 5);
+
+deallocate prunelock_mt_q;
+drop table prunelock_p;
+
+--
+-- Verify that pruning-aware locking falls back to locking all
+-- partitions for multi-statement CachedPlans.  Rule rewriting can
+-- expand a single statement into multiple PlannedStmts, and later
+-- statements must not have their pruning evaluated before earlier
+-- ones have executed, since CCI between statements can change what
+-- pruning expressions see.
+--
+create table prune_config (val int);
+insert into prune_config values (1);
+
+create table multistmt_pt (a int, b int) partition by list (a);
+create table multistmt_pt_1 partition of multistmt_pt for values in (1);
+create table multistmt_pt_2 partition of multistmt_pt for values in (2);
+insert into multistmt_pt values (1, 0), (2, 0);
+
+create function get_prune_val() returns int as $$
+  select val from prune_config;
+$$ language sql stable;
+
+create rule config_upd_rule as on update to multistmt_pt
+  do also update prune_config set val = 2;
+
+set plan_cache_mode to force_generic_plan;
+prepare multi_q as update multistmt_pt set b = b + 1 where a = get_prune_val();
+-- first execute creates the generic plan
+execute multi_q;
+-- reset for the real test
+update prune_config set val = 1;
+update multistmt_pt set b = 0;
+-- second execute reuses the plan; pruning-aware locking kicks in
+execute multi_q;
+select * from multistmt_pt order by a;
+
+deallocate multi_q;
+drop rule config_upd_rule on multistmt_pt;
+drop function get_prune_val;
+drop table multistmt_pt, prune_config;
+reset plan_cache_mode;
diff --git a/src/test/regress/sql/plancache.sql b/src/test/regress/sql/plancache.sql
index 4b2f11dcc64..6a8b8787de6 100644
--- a/src/test/regress/sql/plancache.sql
+++ b/src/test/regress/sql/plancache.sql
@@ -223,3 +223,55 @@ select name, generic_plans, custom_plans from pg_prepared_statements
   where  name = 'test_mode_pp';
 
 drop table test_mode;
+
+-- This exercises the CachedPlanPrepCleanup() path, which must free
+-- the EState created by ExecutorPrep() when the plan is invalidated
+-- before execution begins.  The pruning expression uses a stable SQL
+-- function that calls a volatile plpgsql function.  That function
+-- performs DDL on a partition when a separate "signal" table says to
+-- do so.  The second EXECUTE should replan cleanly after the DDL.
+set plan_cache_mode to force_generic_plan;
+create table inval_during_pruning_p (a int) partition by list (a);
+create table inval_during_pruning_p1 partition of inval_during_pruning_p for values in (1);
+create table inval_during_pruning_p2 partition of inval_during_pruning_p for values in (2);
+insert into inval_during_pruning_p values (1), (2);
+
+create table inval_during_pruning_signal (create_idx bool not null);
+insert into inval_during_pruning_signal values (false);
+create or replace function invalidate_plancache_func() returns int
+as $$
+declare
+	create_index bool;
+begin
+	-- Perform DDL on a partition if asked to
+	select create_idx into create_index from inval_during_pruning_signal for update;
+	if create_index = true then
+		raise notice 'creating index on partition inval_during_pruning_p1';
+		create index on inval_during_pruning_p1 (a);
+		update inval_during_pruning_signal set create_idx = false;
+	end if;
+	-- value that pruning will match against partition bounds
+	return 1;
+end;
+$$ language plpgsql volatile;
+
+create or replace function stable_pruning_val() returns int as $$
+	select invalidate_plancache_func();
+$$ language sql stable;
+
+prepare inval_during_pruning_q as select * from inval_during_pruning_p where a = stable_pruning_val();
+
+-- Build a generic plan and run pruning once, but don't set the signal
+-- for invalidate_plancache_func() to perform the DDL.
+explain (verbose, costs off) execute inval_during_pruning_q;
+
+-- Reuse the generic plan.  Make invalidate_plancache_func() perform DDL
+-- during this execution, which should force replanning without errors.
+update inval_during_pruning_signal set create_idx = true;
+explain (verbose, costs off) execute inval_during_pruning_q;
+
+deallocate inval_during_pruning_q;
+drop table inval_during_pruning_p, inval_during_pruning_signal;
+drop function invalidate_plancache_func, stable_pruning_val;
+
+reset plan_cache_mode;
-- 
2.47.3

