v4-0001-Explicitly-forbid-non-top-level-WAIT-FOR-executio.patch

From c98351d28b158fc0b34e60a16639d35e7f358f72 Mon Sep 17 00:00:00 2001
From: alterego655 <824662526@qq.com>
Date: Fri, 10 Apr 2026 14:46:32 +0800
Subject: [PATCH v4] Explicitly forbid non-top-level WAIT FOR execution

Previously we were relying on a snapshot-based check to detect invalid
execution contexts.  However, when WAIT FOR is wrapped into a stored
procedure or a DO block, it could pass this check, causing an error
elsewhere.

This commit implements an explicit isTopLevel check to reject WAIT FOR
when called from within a function, procedure, or DO block.  The
isTopLevel check catches these cases 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:
top-level execution within a transaction block using an isolation level
higher than READ COMMITTED.

Also adds tests for WAIT FOR LSN wrapped in a procedure and DO block,
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>
---
 doc/src/sgml/ref/wait_for.sgml          |  8 +++++++
 src/backend/commands/wait.c             | 22 +++++++++++++-----
 src/backend/tcop/utility.c              |  3 ++-
 src/include/commands/wait.h             |  3 ++-
 src/test/recovery/t/049_wait_for_lsn.pl | 30 ++++++++++++++++++++++---
 5 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/doc/src/sgml/ref/wait_for.sgml b/doc/src/sgml/ref/wait_for.sgml
index c30fba6f05a..9ba785ea321 100644
--- a/doc/src/sgml/ref/wait_for.sgml
+++ b/doc/src/sgml/ref/wait_for.sgml
@@ -221,6 +221,14 @@ WAIT FOR LSN '<replaceable class="parameter">lsn</replaceable>'
 
  <refsect1>
   <title>Notes</title>
+  <para>
+   <command>WAIT FOR</command> must be executed as a top-level command.
+   It cannot be executed from a function, procedure, or
+   <command>DO</command> block. It also requires that no active or
+   registered snapshot be held, and therefore cannot be used in contexts
+   where such a snapshot must remain active, including transactions running
+   at isolation levels higher than <literal>READ COMMITTED</literal>.
+  </para>
 
   <para>
    <command>WAIT FOR</command> waits until the specified
diff --git a/src/backend/commands/wait.c b/src/backend/commands/wait.c
index 85fcd463b4c..12888281d9b 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;
@@ -45,6 +46,16 @@ ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest)
 	bool		no_throw_specified = false;
 	bool		mode_specified = false;
 
+	/*
+	 * WAIT FOR must not be run as a non-top-level statement (e.g., inside a
+	 * function, procedure, or DO block). Forbid this case upfront.
+	 */
+	if (!isTopLevel)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("%s can only be executed as a top-level statement",
+						"WAIT FOR")));
+
 	/* Parse and validate the mandatory LSN */
 	lsn = DatumGetLSN(DirectFunctionCall1(pg_lsn_in,
 										  CStringGetDatum(stmt->lsn_literal)));
@@ -142,10 +153,9 @@ ExecWaitStmt(ParseState *pstate, WaitStmt *stmt, DestReceiver *dest)
 	 * implying a kind of self-deadlock.  This is the reason why WAIT FOR is a
 	 * command, not a procedure or function.
 	 *
-	 * At first, we should check there is no active snapshot.  According to
-	 * PlannedStmtRequiresSnapshot(), even in an atomic context, CallStmt is
-	 * processed with a snapshot.  Thankfully, we can pop this snapshot,
-	 * because PortalRunUtility() can tolerate this.
+	 * Non-top-level contexts are rejected above, but be defensive and pop
+	 * any active snapshot if one is present.  PortalRunUtility() can
+	 * tolerate utility commands that remove the active snapshot.
 	 */
 	if (ActiveSnapshotSet())
 		PopActiveSnapshot();
@@ -161,7 +171,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..8358c57f7b7 100644
--- a/src/test/recovery/t/049_wait_for_lsn.pl
+++ b/src/test/recovery/t/049_wait_for_lsn.pl
@@ -215,9 +215,33 @@ $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/,
-	"get an error when running within another function");
+ok($stderr =~ /WAIT FOR can only be executed as a top-level statement/,
+	"get an error when running within a 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 can only be executed as a top-level statement/,
+	"get an error when running within a procedure");
+
+$node_standby->psql(
+	'postgres',
+	"DO \$\$ BEGIN EXECUTE format('WAIT FOR LSN %L;', '${lsn3}'); END \$\$;",
+	stderr => \$stderr);
+ok($stderr =~ /WAIT FOR can only be executed as a top-level statement/,
+	"get an error when running within a DO block");
 
 # 6. Check parameter validation error cases on standby before promotion
 my $test_lsn =
-- 
2.51.0