From 12575c2db5ccfc070b4024d644e4c72e32497486 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-46-230.ec2.internal>
Date: Tue, 6 Jan 2026 19:29:01 +0000
Subject: [PATCH v4 2/2] Make JumbleState const in post_parse_analyze hook

Change the post_parse_analyze_hook_type signature to take a const
JumbleState parameter, preventing hooks from modifying the jumble
state during query analysis. This improves API safety by making it
clear that hooks should only read from the jumble state, not modify
it.

Update pg_stat_statements and related functions to match the new const
signature. Refactor SetConstantLengths() to return a newly allocated
LocationLen array instead of modifying the JumbleState, and update
GenerateNormalizedQuery() to work with the const JumbleState and
manage the returned array properly.

Discussion: https://postgr.es/m/CAA5RZ0tZp5qU0ikZEEqJnxvdSNGh1DWv80sb-k4QAUmiMoOp_Q@mail.gmail.com
---
 .../pg_stat_statements/pg_stat_statements.c   |  8 ++--
 src/backend/nodes/queryjumblefuncs.c          | 39 +++++++++++++------
 src/include/nodes/queryjumble.h               |  2 +-
 src/include/parser/analyze.h                  |  2 +-
 4 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index b90891859d9..b796d839cec 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -332,7 +332,7 @@ static void pgss_shmem_request(void);
 static void pgss_shmem_startup(void);
 static void pgss_shmem_shutdown(int code, Datum arg);
 static void pgss_post_parse_analyze(ParseState *pstate, Query *query,
-									JumbleState *jstate);
+									const JumbleState *jstate);
 static PlannedStmt *pgss_planner(Query *parse,
 								 const char *query_string,
 								 int cursorOptions,
@@ -356,7 +356,7 @@ static void pgss_store(const char *query, int64 queryId,
 					   const BufferUsage *bufusage,
 					   const WalUsage *walusage,
 					   const struct JitInstrumentation *jitusage,
-					   JumbleState *jstate,
+					   const JumbleState *jstate,
 					   int parallel_workers_to_launch,
 					   int parallel_workers_launched,
 					   PlannedStmtOrigin planOrigin);
@@ -832,7 +832,7 @@ error:
  * Post-parse-analysis hook: mark query with a queryId
  */
 static void
-pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
+pgss_post_parse_analyze(ParseState *pstate, Query *query, const JumbleState *jstate)
 {
 	if (prev_post_parse_analyze_hook)
 		prev_post_parse_analyze_hook(pstate, query, jstate);
@@ -1289,7 +1289,7 @@ pgss_store(const char *query, int64 queryId,
 		   const BufferUsage *bufusage,
 		   const WalUsage *walusage,
 		   const struct JitInstrumentation *jitusage,
-		   JumbleState *jstate,
+		   const JumbleState *jstate,
 		   int parallel_workers_to_launch,
 		   int parallel_workers_launched,
 		   PlannedStmtOrigin planOrigin)
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index dd53c82b2a8..4dfe9087547 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -789,8 +789,10 @@ CompLocation(const void *a, const void *b)
 }
 
 /*
- * Given a valid SQL string and an array of constant-location records,
- * fill in the textual lengths of those constants in JumbleState.
+ * Given a valid SQL string and an array of constant-location records, return
+ * the textual lengths of those constants in a newly allocated LocationLen
+ * array, or NULL if there are no constants. It is the caller's responsibility
+ * to pfree the result, if necessary.
  *
  * The constants may use any allowed constant syntax, such as float literals,
  * bit-strings, single-quoted strings and dollar-quoted strings.  This is
@@ -803,8 +805,8 @@ CompLocation(const void *a, const void *b)
  * a negative numeric constant.  This precludes there ever being another
  * reason for a constant to start with a '-'.
  */
-static void
-SetConstantLengths(JumbleState *jstate, const char *query,
+static LocationLen *
+SetConstantLengths(const JumbleState *jstate, const char *query,
 				   int query_loc)
 {
 	LocationLen *locs;
@@ -813,14 +815,20 @@ SetConstantLengths(JumbleState *jstate, const char *query,
 	core_YYSTYPE yylval;
 	YYLTYPE		yylloc;
 
+	if (jstate->clocations_count == 0)
+		return NULL;
+
+	/* Copy constant locations to avoid modifying jstate */
+	locs = palloc(jstate->clocations_count * sizeof(LocationLen));
+	memcpy(locs, jstate->clocations, jstate->clocations_count * sizeof(LocationLen));
+
 	/*
 	 * Sort the records by location so that we can process them in order while
 	 * scanning the query text.
 	 */
 	if (jstate->clocations_count > 1)
-		qsort(jstate->clocations, jstate->clocations_count,
+		qsort(locs, jstate->clocations_count,
 			  sizeof(LocationLen), CompLocation);
-	locs = jstate->clocations;
 
 	/* initialize the flex scanner --- should match raw_parser() */
 	yyscanner = scanner_init(query,
@@ -908,6 +916,8 @@ SetConstantLengths(JumbleState *jstate, const char *query,
 	}
 
 	scanner_finish(yyscanner);
+
+	return locs;
 }
 
 /*
@@ -926,7 +936,7 @@ SetConstantLengths(JumbleState *jstate, const char *query,
  * Returns a palloc'd string.
  */
 char *
-GenerateNormalizedQuery(JumbleState *jstate, const char *query,
+GenerateNormalizedQuery(const JumbleState *jstate, const char *query,
 						int query_loc, int *query_len_p)
 {
 	char	   *norm_query;
@@ -938,12 +948,13 @@ GenerateNormalizedQuery(JumbleState *jstate, const char *query,
 				last_off = 0,	/* Offset from start for previous tok */
 				last_tok_len = 0;	/* Length (in bytes) of that tok */
 	int			num_constants_replaced = 0;
+	LocationLen *locs = NULL;
 
 	/*
 	 * Set constants' lengths in JumbleState, as only locations are set during
 	 * DoJumble(). Note this also ensures the items are sorted by location.
 	 */
-	SetConstantLengths(jstate, query, query_loc);
+	locs = SetConstantLengths(jstate, query, query_loc);
 
 	/*
 	 * Allow for $n symbols to be longer than the constants they replace.
@@ -969,15 +980,15 @@ GenerateNormalizedQuery(JumbleState *jstate, const char *query,
 		 * the parameter in the next iteration (or after the loop is done),
 		 * which is a bit odd but seems to work okay in most cases.
 		 */
-		if (jstate->clocations[i].extern_param && !jstate->has_squashed_lists)
+		if (locs[i].extern_param && !jstate->has_squashed_lists)
 			continue;
 
-		off = jstate->clocations[i].location;
+		off = locs[i].location;
 
 		/* Adjust the constant's location (see SetConstantLengths) */
 		off -= query_loc;
 
-		tok_len = jstate->clocations[i].length;
+		tok_len = locs[i].length;
 
 		if (tok_len < 0)
 			continue;			/* ignore any duplicates */
@@ -996,7 +1007,7 @@ GenerateNormalizedQuery(JumbleState *jstate, const char *query,
 		 */
 		n_quer_loc += sprintf(norm_query + n_quer_loc, "$%d%s",
 							  num_constants_replaced + 1 + jstate->highest_extern_param_id,
-							  jstate->clocations[i].squashed ? " /*, ... */" : "");
+							  locs[i].squashed ? " /*, ... */" : "");
 		num_constants_replaced++;
 
 		/* move forward */
@@ -1005,6 +1016,10 @@ GenerateNormalizedQuery(JumbleState *jstate, const char *query,
 		last_tok_len = tok_len;
 	}
 
+	/* Clean up temporary location array before any assertions */
+	if (locs)
+		pfree(locs);
+
 	/*
 	 * We've copied up until the last ignorable constant.  Copy over the
 	 * remaining bytes of the original query string.
diff --git a/src/include/nodes/queryjumble.h b/src/include/nodes/queryjumble.h
index cb8741b1d05..bc3dc4d846f 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -91,7 +91,7 @@ extern PGDLLIMPORT int compute_query_id;
 
 
 extern const char *CleanQuerytext(const char *query, int *location, int *len);
-extern char *GenerateNormalizedQuery(JumbleState *jstate, const char *query,
+extern char *GenerateNormalizedQuery(const JumbleState *jstate, const char *query,
 									 int query_loc, int *query_len_p);
 extern JumbleState *JumbleQuery(Query *query);
 extern void EnableQueryId(void);
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index f29ed03b476..f80da72933c 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -21,7 +21,7 @@
 /* Hook for plugins to get control at end of parse analysis */
 typedef void (*post_parse_analyze_hook_type) (ParseState *pstate,
 											  Query *query,
-											  JumbleState *jstate);
+											  const JumbleState *jstate);
 extern PGDLLIMPORT post_parse_analyze_hook_type post_parse_analyze_hook;
 
 
-- 
2.43.0

