From 638d42a3c52081cf4882520f0622161bad69b40f Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Sat, 12 Aug 2023 10:44:15 +0900
Subject: [PATCH] Show GIDs of two-phase commit commands as constants in
 pg_stat_statements

This relies on the "location" field added to TransactionStmt in 31de7e6,
now applied to the "gid" field used by 2PC commands.  These commands are
now reported like:
COMMIT PREPARED $1
PREPARE TRANSACTION $1
ROLLBACK PREPARED $1

Applying constants for these commands is a huge advantage for workloads
that rely a lot on 2PC commands with different GIDs.  Some tests are
added to track the new behavior.

Reviewed-by: Julien Rouhaud
Discussion: https://postgr.es/m/ZMhT9kNtJJsHw6jK@paquier.xyz
---
 .../pg_stat_statements/expected/utility.out   | 23 +++++++++++++++++++
 .../pg_stat_statements.conf                   |  1 +
 contrib/pg_stat_statements/sql/utility.sql    | 10 ++++++++
 src/backend/parser/gram.y                     |  6 ++---
 src/include/nodes/parsenodes.h                |  3 ++-
 5 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 3d920fb5f7..93735d5d85 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -197,6 +197,29 @@ SELECT pg_stat_statements_reset();
  
 (1 row)
 
+-- Two-phase transactions
+BEGIN;
+PREPARE TRANSACTION 'stat_trans1';
+COMMIT PREPARED 'stat_trans1';
+BEGIN;
+PREPARE TRANSACTION 'stat_trans2';
+ROLLBACK PREPARED 'stat_trans2';
+SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+ calls | rows |               query               
+-------+------+-----------------------------------
+     2 |    0 | BEGIN
+     1 |    0 | COMMIT PREPARED $1
+     2 |    0 | PREPARE TRANSACTION $1
+     1 |    0 | ROLLBACK PREPARED $1
+     1 |    1 | SELECT pg_stat_statements_reset()
+(5 rows)
+
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
 -- Savepoints
 BEGIN;
 SAVEPOINT sp1;
diff --git a/contrib/pg_stat_statements/pg_stat_statements.conf b/contrib/pg_stat_statements/pg_stat_statements.conf
index 13346e2807..0e900d7119 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.conf
+++ b/contrib/pg_stat_statements/pg_stat_statements.conf
@@ -1 +1,2 @@
 shared_preload_libraries = 'pg_stat_statements'
+max_prepared_transactions = 5
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index 859e57955e..87666d9135 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -115,6 +115,16 @@ COMMIT;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 SELECT pg_stat_statements_reset();
 
+-- Two-phase transactions
+BEGIN;
+PREPARE TRANSACTION 'stat_trans1';
+COMMIT PREPARED 'stat_trans1';
+BEGIN;
+PREPARE TRANSACTION 'stat_trans2';
+ROLLBACK PREPARED 'stat_trans2';
+SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+SELECT pg_stat_statements_reset();
+
 -- Savepoints
 BEGIN;
 SAVEPOINT sp1;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 15ece871a0..b3bdf947b6 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -10924,7 +10924,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_PREPARE;
 					n->gid = $3;
-					n->location = -1;
+					n->location = @3;
 					$$ = (Node *) n;
 				}
 			| COMMIT PREPARED Sconst
@@ -10933,7 +10933,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_COMMIT_PREPARED;
 					n->gid = $3;
-					n->location = -1;
+					n->location = @3;
 					$$ = (Node *) n;
 				}
 			| ROLLBACK PREPARED Sconst
@@ -10942,7 +10942,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_ROLLBACK_PREPARED;
 					n->gid = $3;
-					n->location = -1;
+					n->location = @3;
 					$$ = (Node *) n;
 				}
 		;
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index fe003ded50..2565348303 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -3540,7 +3540,8 @@ typedef struct TransactionStmt
 	List	   *options;		/* for BEGIN/START commands */
 	/* for savepoint commands */
 	char	   *savepoint_name pg_node_attr(query_jumble_ignore);
-	char	   *gid;			/* for two-phase-commit related commands */
+	/* for two-phase-commit related commands */
+	char	   *gid pg_node_attr(query_jumble_ignore);
 	bool		chain;			/* AND CHAIN option */
 	/* token location, or -1 if unknown */
 	int			location pg_node_attr(query_jumble_location);