v3-0001-Explicitly-forbid-WAIT-FOR-inside-functions-and-p.patch

From 0d3b3a12357f65dc6ff0af54aabdaf88dbcd2a9b Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Thu, 9 Apr 2026 15:02:32 +0300
Subject: [PATCH v3] Explicitly forbid WAIT FOR inside functions and procedures

Previously we were relying on snapshot-based check to detect such cases.
However, it appears that when WAIT FOR is wrapped into a stored procedure
it could pass this check causing an error elsewhere.  This commit implments
an explicit isTopLevel check to reject WAIT FOR when called from within
a function or procedure.  The isTopLevel check catches this case early with
a clear error message, matching the pattern used by other utility commands
like VACUUM and REINDEX.  The snapshot check is retained for the remaining
case: transactions with isolation level higher than READ COMMITTED.

Also add a test for WAIT FOR LSN wrapped in a procedure, complementing
the existing test that uses a function wrapper.

Reported-by: Satyanarayana Narlapuram <satyanarlapuram@gmail.com>
Discussion: https://postgr.es/m/CAHg%2BQDcN-n3NUqgRtj%3DBQb9fFQmH8-DeEROCr%3DPDbw_BBRKOYA%40mail.gmail.com
Author: Satyanarayana Narlapuram <satyanarlapuram@gmail.com>
Reviewed-by: Alexander Korotkov <aekorotkov@gmail.com>
---
 src/backend/commands/wait.c             | 15 +++++++++++++--
 src/backend/tcop/utility.c              |  3 ++-
 src/include/commands/wait.h             |  3 ++-
 src/test/recovery/t/049_wait_for_lsn.pl | 21 +++++++++++++++++++--
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/backend/commands/wait.c b/src/backend/commands/wait.c
index 85fcd463b4c..ec132e125a5 100644
--- a/src/backend/commands/wait.c
+++ b/src/backend/commands/wait.c
@@ -31,7 +31,8 @@
 
 
 void
-ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest)
+ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, bool isTopLevel,
+			 DestReceiver *dest)
 {
 	XLogRecPtr	lsn;
 	int64		timeout = 0;
@@ -135,6 +136,16 @@ ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest)
 		}
 	}
 
+	/*
+	 * WAIT FOR must not run inside a function or procedure. Forbid this case
+	 * upfront.
+	 */
+	if (!isTopLevel)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("%s cannot be executed from a function or procedure",
+						"WAIT FOR")));
+
 	/*
 	 * We are going to wait for the LSN.  We should first care that we don't
 	 * hold a snapshot and correspondingly our MyProc->xmin is invalid.
@@ -161,7 +172,7 @@ ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest)
 		ereport(ERROR,
 				errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 				errmsg("WAIT FOR must be called without an active or registered snapshot"),
-				errdetail("WAIT FOR cannot be executed from a function or procedure, nor within a transaction with an isolation level higher than READ COMMITTED."));
+				errdetail("WAIT FOR cannot be executed within a transaction with an isolation level higher than READ COMMITTED."));
 
 	/*
 	 * As the result we should hold no snapshot, and correspondingly our xmin
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
index 1d34c19913e..73a56f1df1d 100644
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -1062,7 +1062,8 @@ standard_ProcessUtility(PlannedStmt *pstmt,
 
 		case T_WaitStmt:
 			{
-				ExecWaitStmt(pstate, (WaitStmt *) parsetree, dest);
+				ExecWaitStmt(pstate, (WaitStmt *) parsetree, isTopLevel,
+							 dest);
 			}
 			break;
 
diff --git a/src/include/commands/wait.h b/src/include/commands/wait.h
index 521a312908d..a563579695c 100644
--- a/src/include/commands/wait.h
+++ b/src/include/commands/wait.h
@@ -16,7 +16,8 @@
 #include "parser/parse_node.h"
 #include "tcop/dest.h"
 
-extern void ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest);
+extern void ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, bool isTopLevel,
+						 DestReceiver *dest);
 extern TupleDesc WaitStmtResultDesc(WaitStmt *stmt);
 
 #endif							/* WAIT_H */
diff --git a/src/test/recovery/t/049_wait_for_lsn.pl b/src/test/recovery/t/049_wait_for_lsn.pl
index bf61b8c47cf..e108301ccdb 100644
--- a/src/test/recovery/t/049_wait_for_lsn.pl
+++ b/src/test/recovery/t/049_wait_for_lsn.pl
@@ -215,10 +215,27 @@ $node_standby->psql(
 	'postgres',
 	"SELECT pg_wal_replay_wait_wrap('${lsn3}');",
 	stderr => \$stderr);
-ok( $stderr =~
-	  /WAIT FOR must be called without an active or registered snapshot/,
+ok($stderr =~ /WAIT FOR cannot be executed from a function or procedure/,
 	"get an error when running within another function");
 
+$node_primary->safe_psql(
+	'postgres', qq[
+CREATE PROCEDURE pg_wal_replay_wait_proc(target_lsn pg_lsn) AS \$\$
+  BEGIN
+    EXECUTE format('WAIT FOR LSN %L;', target_lsn);
+  END
+\$\$
+LANGUAGE plpgsql;
+]);
+
+$node_primary->wait_for_catchup($node_standby);
+$node_standby->psql(
+	'postgres',
+	"CALL pg_wal_replay_wait_proc('${lsn3}');",
+	stderr => \$stderr);
+ok($stderr =~ /WAIT FOR cannot be executed from a function or procedure/,
+	"get an error when running within a procedure");
+
 # 6. Check parameter validation error cases on standby before promotion
 my $test_lsn =
   $node_primary->safe_psql('postgres', "SELECT pg_current_wal_insert_lsn()");
-- 
2.39.5 (Apple Git-154)