From 1c5b1a99aee4d0872d42b6edfdaab1266d13a522 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 21 Mar 2025 09:43:48 +0900
Subject: [PATCH v3 2/2] Add custom query jumble function for
 RangeTblEntry.relid

---
 src/include/nodes/parsenodes.h                |  9 +++--
 src/backend/nodes/queryjumblefuncs.c          | 27 +++++++++++++++
 .../pg_stat_statements/expected/utility.out   | 33 +++++++++++++++++++
 contrib/pg_stat_statements/sql/utility.sql    | 12 +++++++
 4 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 23c9e3c5abf2..3eb16846e0e1 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1093,8 +1093,13 @@ typedef struct RangeTblEntry
 	 * relation.  This allows plans referencing AFTER trigger transition
 	 * tables to be invalidated if the underlying table is altered.
 	 */
-	/* OID of the relation */
-	Oid			relid;
+
+	/*
+	 * OID of the relation.  A custom query jumble function is used here for
+	 * temporary tables, where the computation uses the relation name instead
+	 * of the OID.
+	 */
+	Oid			relid pg_node_attr(custom_query_jumble);
 	/* inheritance requested? */
 	bool		inh;
 	/* relation kind (see pg_class.relkind) */
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 9b9cf6f34381..fbb05eab1bbe 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -33,6 +33,7 @@
 #include "postgres.h"
 
 #include "access/transam.h"
+#include "catalog/namespace.h"
 #include "catalog/pg_proc.h"
 #include "common/hashfn.h"
 #include "miscadmin.h"
@@ -67,6 +68,9 @@ static void _jumbleElements(JumbleState *jstate, List *elements);
 static void _jumbleA_Const(JumbleState *jstate, Node *node);
 static void _jumbleList(JumbleState *jstate, Node *node);
 static void _jumbleVariableSetStmt(JumbleState *jstate, Node *node);
+static void _jumbleRangeTblEntry_relid(JumbleState *jstate,
+									   RangeTblEntry *expr,
+									   Oid relid);
 
 /*
  * Given a possibly multi-statement source string, confine our attention to the
@@ -519,3 +523,26 @@ _jumbleVariableSetStmt(JumbleState *jstate, Node *node)
 	JUMBLE_FIELD(is_local);
 	JUMBLE_LOCATION(location);
 }
+
+/*
+ * Custom query jumble function for RangeTblEntry.relid.
+ */
+static void
+_jumbleRangeTblEntry_relid(JumbleState *jstate,
+						   RangeTblEntry *expr,
+						   Oid relid)
+{
+	/*
+	 * If this is a temporary table, jumble its name instead of the table OID.
+	 */
+	if (expr->rtekind == RTE_RELATION &&
+		isAnyTempNamespace(get_rel_namespace(expr->relid)))
+	{
+		char	   *relname = get_rel_name(expr->relid);
+
+		AppendJumble(jstate, (const unsigned char *) "pg_temp", sizeof("pg_temp"));
+		AppendJumble(jstate, (const unsigned char *) relname, strlen(relname));
+	}
+	else
+		JUMBLE_FIELD(relid);
+}
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index aa4f0f7e6280..941ba0e66deb 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -9,6 +9,39 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t
 (1 row)
 
+-- Temporary tables, grouped together.
+BEGIN;
+  CREATE TEMP TABLE temp_t (id int) ON COMMIT DROP;
+  SELECT * FROM temp_t;
+ id 
+----
+(0 rows)
+
+COMMIT;
+BEGIN;
+  CREATE TEMP TABLE temp_t (id int) ON COMMIT DROP;
+  SELECT * FROM temp_t;
+ id 
+----
+(0 rows)
+
+COMMIT;
+SELECT calls, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+ calls |                       query                        
+-------+----------------------------------------------------
+     2 | BEGIN
+     2 | COMMIT
+     2 | CREATE TEMP TABLE temp_t (id int) ON COMMIT DROP
+     2 | SELECT * FROM temp_t
+     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(5 rows)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
 -- Tables, indexes, triggers
 CREATE TEMP TABLE tab_stats (a int, b char(20));
 CREATE INDEX index_stats ON tab_stats(b, (b || 'data1'), (b || 'data2')) WHERE a > 0;
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index dd97203c2102..e21b656d44a8 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -6,6 +6,18 @@
 SET pg_stat_statements.track_utility = TRUE;
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 
+-- Temporary tables, grouped together.
+BEGIN;
+  CREATE TEMP TABLE temp_t (id int) ON COMMIT DROP;
+  SELECT * FROM temp_t;
+COMMIT;
+BEGIN;
+  CREATE TEMP TABLE temp_t (id int) ON COMMIT DROP;
+  SELECT * FROM temp_t;
+COMMIT;
+SELECT calls, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
 -- Tables, indexes, triggers
 CREATE TEMP TABLE tab_stats (a int, b char(20));
 CREATE INDEX index_stats ON tab_stats(b, (b || 'data1'), (b || 'data2')) WHERE a > 0;
-- 
2.49.0