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.  Relevant
documentation paragraph is also added.

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>
Reviewed-by: Xuneng Zhou <xunengzhou@gmail.com>

diff --git a/doc/src/sgml/ref/wait_for.sgml b/doc/src/sgml/ref/wait_for.sgml
index c30fba6f05ac23abacc8b115af74d625d882a9ac..9ba785ea3217b94a770e7855c65c6d21c9f2e40e 100644 (file)
--- 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 85fcd463b4c59431ab0a0236a4335a58caada459..382d5c2d44ffbe1f440bb4edc601c326cd1d4d35 100644 (file)
--- 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,17 @@ 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"),
+                                errdetail("WAIT FOR cannot be used within a function, procedure, or DO block.")));
+
        /* Parse and validate the mandatory LSN */
        lsn = DatumGetLSN(DirectFunctionCall1(pg_lsn_in,
                                                                                  CStringGetDatum(stmt->lsn_literal)));
@@ -142,10 +154,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 +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 1d34c19913ee68a84449337c5a979859b6b2d2e1..73a56f1df1dc3c265497df3b74dd0b643bc6e73c 100644 (file)
--- 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 521a312908dee6fe189b9895d370190baa253fdb..a563579695c95f7e929609e65053ab7a4d04edea 100644 (file)
--- 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 bf61b8c47cf55fa05473c73db38938ff688dd56e..8358c57f7b7b83127db62758722839dc3fd73833 100644 (file)
--- 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 =