From 8fa533bf528e420ec322ec17000617c75e2638a0 Mon Sep 17 00:00:00 2001
From: Sami Imseih <samimseih.pg@gmail.com>
Date: Sat, 29 Aug 2026 19:40:14 +0000
Subject: [PATCH v2 2/4] Normalize WAIT FOR WITH clause values in
 pg_stat_statements

Teach WaitStmt option jumbling to record DefElem argument
locations, so MODE and TIMEOUT values normalize to one
pg_stat_statements entry.

Introduce JumbleDefElemOptions() as a small helper for
other statements that need the same DefElem option
jumbling.
---
 .../pg_stat_statements/expected/utility.out   | 26 +++++++++++++++++
 contrib/pg_stat_statements/sql/utility.sql    |  6 ++++
 src/backend/nodes/queryjumblefuncs.c          | 28 +++++++++++++++++++
 src/backend/parser/gram.y                     |  1 +
 src/include/nodes/parsenodes.h                | 11 ++++++--
 5 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index e4d6564ea5b..f2e53e47272 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -719,6 +719,32 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t
 (1 row)
 
+-- WAIT FOR
+WAIT FOR LSN 'FFFFFFFF/FFFFFFFF' WITH (MODE 'primary_flush', TIMEOUT '1ms', NO_THROW);
+ status  
+---------
+ timeout
+(1 row)
+
+WAIT FOR LSN 'FFFFFFFE/FFFFFFFF' WITH (MODE 'primary_flush', TIMEOUT '2ms', NO_THROW);
+ status  
+---------
+ timeout
+(1 row)
+
+SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+ calls | rows |                        query                         
+-------+------+------------------------------------------------------
+     1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+     2 |    0 | WAIT FOR LSN $1 WITH (MODE $2, TIMEOUT $3, NO_THROW)
+(2 rows)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
 -- Special cases.  Keep these ones at the end to avoid conflicts.
 SET SCHEMA 'foo';
 SET SCHEMA 'public';
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index dd97203c210..fb843fffa6c 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -365,6 +365,12 @@ DROP TABLE pgss_select_into;
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 
+-- WAIT FOR
+WAIT FOR LSN 'FFFFFFFF/FFFFFFFF' WITH (MODE 'primary_flush', TIMEOUT '1ms', NO_THROW);
+WAIT FOR LSN 'FFFFFFFE/FFFFFFFF' WITH (MODE 'primary_flush', TIMEOUT '2ms', NO_THROW);
+SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
 -- Special cases.  Keep these ones at the end to avoid conflicts.
 SET SCHEMA 'foo';
 SET SCHEMA 'public';
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 094d2872644..53b08285fcf 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -70,6 +70,7 @@ static void FlushPendingNulls(JumbleState *jstate);
 static void RecordConstLocation(JumbleState *jstate,
 								bool extern_param,
 								int location, int len);
+static void JumbleDefElemOptions(JumbleState *jstate, List *options);
 static void _jumbleNode(JumbleState *jstate, Node *node);
 static void _jumbleList(JumbleState *jstate, Node *node);
 static void _jumbleElements(JumbleState *jstate, List *elements, Node *node);
@@ -79,6 +80,9 @@ static void _jumbleVariableSetStmt(JumbleState *jstate, Node *node);
 static void _jumbleRangeTblEntry_eref(JumbleState *jstate,
 									  RangeTblEntry *rte,
 									  Alias *expr);
+static void _jumbleWaitStmt_options(JumbleState *jstate,
+									WaitStmt *ws,
+									List *options);
 
 /*
  * Given a possibly multi-statement source string, confine our attention to the
@@ -774,6 +778,30 @@ _jumbleRangeTblEntry_eref(JumbleState *jstate,
 	JUMBLE_STRING(aliasname);
 }
 
+static void
+JumbleDefElemOptions(JumbleState *jstate, List *options)
+{
+	ListCell   *l;
+
+	foreach(l, options)
+	{
+		DefElem    *expr = lfirst_node(DefElem, l);
+
+		JUMBLE_STRING(defnamespace);
+		JUMBLE_STRING(defname);
+		JUMBLE_FIELD(defaction);
+		RecordConstLocation(jstate, false, expr->arg_location, -1);
+	}
+}
+
+static void
+_jumbleWaitStmt_options(JumbleState *jstate,
+						WaitStmt *ws,
+						List *options)
+{
+	JumbleDefElemOptions(jstate, options);
+}
+
 /*
  * CompLocation: comparator for qsorting LocationLen structs by location
  */
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 65b2b4384f8..bff0b2a46b0 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -17300,6 +17300,7 @@ WaitStmt:
 				{
 					WaitStmt *n = makeNode(WaitStmt);
 					n->lsn_literal = $4;
+					n->lsn_location = @4;
 					n->options = $5;
 					$$ = (Node *) n;
 				}
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index afe09b1f7a9..86eb9512d08 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -4608,8 +4608,15 @@ typedef struct DropSubscriptionStmt
 typedef struct WaitStmt
 {
 	NodeTag		type;
-	char	   *lsn_literal;	/* LSN string from grammar */
-	List	   *options;		/* List of DefElem nodes */
+
+	/* LSN string from grammar */
+	char	   *lsn_literal pg_node_attr(query_jumble_ignore);
+
+	/* List of DefElem nodes */
+	List	   *options pg_node_attr(custom_query_jumble);
+
+	/* LSN literal token location, or -1 if unknown */
+	ParseLoc	lsn_location pg_node_attr(query_jumble_location);
 } WaitStmt;
 
 
-- 
2.47.3

