v2-0002-Improve-cursor-handling-in-pg_stat_statements.patch

From 1746dff6639fadcbd0f530144483ccb83e44123f Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Wed, 30 Apr 2025 14:27:39 -0500
Subject: [PATCH v2 2/2] Improve cursor handling in pg_stat_statements

---
 .../pg_stat_statements/expected/cursors.out   |  6 ++-
 .../expected/level_tracking.out               |  5 +-
 .../pg_stat_statements/expected/utility.out   |  3 +-
 .../pg_stat_statements/pg_stat_statements.c   | 54 +++++++++++++------
 4 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/cursors.out b/contrib/pg_stat_statements/expected/cursors.out
index 1f4de404707..5c4103e33d6 100644
--- a/contrib/pg_stat_statements/expected/cursors.out
+++ b/contrib/pg_stat_statements/expected/cursors.out
@@ -20,8 +20,9 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | CLOSE $1
      2 |    0 | DECLARE $1 CURSOR WITH HOLD FOR SELECT $2
+     2 |    2 | SELECT $1
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(3 rows)
+(4 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t 
@@ -57,8 +58,9 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    0 | DECLARE $1 CURSOR WITH HOLD FOR SELECT $2
      1 |    1 | FETCH $1 IN cursor_stats_1
      1 |    1 | FETCH $1 IN cursor_stats_2
+     2 |    2 | SELECT $1
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(7 rows)
+(8 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t 
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 2b2f1d5ee42..e587c3269ac 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -1143,7 +1143,7 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | COMMIT
  t        |     1 | DECLARE $1 CURSOR FOR SELECT * from stats_track_tab
  t        |     1 | FETCH FORWARD $1 FROM foocur
- f        |     1 | SELECT * from stats_track_tab
+ t        |     1 | SELECT * from stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (7 rows)
 
@@ -1173,8 +1173,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | COMMIT
  t        |     1 | DECLARE $1 CURSOR FOR SELECT * FROM stats_track_tab
  t        |     1 | FETCH FORWARD $1 FROM foocur
+ t        |     1 | SELECT * FROM stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(6 rows)
+(7 rows)
 
 -- COPY - all-level tracking.
 SET pg_stat_statements.track = 'all';
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index e8b27d8304b..bb91e97941a 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -704,9 +704,10 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      1 |    0 | DECLARE $1 CURSOR FOR SELECT * FROM pgss_matv
      3 |   13 | FETCH NEXT pgss_cursor
      1 |   13 | REFRESH MATERIALIZED VIEW pgss_matv
+     1 |   13 | SELECT * FROM pgss_matv
      1 |   10 | SELECT generate_series($1, $2) c INTO pgss_select_into
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(10 rows)
+(11 rows)
 
 DROP MATERIALIZED VIEW pgss_matv;
 DROP TABLE pgss_ctas;
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 9778407cba3..91d26518512 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -311,6 +311,11 @@ static bool pgss_save = true;	/* whether to save stats across shutdown */
 		SpinLockRelease(&pgss->mutex); \
 	} while(0)
 
+#define is_cursor_utility() \
+	(IsA(parsetree, DeclareCursorStmt) || \
+	IsA(parsetree, FetchStmt) || \
+	IsA(parsetree, ClosePortalStmt))
+
 /*---- Function declarations ----*/
 
 PG_FUNCTION_INFO_V1(pg_stat_statements_reset);
@@ -1123,6 +1128,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	int			saved_stmt_location = pstmt->stmt_location;
 	int			saved_stmt_len = pstmt->stmt_len;
 	bool		enabled = pgss_track_utility && pgss_enabled(nesting_level);
+	bool		bump_level;
 
 	/*
 	 * Force utility statements to get queryId zero.  We do this even in cases
@@ -1153,8 +1159,14 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 * ensuing EXECUTEs.  This would be confusing.  Since PREPARE doesn't
 	 * actually run the planner (only parse+rewrite), its costs are generally
 	 * pretty negligible and it seems okay to just ignore it.
+	 *
+	 * If it's a cursor related utility command, don't bump the level and
+	 * force the tracking of the underlying sql.
 	 */
-	if (enabled &&
+	bump_level =
+		!is_cursor_utility();
+
+	if ((enabled || !bump_level) &&
 		!IsA(parsetree, ExecuteStmt) &&
 		!IsA(parsetree, PrepareStmt))
 	{
@@ -1165,12 +1177,14 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 					bufusage;
 		WalUsage	walusage_start,
 					walusage;
+		bool		store_utility = true;
 
 		bufusage_start = pgBufferUsage;
 		walusage_start = pgWalUsage;
 		INSTR_TIME_SET_CURRENT(start);
 
-		nesting_level++;
+		if (bump_level)
+			nesting_level++;
 		PG_TRY();
 		{
 			if (prev_ProcessUtility)
@@ -1184,7 +1198,8 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		}
 		PG_FINALLY();
 		{
-			nesting_level--;
+			if (bump_level)
+				nesting_level--;
 		}
 		PG_END_TRY();
 
@@ -1220,19 +1235,24 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		memset(&walusage, 0, sizeof(WalUsage));
 		WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start);
 
-		pgss_store(queryString,
-				   saved_queryId,
-				   saved_stmt_location,
-				   saved_stmt_len,
-				   PGSS_EXEC,
-				   INSTR_TIME_GET_MILLISEC(duration),
-				   rows,
-				   &bufusage,
-				   &walusage,
-				   NULL,
-				   NULL,
-				   0,
-				   0);
+		/*
+		 * In case we got here because of a cursor command, let's make sure we
+		 * have utility tracking enabled, again.
+		 */
+		if (enabled)
+			pgss_store(queryString,
+					   saved_queryId,
+					   saved_stmt_location,
+					   saved_stmt_len,
+					   PGSS_EXEC,
+					   INSTR_TIME_GET_MILLISEC(duration),
+					   rows,
+					   &bufusage,
+					   &walusage,
+					   NULL,
+					   NULL,
+					   0,
+					   0);
 	}
 	else
 	{
@@ -1248,7 +1268,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		 * To be absolutely certain we don't mess up the nesting level,
 		 * evaluate the bump_level condition just once.
 		 */
-		bool		bump_level =
+		bump_level =
 			!IsA(parsetree, ExecuteStmt) &&
 			!IsA(parsetree, PrepareStmt);
 
-- 
2.39.5 (Apple Git-154)