From 0176ea9412c4ad49621fb7e3267eed403d9fe6fb Mon Sep 17 00:00:00 2001 From: Sami Imseih Date: Wed, 30 Apr 2025 10:13:15 -0500 Subject: [PATCH v1 1/1] 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 0fc4b2c098d..b7bda01d8f1 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 cursor_stats_1 2 | 0 | DECLARE cursor_stats_1 CURSOR WITH HOLD FOR SELECT $1 + 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 @@ -59,8 +60,9 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"; 1 | 0 | DECLARE cursor_stats_2 CURSOR WITH HOLD FOR SELECT $1 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 -(9 rows) +(10 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 03bea14d5da..bba78289962 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 FOOCUR 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 FOOCUR 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 aa4f0f7e628..c19454754dc 100644 --- a/contrib/pg_stat_statements/expected/utility.out +++ b/contrib/pg_stat_statements/expected/utility.out @@ -706,9 +706,10 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"; 1 | 7 | FETCH FORWARD ALL pgss_cursor 1 | 1 | 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 -(12 rows) +(13 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)