v8-0001-JumbleQuery-on-Call-and-Set.patch

diff --git a/contrib/pg_stat_statements/expected/pg_stat_statements.out b/contrib/pg_stat_statements/expected/pg_stat_statements.out
index 9ac5c87c3a..5d3330a87d 100644
--- a/contrib/pg_stat_statements/expected/pg_stat_statements.out
+++ b/contrib/pg_stat_statements/expected/pg_stat_statements.out
@@ -311,7 +311,7 @@ FROM pg_stat_statements ORDER BY query COLLATE "C";
  wal_records > $2 as wal_records_generated,               +|       |      |                     |                       | 
  wal_records >= rows as wal_records_ge_rows               +|       |      |                     |                       | 
  FROM pg_stat_statements ORDER BY query COLLATE "C"        |       |      |                     |                       | 
- SET pg_stat_statements.track_utility = FALSE              |     1 |    0 | f                   | f                     | t
+ SET pg_stat_statements.track_utility = $1                 |     1 |    0 | f                   | f                     | t
  UPDATE pgss_test SET b = $1 WHERE a > $2                  |     1 |    3 | t                   | t                     | t
 (7 rows)
 
@@ -462,6 +462,111 @@ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
  SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
 (6 rows)
 
+-- PL/pgSQL procedure and pg_stat_statements.track = top
+CREATE PROCEDURE MINUS_TWO(i INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i - 1 - 1.0)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+CREATE PROCEDURE SUM_TWO(i INTEGER, j INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (j + j)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+SET pg_stat_statements.track = 'top';
+SET pg_stat_statements.track_utility = FALSE;
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+CALL MINUS_TWO(3);
+CALL MINUS_TWO(7);
+CALL SUM_TWO(3, 8);
+CALL SUM_TWO(7, 5);
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                    query                                     | calls | rows 
+------------------------------------------------------------------------------+-------+------
+ CALL MINUS_TWO($1)                                                           |     2 |    0
+ CALL SUM_TWO($1, $2)                                                         |     2 |    0
+ SELECT pg_stat_statements_reset()                                            |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
+(4 rows)
+
+--
+-- pg_stat_statements.track = all
+--
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+-- PL/pgSQL procedure and pg_stat_statements.track = all
+-- we drop and recreate the procedures to avoid any caching funnies
+DROP PROCEDURE MINUS_TWO(INTEGER);
+DROP PROCEDURE SUM_TWO(INTEGER, INTEGER);
+CREATE PROCEDURE MINUS_TWO(i INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i - 1 - 1.0)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+CREATE PROCEDURE SUM_TWO(i INTEGER, j INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (j + j)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+CALL MINUS_TWO(3);
+CALL MINUS_TWO(7);
+CALL SUM_TWO(3, 8);
+CALL SUM_TWO(7, 5);
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                    query                                     | calls | rows 
+------------------------------------------------------------------------------+-------+------
+ CALL MINUS_TWO($1)                                                           |     2 |    0
+ CALL SUM_TWO($1, $2)                                                         |     2 |    0
+ SELECT (i - $2 - $3)::INTEGER                                                |     2 |    2
+ SELECT (j + j)::INTEGER                                                      |     2 |    2
+ SELECT pg_stat_statements_reset()                                            |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
+(6 rows)
+
+SET pg_stat_statements.track_utility = TRUE;
+-- SET
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+set enable_seqscan=false;
+set enable_seqscan=true;
+set seq_page_cost=2.0;
+set seq_page_cost=1.0;
+set enable_seqscan to default;
+set seq_page_cost to default;
+reset seq_page_cost;
+reset enable_seqscan;
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                    query                                     | calls | rows 
+------------------------------------------------------------------------------+-------+------
+ SELECT pg_stat_statements_reset()                                            |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
+ reset enable_seqscan                                                         |     1 |    0
+ reset seq_page_cost                                                          |     1 |    0
+ set enable_seqscan to default                                                |     1 |    0
+ set enable_seqscan=$1                                                        |     2 |    0
+ set seq_page_cost to default                                                 |     1 |    0
+ set seq_page_cost=$1                                                         |     2 |    0
+(8 rows)
+
+SET pg_stat_statements.track_utility = FALSE;
 --
 -- queries with locking clauses
 --
@@ -592,11 +697,49 @@ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
  SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
 (9 rows)
 
+-- 2PC
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+create table test_tx (a int);
+begin;
+prepare transaction 'tx1';
+insert into test_tx values (1);
+commit prepared 'tx1';
+begin;
+prepare transaction 'tx2';
+insert into test_tx values (2);
+commit prepared 'tx2';
+begin;
+prepare transaction 'tx3';
+insert into test_tx values (3);
+rollback prepared 'tx3';
+begin;
+prepare transaction 'tx4';
+insert into test_tx values (4);
+rollback prepared 'tx4';
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                    query                                     | calls | rows 
+------------------------------------------------------------------------------+-------+------
+ SELECT pg_stat_statements_reset()                                            |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C" |     0 |    0
+ begin                                                                        |     4 |    0
+ commit prepared $1                                                           |     2 |    0
+ create table test_tx (a int)                                                 |     1 |    0
+ insert into test_tx values ($1)                                              |     4 |    4
+ prepare transaction $1                                                       |     4 |    0
+ rollback prepared $1                                                         |     2 |    0
+(8 rows)
+
 --
 -- Track the total number of rows retrieved or affected by the utility
 -- commands of COPY, FETCH, CREATE TABLE AS, CREATE MATERIALIZED VIEW,
 -- REFRESH MATERIALIZED VIEW and SELECT INTO
 --
+SET pg_stat_statements.track_utility = TRUE;
 SELECT pg_stat_statements_reset();
  pg_stat_statements_reset 
 --------------------------
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 73439c0199..4aefbe70f0 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -106,6 +106,11 @@ static const uint32 PGSS_PG_MAJOR_VERSION = PG_VERSION_NUM / 100;
 #define PGSS_HANDLED_UTILITY(n)		(!IsA(n, ExecuteStmt) && \
 									!IsA(n, PrepareStmt) && \
 									!IsA(n, DeallocateStmt))
+/*
+ * Force track those utility statements
+ * whatever the value of pgss_track_utility is.
+ */
+#define FORCE_TRACK_UTILITY(n)	(IsA(n, CallStmt))
 
 /*
  * Extension version number, for supporting older extension versions' objects
@@ -832,9 +837,10 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 	 * inherit from the underlying statement's one (except DEALLOCATE which is
 	 * entirely untracked).
 	 */
-	if (query->utilityStmt)
+	if (query->utilityStmt && !jstate)
 	{
-		if (pgss_track_utility && !PGSS_HANDLED_UTILITY(query->utilityStmt))
+		if ((pgss_track_utility || FORCE_TRACK_UTILITY(query->utilityStmt)) &&
+			!PGSS_HANDLED_UTILITY(query->utilityStmt))
 			query->queryId = UINT64CONST(0);
 		return;
 	}
@@ -1097,7 +1103,8 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 * that user configured another extension to handle utility statements
 	 * only.
 	 */
-	if (pgss_enabled(exec_nested_level) && pgss_track_utility)
+	if (pgss_enabled(exec_nested_level) &&
+		(pgss_track_utility || FORCE_TRACK_UTILITY(parsetree)))
 		pstmt->queryId = UINT64CONST(0);
 
 	/*
@@ -1114,7 +1121,8 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 *
 	 * Likewise, we don't track execution of DEALLOCATE.
 	 */
-	if (pgss_track_utility && pgss_enabled(exec_nested_level) &&
+	if ((pgss_track_utility || FORCE_TRACK_UTILITY(parsetree)) &&
+		pgss_enabled(exec_nested_level) &&
 		PGSS_HANDLED_UTILITY(parsetree))
 	{
 		instr_time	start;
diff --git a/contrib/pg_stat_statements/sql/pg_stat_statements.sql b/contrib/pg_stat_statements/sql/pg_stat_statements.sql
index 8f5c866225..34b0d61815 100644
--- a/contrib/pg_stat_statements/sql/pg_stat_statements.sql
+++ b/contrib/pg_stat_statements/sql/pg_stat_statements.sql
@@ -223,6 +223,81 @@ SELECT PLUS_ONE(1);
 
 SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+-- PL/pgSQL procedure and pg_stat_statements.track = top
+CREATE PROCEDURE MINUS_TWO(i INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i - 1 - 1.0)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+
+CREATE PROCEDURE SUM_TWO(i INTEGER, j INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (j + j)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+
+SET pg_stat_statements.track = 'top';
+SET pg_stat_statements.track_utility = FALSE;
+SELECT pg_stat_statements_reset();
+CALL MINUS_TWO(3);
+CALL MINUS_TWO(7);
+CALL SUM_TWO(3, 8);
+CALL SUM_TWO(7, 5);
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- pg_stat_statements.track = all
+--
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset();
+
+-- PL/pgSQL procedure and pg_stat_statements.track = all
+-- we drop and recreate the procedures to avoid any caching funnies
+
+DROP PROCEDURE MINUS_TWO(INTEGER);
+DROP PROCEDURE SUM_TWO(INTEGER, INTEGER);
+
+CREATE PROCEDURE MINUS_TWO(i INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i - 1 - 1.0)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+
+CREATE PROCEDURE SUM_TWO(i INTEGER, j INTEGER) AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (j + j)::INTEGER INTO r;
+END; $$ LANGUAGE plpgsql;
+
+CALL MINUS_TWO(3);
+CALL MINUS_TWO(7);
+CALL SUM_TWO(3, 8);
+CALL SUM_TWO(7, 5);
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+SET pg_stat_statements.track_utility = TRUE;
+
+-- SET
+SELECT pg_stat_statements_reset();
+set enable_seqscan=false;
+set enable_seqscan=true;
+set seq_page_cost=2.0;
+set seq_page_cost=1.0;
+set enable_seqscan to default;
+set seq_page_cost to default;
+reset seq_page_cost;
+reset enable_seqscan;
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+SET pg_stat_statements.track_utility = FALSE;
+
 --
 -- queries with locking clauses
 --
@@ -272,11 +347,38 @@ DROP FUNCTION PLUS_TWO(INTEGER);
 
 SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+-- 2PC
+SELECT pg_stat_statements_reset();
+
+create table test_tx (a int);
+begin;
+prepare transaction 'tx1';
+insert into test_tx values (1);
+commit prepared 'tx1';
+
+begin;
+prepare transaction 'tx2';
+insert into test_tx values (2);
+commit prepared 'tx2';
+
+begin;
+prepare transaction 'tx3';
+insert into test_tx values (3);
+rollback prepared 'tx3';
+
+begin;
+prepare transaction 'tx4';
+insert into test_tx values (4);
+rollback prepared 'tx4';
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
 --
 -- Track the total number of rows retrieved or affected by the utility
 -- commands of COPY, FETCH, CREATE TABLE AS, CREATE MATERIALIZED VIEW,
 -- REFRESH MATERIALIZED VIEW and SELECT INTO
 --
+SET pg_stat_statements.track_utility = TRUE;
 SELECT pg_stat_statements_reset();
 
 CREATE TABLE pgss_ctas AS SELECT a, 'ctas' b FROM generate_series(1, 10) a;
diff --git a/doc/src/sgml/pgstatstatements.sgml b/doc/src/sgml/pgstatstatements.sgml
index ea90365c7f..db91bb20bd 100644
--- a/doc/src/sgml/pgstatstatements.sgml
+++ b/doc/src/sgml/pgstatstatements.sgml
@@ -487,13 +487,15 @@
 
   <para>
    Plannable queries (that is, <command>SELECT</command>, <command>INSERT</command>,
-   <command>UPDATE</command>, <command>DELETE</command>, and <command>MERGE</command>) are combined into a single
-   <structname>pg_stat_statements</structname> entry whenever they have identical query
-   structures according to an internal hash calculation.  Typically, two
-   queries will be considered the same for this purpose if they are
-   semantically equivalent except for the values of literal constants
-   appearing in the query.  Utility commands (that is, all other commands)
-   are compared strictly on the basis of their textual query strings, however.
+   <command>UPDATE</command>, <command>DELETE</command>, and <command>MERGE</command>)
+   as well as <command>CALL</command>, <command>SET</command>, <command>PREPARE TRANSACTION</command>,
+   <command>COMMIT PREPARED</command> and <command>ROLLBACK PREPARED</command>
+   are combined into a single <structname>pg_stat_statements</structname> entry
+   whenever they have identical query structures according to an internal hash calculation.
+   Typically, two queries will be considered the same for this purpose if they are
+   semantically equivalent except for the values of literal constants appearing
+   in the command. All other commands are compared strictly on the basis
+   of their textual query strings, however.
   </para>
 
   <note>
@@ -781,10 +783,10 @@
     <listitem>
      <para>
       <varname>pg_stat_statements.track_utility</varname> controls whether
-      utility commands are tracked by the module.  Utility commands are
+      utility commands are tracked by the module. Tracked utility commands are
       all those other than <command>SELECT</command>, <command>INSERT</command>,
-      <command>UPDATE</command>, <command>DELETE</command>, and <command>MERGE</command>.
-      The default value is <literal>on</literal>.
+      <command>UPDATE</command>, <command>DELETE</command>, <command>MERGE</command>,
+      and <command>CALL</command>. The default value is <literal>on</literal>.
       Only superusers can change this setting.
      </para>
     </listitem>
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 94d5142a4a..7203a9411a 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -10890,6 +10890,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_PREPARE;
 					n->gid = $3;
+					n->gid_location = @3;
 					$$ = (Node *) n;
 				}
 			| COMMIT PREPARED Sconst
@@ -10898,6 +10899,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_COMMIT_PREPARED;
 					n->gid = $3;
+					n->gid_location = @3;
 					$$ = (Node *) n;
 				}
 			| ROLLBACK PREPARED Sconst
@@ -10906,6 +10908,7 @@ TransactionStmt:
 
 					n->kind = TRANS_STMT_ROLLBACK_PREPARED;
 					n->gid = $3;
+					n->gid_location = @3;
 					$$ = (Node *) n;
 				}
 		;
