From f4715162978951eb4513b6963b5cc7cd24d5a5d9 Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Wed, 30 Apr 2025 15:46:58 -0500
Subject: [PATCH v1 1/1] Allow query jumble to squash a list external
 parameters

62d712ecf now allows query jumbling to squash a list of constants,
but not constants that are passed as external parameters. This patch
now allows the squashing of constant values supplied as external parameters
(e.g., $1, $2), as is the case with prepared statements.
---
 .../pg_stat_statements/expected/squashing.out | 33 +++++++++++++++++++
 contrib/pg_stat_statements/sql/squashing.sql  | 11 +++++++
 src/backend/nodes/queryjumblefuncs.c          | 20 ++++++++---
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/squashing.out b/contrib/pg_stat_statements/expected/squashing.out
index 7b138af098c..8dc98bad6d5 100644
--- a/contrib/pg_stat_statements/expected/squashing.out
+++ b/contrib/pg_stat_statements/expected/squashing.out
@@ -301,6 +301,39 @@ SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
  SELECT pg_stat_statements_reset() IS NOT NULL AS t                   |     1
 (2 rows)
 
+-- Test bind parameters
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SELECT * FROM test_squash_bigint WHERE data IN ($1, $2, $3) \bind 1 2 3
+;
+ id | data 
+----+------
+(0 rows)
+
+SELECT * FROM test_squash_bigint WHERE data IN ($1, $2, $3, $4) \bind 1 2 3 4
+;
+ id | data 
+----+------
+(0 rows)
+
+SELECT * FROM test_squash_bigint WHERE data IN
+	($1::bigint, $2::bigint, $3::bigint, $4::bigint) \bind 1 2 3 4
+;
+ id | data 
+----+------
+(0 rows)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+                             query                              | calls 
+----------------------------------------------------------------+-------
+ SELECT * FROM test_squash_bigint WHERE data IN ($1 /*, ... */) |     3
+ SELECT pg_stat_statements_reset() IS NOT NULL AS t             |     1
+(2 rows)
+
 -- CoerceViaIO
 -- Create some dummy type to force CoerceViaIO
 CREATE TYPE casttesttype;
diff --git a/contrib/pg_stat_statements/sql/squashing.sql b/contrib/pg_stat_statements/sql/squashing.sql
index 908be81ff2b..ce0bcbc4121 100644
--- a/contrib/pg_stat_statements/sql/squashing.sql
+++ b/contrib/pg_stat_statements/sql/squashing.sql
@@ -97,6 +97,17 @@ SELECT * FROM test_squash_jsonb WHERE data IN
 	 (SELECT '"10"')::jsonb);
 SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+-- Test bind parameters
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+SELECT * FROM test_squash_bigint WHERE data IN ($1, $2, $3) \bind 1 2 3
+;
+SELECT * FROM test_squash_bigint WHERE data IN ($1, $2, $3, $4) \bind 1 2 3 4
+;
+SELECT * FROM test_squash_bigint WHERE data IN
+	($1::bigint, $2::bigint, $3::bigint, $4::bigint) \bind 1 2 3 4
+;
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+
 -- CoerceViaIO
 
 -- Create some dummy type to force CoerceViaIO
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index d1e82a63f09..7468583edc8 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -410,7 +410,8 @@ RecordConstLocation(JumbleState *jstate, int location, bool squashed)
  * - Ignore a possible wrapping RelabelType and CoerceViaIO.
  * - If it's a FuncExpr, check that the function is an implicit
  *   cast and its arguments are Const.
- * - Otherwise test if the expression is a simple Const.
+ * - Otherwise test if the expression is a simple Const or an
+ *   external parameter.
  */
 static bool
 IsSquashableConst(Node *element)
@@ -444,10 +445,21 @@ IsSquashableConst(Node *element)
 		return true;
 	}
 
-	if (!IsA(element, Const))
-		return false;
+	switch (nodeTag(element))
+	{
+		case T_Const:
+			return true;
+		case T_Param:
+			{
+				Param      *param = (Param *) element;
 
-	return true;
+				return param->paramkind == PARAM_EXTERN;
+			}
+		default:
+			break;
+	}
+
+	return false;
 }
 
 /*
-- 
2.39.5 (Apple Git-154)