From 42360a306bfed6b899bd03e60c69b31db59e2ec0 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 v7 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.

Also, rename SetConstantLengths() to ComputeConstantLengths() as it is
no longer setting constant lengths, and it's also computing lengths
based on traversing through the LocationLen data.

Furthermore, ComputeConstantLengths is also exported to allow plug-ins
that wish to perform custom processing of literals from a query.

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

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index c9338442d96..386a1419232 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -333,7 +333,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,
@@ -357,7 +357,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);
@@ -833,7 +833,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);
@@ -1290,7 +1290,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 c997a1ef609..af9417797a3 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
@@ -808,9 +810,9 @@ 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,
-				   int query_loc)
+LocationLen *
+ComputeConstantLengths(const JumbleState *jstate, const char *query,
+					   int query_loc)
 {
 	LocationLen *locs;
 	core_yyscan_t yyscanner;
@@ -818,14 +820,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_array(LocationLen, jstate->clocations_count);
+	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,
@@ -910,6 +918,8 @@ SetConstantLengths(JumbleState *jstate, const char *query,
 	}
 
 	scanner_finish(yyscanner);
+
+	return locs;
 }
 
 /*
@@ -920,7 +930,7 @@ SetConstantLengths(JumbleState *jstate, const char *query,
  * the result string length on exit.  The resulting string might be longer
  * or shorter depending on what happens with replacement of constants.
  *
- * Similar to SetConstantLengths, we must also translate the provided location
+ * Similar to ComputeConstantLengths, we must also translate the provided location
  * to compensate for multi-statement strings.
  *
  * It is the caller's job to ensure that the string is a valid SQL statement
@@ -931,7 +941,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;
@@ -943,12 +953,14 @@ 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.
+	 * Determine constants' lengths, which are not set during DoJumble(), and
+	 * return a sorted copy of jstate's LocationLen data with lengths filled
+	 * in.
 	 */
-	SetConstantLengths(jstate, query, query_loc);
+	locs = ComputeConstantLengths(jstate, query, query_loc);
 
 	/*
 	 * Allow for $n symbols to be longer than the constants they replace.
@@ -974,15 +986,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) */
+		/* Adjust the constant's location (see ComputeConstantLengths) */
 		off -= query_loc;
 
-		tok_len = jstate->clocations[i].length;
+		tok_len = locs[i].length;
 
 		if (tok_len < 0)
 			continue;			/* ignore any duplicates */
@@ -1001,7 +1013,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 */
@@ -1010,6 +1022,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 e354f857c37..8364d58e4cb 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -91,8 +91,10 @@ 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 LocationLen *ComputeConstantLengths(const JumbleState *jstate, const char *query,
+										   int query_loc);
 extern JumbleState *JumbleQuery(Query *query);
 extern void EnableQueryId(void);
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index e10270ff0ff..b2db6fa4e8c 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.47.3

