From 836f0b63ced2546b594643043b7d0055ffaa7b66 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Tue, 10 Feb 2026 22:09:23 +0900
Subject: [PATCH v6 5/6] Make SQL function executor track ExecutorPrep state

Extend the SQL function executor to use the ExecutorPrep results
returned by GetCachedPlan().  init_execution_state() now passes a
CachedPlanPrepData to GetCachedPlan() and stores the per statement
ExecPrep pointers in the execution_state nodes.

At execution time, postquel_start() reparents the prep estate's
es_query_cxt under the function's subcontext so that prep state
follows the usual per call context hierarchy.

This allows SQL language functions to participate in the same
ExecutorPrep machinery as other plan cache users.

Add a regression test where rule rewrite expands a single UPDATE
into multiple PlannedStmts, exercising the SQL function plan cache
and the generic plan reuse path that now invokes ExecutorPrep.
---
 src/backend/executor/functions.c        | 29 +++++++++++++--
 src/test/regress/expected/plancache.out | 48 +++++++++++++++++++++++++
 src/test/regress/sql/plancache.sql      | 34 ++++++++++++++++++
 3 files changed, 109 insertions(+), 2 deletions(-)

diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 65dfae58dcf..c70e06d8886 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -72,6 +72,7 @@ typedef struct execution_state
 	bool		setsResult;		/* true if this query produces func's result */
 	bool		lazyEval;		/* true if should fetch one row at a time */
 	PlannedStmt *stmt;			/* plan for this query */
+	EState	   *prep_estate;	/* EState created in ExecutorPrep() for this plan */
 	QueryDesc  *qd;				/* null unless status == RUN */
 } execution_state;
 
@@ -657,6 +658,8 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	execution_state *lasttages = NULL;
 	int			nstmts;
 	ListCell   *lc;
+	CachedPlanPrepData cprep = {0};
+	ListCell   *prep_lc;
 
 	/*
 	 * Clean up after previous query, if there was one.
@@ -695,11 +698,20 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	 * CurrentResourceOwner will be the same when ShutdownSQLFunction runs.)
 	 */
 	fcache->cowner = CurrentResourceOwner;
+
+	/*
+	 * Have ExecutorPrep() allocate under fcache->fcontext.  The prep
+	 * EStates it creates will initially live there; postquel_start()
+	 * will later reparent their es_query_cxt into fcache->subcontext
+	 * when using them for execution.
+	 */
+	cprep.context = fcache->fcontext;
+	cprep.owner = fcache->cowner;
 	fcache->cplan = GetCachedPlan(plansource,
 								  fcache->paramLI,
 								  fcache->cowner,
 								  NULL,
-								  NULL);
+								  &cprep);
 
 	/*
 	 * If necessary, make esarray[] bigger to hold the needed state.
@@ -720,9 +732,11 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	/*
 	 * Build execution_state list to match the number of contained plans.
 	 */
+	prep_lc = list_head(cprep.prep_estates);
 	foreach(lc, fcache->cplan->stmt_list)
 	{
 		PlannedStmt *stmt = lfirst_node(PlannedStmt, lc);
+		EState *prep_estate = next_prep_estate(cprep.prep_estates, &prep_lc);
 		execution_state *newes;
 
 		/*
@@ -764,6 +778,7 @@ init_execution_state(SQLFunctionCachePtr fcache)
 		newes->setsResult = false;	/* might change below */
 		newes->lazyEval = false;	/* might change below */
 		newes->stmt = stmt;
+		newes->prep_estate = prep_estate;
 		newes->qd = NULL;
 
 		if (stmt->canSetTag)
@@ -1362,6 +1377,15 @@ postquel_start(execution_state *es, SQLFunctionCachePtr fcache)
 	else
 		dest = None_Receiver;
 
+	/*
+	 * Prep EStates were built under fcache->fcontext.  For execution,
+	 * make their es_query_cxt a child of fcache->subcontext so they
+	 * follow the usual per call lifetime.
+	 */
+	if (es->prep_estate)
+		MemoryContextSetParent(es->prep_estate->es_query_cxt,
+							   fcache->subcontext);
+
 	es->qd = CreateQueryDesc(es->stmt,
 							 fcache->func->src,
 							 GetActiveSnapshot(),
@@ -1370,7 +1394,7 @@ postquel_start(execution_state *es, SQLFunctionCachePtr fcache)
 							 fcache->paramLI,
 							 es->qd ? es->qd->queryEnv : NULL,
 							 0,
-							 NULL);
+							 es->prep_estate);
 
 	/* Utility commands don't need Executor. */
 	if (es->qd->operation != CMD_UTILITY)
