From 76afc3c74f222d3bf5551e7b45cf736453ae91c9 Mon Sep 17 00:00:00 2001
From: "Sami Imseih (AWS)"
 <simseih@dev-dsk-simseih-1d-3940b79e.us-east-1.amazon.com>
Date: Fri, 21 Mar 2025 05:37:59 +0000
Subject: [PATCH 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 | 38 +++++++++++++++++++
 contrib/pg_stat_statements/sql/squashing.sql  | 12 ++++++
 src/backend/nodes/queryjumblefuncs.c          | 20 ++++++++--
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/squashing.out b/contrib/pg_stat_statements/expected/squashing.out
index 55aa5109433..370d91642d4 100644
--- a/contrib/pg_stat_statements/expected/squashing.out
+++ b/contrib/pg_stat_statements/expected/squashing.out
@@ -333,6 +333,44 @@ 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 * FROM test_squash_bigint WHERE data IN (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 /*, ... */) |     4
+ 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 56ee8ccb9a1..3ff251a59ee 100644
--- a/contrib/pg_stat_statements/sql/squashing.sql
+++ b/contrib/pg_stat_statements/sql/squashing.sql
@@ -106,6 +106,18 @@ 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 * FROM test_squash_bigint WHERE data IN (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 189bfda610a..de405ab08f3 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -244,7 +244,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)
@@ -278,10 +279,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)