v6-0005-Extract-nested-query-from-PrepareStmt.patch

From 0a0a481cc0218bf59b290a4097fb08da6ef696cf Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Tue, 8 Oct 2024 08:45:39 +0200
Subject: Extract nested query from PrepareStmt

Previously, executing a prepared query would report the full "PREPARE x
AS ..." in pg_stat_statements. This patch extracts the
Insert/Select/Delete/Update statement nested in prepare and report only
the relevant query part.
---
 .../expected/level_tracking.out               |  2 +-
 .../pg_stat_statements/expected/planning.out  | 10 +++---
 .../pg_stat_statements/expected/select.out    |  2 +-
 .../pg_stat_statements/expected/utility.out   | 19 +++++++++--
 contrib/pg_stat_statements/sql/planning.sql   |  4 +--
 contrib/pg_stat_statements/sql/utility.sql    |  4 +++
 src/backend/commands/prepare.c                | 33 +++++++++++++++++--
 7 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 320cb93b597..d36bdaa665b 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -886,8 +886,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  f        |     1 | SELECT $1
+ f        |     1 | select generate_series($1, $2)
 (5 rows)
 
 -- Create Table As, top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..6c0a4a543c2 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -532,6 +532,19 @@ EXECUTE stat_select (2);
  2
 (1 row)
 
+PREPARE stat_select_2 AS SELECT $1, $2 AS a\; PREPARE stat_select_3 AS SELECT $1, $2, $3 AS a;
+EXECUTE stat_select_2 (2, 3);
+ ?column? | a 
+----------+---
+ 2        | 3
+(1 row)
+
+EXECUTE stat_select_3 (2, 3, 4);
+ ?column? | ?column? | a 
+----------+----------+---
+ 2        | 3        | 4
+(1 row)
+
 DEALLOCATE PREPARE stat_select;
 DEALLOCATE ALL;
 DEALLOCATE PREPARE ALL;
@@ -540,10 +553,12 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
+     1 |    1 | SELECT $1, $2 AS a
+     1 |    1 | SELECT $1, $2, $3 AS a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(5 rows)
+(7 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t 
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index dd97203c210..7a7df695910 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -270,6 +270,10 @@ EXECUTE stat_select (1);
 DEALLOCATE stat_select;
 PREPARE stat_select AS SELECT $1 AS a;
 EXECUTE stat_select (2);
+PREPARE stat_select_2 AS SELECT $1, $2 AS a\; PREPARE stat_select_3 AS SELECT $1, $2, $3 AS a;
+EXECUTE stat_select_2 (2, 3);
+EXECUTE stat_select_3 (2, 3, 4);
+
 DEALLOCATE PREPARE stat_select;
 DEALLOCATE ALL;
 DEALLOCATE PREPARE ALL;
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index a93f970a292..39aa447cf98 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -77,8 +77,37 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 	 */
 	rawstmt = makeNode(RawStmt);
 	rawstmt->stmt = stmt->query;
-	rawstmt->stmt_location = stmt_location;
-	rawstmt->stmt_len = stmt_len;
+
+	switch (nodeTag(stmt->query))
+	{
+		case T_InsertStmt:
+			rawstmt->stmt_location = ((InsertStmt *) stmt->query)->location;
+			break;
+		case T_DeleteStmt:
+			rawstmt->stmt_location = ((DeleteStmt *) stmt->query)->location;
+			break;
+		case T_UpdateStmt:
+			rawstmt->stmt_location = ((UpdateStmt *) stmt->query)->location;
+			break;
+		case T_MergeStmt:
+			rawstmt->stmt_location = ((MergeStmt *) stmt->query)->location;
+			break;
+		case T_SelectStmt:
+			rawstmt->stmt_location = ((SelectStmt *) stmt->query)->location;
+			rawstmt->stmt_len = ((SelectStmt *) stmt->query)->stmt_len;
+			break;
+		default:
+			elog(ERROR, "unexpected node type: %d", (int) nodeTag(stmt->query));
+			break;
+	}
+
+	/*
+	 * stmt_len will be defined for SelectStmt within parentheses. If it's
+	 * defined, use it. Otherwise, we need to compute the new length based on
+	 * the new statement location and the initial location+length.
+	 */
+	if (stmt_len > 0 && rawstmt->stmt_len == 0)
+		rawstmt->stmt_len = stmt_len - (rawstmt->stmt_location - stmt_location);
 
 	/*
 	 * Create the CachedPlanSource before we do parse analysis, since it needs
-- 
2.39.5 (Apple Git-154)