diff --git a/src/backend/utils/misc/queryjumble.c b/src/backend/utils/misc/queryjumble.c
index a8508463e7..02f5178b6b 100644
--- a/src/backend/utils/misc/queryjumble.c
+++ b/src/backend/utils/misc/queryjumble.c
@@ -105,7 +105,7 @@ JumbleQuery(Query *query, const char *querytext)
 
 	Assert(IsQueryIdEnabled());
 
-	if (query->utilityStmt)
+	if (query->utilityStmt && !IsJumbleUtilityAllowed(query->utilityStmt))
 	{
 		query->queryId = compute_utility_query_id(querytext,
 												  query->stmt_location,
@@ -241,10 +241,11 @@ static void
 JumbleQueryInternal(JumbleState *jstate, Query *query)
 {
 	Assert(IsA(query, Query));
-	Assert(query->utilityStmt == NULL);
+	Assert(query->utilityStmt == NULL || IsJumbleUtilityAllowed(query->utilityStmt));
 
 	APP_JUMB(query->commandType);
 	/* resultRelation is usually predictable from commandType */
+	JumbleExpr(jstate, (Node *) query->utilityStmt);
 	JumbleExpr(jstate, (Node *) query->cteList);
 	JumbleRangeTable(jstate, query->rtable);
 	JumbleExpr(jstate, (Node *) query->jointree);
@@ -385,6 +386,47 @@ JumbleExpr(JumbleState *jstate, Node *node)
 				APP_JUMB(var->varlevelsup);
 			}
 			break;
+		case T_CallStmt:
+			{
+				CallStmt   *stmt = (CallStmt *) node;
+				FuncExpr   *expr = stmt->funcexpr;
+
+				APP_JUMB(expr->funcid);
+				JumbleExpr(jstate, (Node *) expr->args);
+			}
+			break;
+		case T_VariableSetStmt:
+			{
+				VariableSetStmt *stmt = (VariableSetStmt *) node;
+
+				/* stmt->name is NULL for RESET ALL */
+				if (stmt->name)
+				{
+					APP_JUMB(stmt->kind);
+					APP_JUMB_STRING(stmt->name);
+					JumbleExpr(jstate, (Node *) stmt->args);
+				}
+			}
+			break;
+		case T_A_Const:
+			{
+				int			loc = ((const A_Const *) node)->location;
+
+				RecordConstLocation(jstate, loc);
+			}
+			break;
+		case T_TransactionStmt:
+			{
+				TransactionStmt *stmt = (TransactionStmt *) node;
+
+				Assert(stmt->kind == TRANS_STMT_PREPARE ||
+					   stmt->kind == TRANS_STMT_COMMIT_PREPARED ||
+					   stmt->kind == TRANS_STMT_ROLLBACK_PREPARED);
+
+				APP_JUMB(stmt->kind);
+				RecordConstLocation(jstate, stmt->gid_location);
+			}
+			break;
 		case T_Const:
 			{
 				Const	   *c = (Const *) node;
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 633e7671b3..d62fac9d86 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -3282,6 +3282,7 @@ typedef struct TransactionStmt
 	char	   *savepoint_name; /* for savepoint commands */
 	char	   *gid;			/* for two-phase-commit related commands */
 	bool		chain;			/* AND CHAIN option */
+	int			gid_location;	/* gid location */
 } TransactionStmt;
 
 /* ----------------------
diff --git a/src/include/utils/queryjumble.h b/src/include/utils/queryjumble.h
index f799fb19e4..bb15c7e1d8 100644
--- a/src/include/utils/queryjumble.h
+++ b/src/include/utils/queryjumble.h
@@ -83,4 +83,24 @@ IsQueryIdEnabled(void)
 	return query_id_enabled;
 }
 
+/*
+ * Jumble those utility statements
+ */
+static inline bool
+IsJumbleUtilityAllowed(Node *n)
+{
+	if (IsA(n, CallStmt) || IsA(n, VariableSetStmt))
+		return true;
+	else if (IsA(n, TransactionStmt))
+	{
+		TransactionStmt *stmt = (TransactionStmt *) n;
+
+		return (stmt->kind == TRANS_STMT_PREPARE ||
+				stmt->kind == TRANS_STMT_COMMIT_PREPARED ||
+				stmt->kind == TRANS_STMT_ROLLBACK_PREPARED);
+	}
+
+	return false;
+}
+
 #endif							/* QUERYJUMBLE_H */