@@ -1461,6 +1485,7 @@ postquel_end(execution_state *es, SQLFunctionCachePtr fcache)
 
 	FreeQueryDesc(es->qd);
 	es->qd = NULL;
+	es->prep_estate = NULL;
 
 	MemoryContextSwitchTo(oldcontext);
 
diff --git a/src/test/regress/expected/plancache.out b/src/test/regress/expected/plancache.out
index 1d69ab0a1c2..371673a6e96 100644
--- a/src/test/regress/expected/plancache.out
+++ b/src/test/regress/expected/plancache.out
@@ -459,4 +459,52 @@ NOTICE:  creating index on partition inval_during_pruning_p1
 drop table inval_during_pruning_p, inval_during_pruning_signal;
 drop function invalidate_plancache_func, stable_pruning_val;
 deallocate inval_during_pruning_q;
+-- exercise sql-function plan cache when rewrite expands a single statement
+-- into multiple planned statements. this forces cachedplan->stmt_list to
+-- contain more than one entry and checks that executor state for the first
+-- rewritten statement does not destroy state needed by the second one.
+set plan_cache_mode = force_generic_plan;
+create table sqlf_base(id int, val int) partition by list (id);
+create table sqlf_base_1 partition of sqlf_base for values in (1);
+create table sqlf_base_2 partition of sqlf_base for values in (2);
+create table sqlf_log(id int, note text);
+insert into sqlf_base values (1, 10);
+create rule sqlf_base_upd_log as
+on update to sqlf_base do also
+	insert into sqlf_log(id, note)
+	values (new.id, 'logged by rule');
+create or replace function sqlf_execprep_test(a int, v int)
+returns void
+language sql
+as $$
+	update sqlf_base set val = v where id = a;
+$$;
+select sqlf_execprep_test(1, 20);
+ sqlf_execprep_test 
+--------------------
+ 
+(1 row)
+
+select sqlf_execprep_test(1, 30);
+ sqlf_execprep_test 
+--------------------
+ 
+(1 row)
+
+select * from sqlf_base order by 1;
+ id | val 
+----+-----
+  1 |  30
+(1 row)
+
+select * from sqlf_log order by 1;
+ id |      note      
+----+----------------
+  1 | logged by rule
+  1 | logged by rule
+(2 rows)
+
+drop rule sqlf_base_upd_log on sqlf_base;
+drop table sqlf_base, sqlf_log;
+drop function sqlf_execprep_test;
 reset plan_cache_mode;
diff --git a/src/test/regress/sql/plancache.sql b/src/test/regress/sql/plancache.sql
index 139b4688fd6..b89c9ad69a4 100644
--- a/src/test/regress/sql/plancache.sql
+++ b/src/test/regress/sql/plancache.sql
@@ -273,4 +273,38 @@ drop table inval_during_pruning_p, inval_during_pruning_signal;
 drop function invalidate_plancache_func, stable_pruning_val;
 deallocate inval_during_pruning_q;
 
+-- exercise sql-function plan cache when rewrite expands a single statement
+-- into multiple planned statements. this forces cachedplan->stmt_list to
+-- contain more than one entry and checks that executor state for the first
+-- rewritten statement does not destroy state needed by the second one.
+
+set plan_cache_mode = force_generic_plan;
+
+create table sqlf_base(id int, val int) partition by list (id);
+create table sqlf_base_1 partition of sqlf_base for values in (1);
+create table sqlf_base_2 partition of sqlf_base for values in (2);
+create table sqlf_log(id int, note text);
+
+insert into sqlf_base values (1, 10);
+
+create rule sqlf_base_upd_log as
+on update to sqlf_base do also
+	insert into sqlf_log(id, note)
+	values (new.id, 'logged by rule');
+
+create or replace function sqlf_execprep_test(a int, v int)
+returns void
+language sql
+as $$
+	update sqlf_base set val = v where id = a;
+$$;
+
+select sqlf_execprep_test(1, 20);
+select sqlf_execprep_test(1, 30);
+select * from sqlf_base order by 1;
+select * from sqlf_log order by 1;
+
+drop rule sqlf_base_upd_log on sqlf_base;
+drop table sqlf_base, sqlf_log;
+drop function sqlf_execprep_test;
 reset plan_cache_mode;
-- 
2.47.3

