Fix nesting level tracking

Makefile

@@ -12,7 +12,7 @@ LDFLAGS_SL += $(filter -lm, $(LIBS))
 
 TAP_TESTS = 1
 REGRESS_OPTS = --temp-config $(top_srcdir)/contrib/pg_stat_monitor/pg_stat_monitor.conf --inputdir=regression
-REGRESS = basic version guc pgsm_query_id functions counters relations database error_insert application_name application_name_unique top_query cmd_type error rows tags user
+REGRESS = basic version guc pgsm_query_id functions counters relations database error_insert application_name application_name_unique top_query cmd_type error rows tags user level_tracking
 
 # Disabled because these tests require "shared_preload_libraries=pg_stat_statements",
 # which typical installcheck users do not have (e.g. buildfarm clients).

pg_stat_monitor.c

@@ -73,11 +73,12 @@ do                                                      \
 /*---- Initicalization Function Declarations ----*/
 void		_PG_init(void);
 
-/* Current nesting depth of ExecutorRun+ProcessUtility calls */
+/* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */
-static int	exec_nested_level = 0;
+static int	nesting_level = 0;
 volatile bool __pgsm_do_not_capture_error = false;
 
-#if PG_VERSION_NUM >= 130000
+#if PG_VERSION_NUM >= 130000 && PG_VERSION_NUM < 170000
+/* Before planner nesting level was conunted separately */
 static int	plan_nested_level = 0;
 #endif
 
@@ -200,11 +201,6 @@ DECLARE_HOOK(void pgsm_ProcessUtility, PlannedStmt *pstmt, const char *queryStri
 static uint64 pgsm_hash_string(const char *str, int len);
 char	   *unpack_sql_state(int sql_state);
 
-#define PGSM_HANDLED_UTILITY(n)  (!IsA(n, ExecuteStmt) && \
-									!IsA(n, PrepareStmt) && \
-									!IsA(n, DeallocateStmt))
-
-
 static pgsmEntry * pgsm_create_hash_entry(uint64 bucket_id, uint64 queryid, PlanInfo * plan_info);
 static void pgsm_add_to_list(pgsmEntry * entry, char *query_text, int query_len);
 static pgsmEntry * pgsm_get_entry_for_query(uint64 queryid, PlanInfo * plan_info, const char *query_text, int query_len, bool create);
@@ -428,17 +424,18 @@ pgsm_post_parse_analyze_internal(ParseState *pstate, Query *query, JumbleState *
 		}
 	}
 
-	if (!pgsm_enabled(exec_nested_level))
+	if (!pgsm_enabled(nesting_level))
 		return;
 
 	/*
-	 * Clear queryId for prepared statements related utility, as those will
+	 * If it's EXECUTE, clear the queryId so that stats will accumulate for
-	 * inherit from the underlying statement's one (except DEALLOCATE which is
+	 * the underlying PREPARE.  But don't do this if we're not tracking
-	 * entirely untracked).
+	 * utility statements, to avoid messing up another extension that might be
+	 * tracking them.
 	 */
 	if (query->utilityStmt)
 	{
-		if (pgsm_track_utility && !PGSM_HANDLED_UTILITY(query->utilityStmt))
+		if (pgsm_track_utility && IsA(query->utilityStmt, ExecuteStmt))
 			query->queryId = UINT64CONST(0);
 
 		return;
@@ -571,7 +568,7 @@ pgsm_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 * counting of optimizable statements that are directly contained in
 	 * utility statements.
 	 */
-	if (pgsm_enabled(exec_nested_level) &&
+	if (pgsm_enabled(nesting_level) &&
 		queryDesc->plannedstmt->queryId != UINT64CONST(0))
 	{
 		/*
@@ -602,37 +599,37 @@ static void
 pgsm_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count,
 				 bool execute_once)
 {
-	if (exec_nested_level >= 0 && exec_nested_level < max_stack_depth)
+	if (nesting_level >= 0 && nesting_level < max_stack_depth)
 	{
-		nested_queryids[exec_nested_level] = queryDesc->plannedstmt->queryId;
+		nested_queryids[nesting_level] = queryDesc->plannedstmt->queryId;
-		nested_query_txts[exec_nested_level] = strdup(queryDesc->sourceText);
+		nested_query_txts[nesting_level] = strdup(queryDesc->sourceText);
 	}
 
-	exec_nested_level++;
+	nesting_level++;
 	PG_TRY();
 	{
 		if (prev_ExecutorRun)
 			prev_ExecutorRun(queryDesc, direction, count, execute_once);
 		else
 			standard_ExecutorRun(queryDesc, direction, count, execute_once);
-		exec_nested_level--;
+		nesting_level--;
-		if (exec_nested_level >= 0 && exec_nested_level < max_stack_depth)
+		if (nesting_level >= 0 && nesting_level < max_stack_depth)
 		{
-			nested_queryids[exec_nested_level] = UINT64CONST(0);
+			nested_queryids[nesting_level] = UINT64CONST(0);
-			if (nested_query_txts[exec_nested_level])
+			if (nested_query_txts[nesting_level])
-				free(nested_query_txts[exec_nested_level]);
+				free(nested_query_txts[nesting_level]);
-			nested_query_txts[exec_nested_level] = NULL;
+			nested_query_txts[nesting_level] = NULL;
 		}
 	}
 	PG_CATCH();
 	{
-		exec_nested_level--;
+		nesting_level--;
-		if (exec_nested_level >= 0 && exec_nested_level < max_stack_depth)
+		if (nesting_level >= 0 && nesting_level < max_stack_depth)
 		{
-			nested_queryids[exec_nested_level] = UINT64CONST(0);
+			nested_queryids[nesting_level] = UINT64CONST(0);
-			if (nested_query_txts[exec_nested_level])
+			if (nested_query_txts[nesting_level])
-				free(nested_query_txts[exec_nested_level]);
+				free(nested_query_txts[nesting_level]);
-			nested_query_txts[exec_nested_level] = NULL;
+			nested_query_txts[nesting_level] = NULL;
 		}
 		PG_RE_THROW();
 	}
@@ -645,7 +642,7 @@ pgsm_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count,
 static void
 pgsm_ExecutorFinish(QueryDesc *queryDesc)
 {
-	exec_nested_level++;
+	nesting_level++;
 
 	PG_TRY();
 	{
@@ -653,11 +650,11 @@ pgsm_ExecutorFinish(QueryDesc *queryDesc)
 			prev_ExecutorFinish(queryDesc);
 		else
 			standard_ExecutorFinish(queryDesc);
-		exec_nested_level--;
+		nesting_level--;
 	}
 	PG_CATCH();
 	{
-		exec_nested_level--;
+		nesting_level--;
 		PG_RE_THROW();
 	}
 	PG_END_TRY();
@@ -725,7 +722,7 @@ pgsm_ExecutorEnd(QueryDesc *queryDesc)
 		MemoryContextSwitchTo(oldctx);
 	}
 
-	if (queryId != UINT64CONST(0) && queryDesc->totaltime && pgsm_enabled(exec_nested_level))
+	if (queryId != UINT64CONST(0) && queryDesc->totaltime && pgsm_enabled(nesting_level))
 	{
 		entry = pgsm_get_entry_for_query(queryId, plan_ptr, (char *) queryDesc->sourceText, strlen(queryDesc->sourceText), true);
 		if (!entry)
@@ -871,8 +868,14 @@ pgsm_planner_hook(Query *parse, const char *query_string, int cursorOptions, Par
 	 * top level planner call.
 	 */
 
-	if (pgsm_enabled(plan_nested_level + exec_nested_level) &&
+	bool enabled;
-		pgsm_track_planning && query_string && parse->queryId != UINT64CONST(0))
+#if PG_VERSION_NUM >= 170000
+	enabled = pgsm_enabled(nesting_level);
+#else
+	enabled = pgsm_enabled(plan_nested_level + nesting_level);
+#endif
+
+	if (enabled && pgsm_track_planning && query_string && parse->queryId != UINT64CONST(0))
 	{
 		pgsmEntry  *entry = NULL;
 		instr_time	start;
@@ -895,7 +898,11 @@ pgsm_planner_hook(Query *parse, const char *query_string, int cursorOptions, Par
 		if (MemoryContextIsValid(MessageContext))
 			entry = pgsm_get_entry_for_query(parse->queryId, NULL, query_string, strlen(query_string), true);
 
+#if PG_VERSION_NUM >= 170000
+		nesting_level++;
+#else
 		plan_nested_level++;
+#endif
 		PG_TRY();
 		{
 			/*
@@ -912,7 +919,11 @@ pgsm_planner_hook(Query *parse, const char *query_string, int cursorOptions, Par
 		}
 		PG_FINALLY();
 		{
-			plan_nested_level--;
+#if PG_VERSION_NUM >= 170000
+		nesting_level--;
+#else
+		plan_nested_level--;
+#endif
 		}
 		PG_END_TRY();
 
@@ -948,19 +959,38 @@ pgsm_planner_hook(Query *parse, const char *query_string, int cursorOptions, Par
 	else
 	{
 		/*
+		 * Even though we're not tracking plan time for this statement, we
+		 * must still increment the nesting level, to ensure that functions
+		 * evaluated during planning are not seen as top-level calls.
+		 * 
 		 * If there is a previous installed hook, then assume it's going to
 		 * call standard_planner() function, otherwise we call the function
 		 * here. This is to avoid calling standard_planner() function twice,
 		 * since it modifies the first argument (Query *), the second call
 		 * would trigger an assertion failure.
 		 */
-		plan_nested_level++;
 
-		if (planner_hook_next)
+#if PG_VERSION_NUM >= 170000
-			result = planner_hook_next(parse, query_string, cursorOptions, boundParams);
+		nesting_level++;
-		else
+#else
-			result = standard_planner(parse, query_string, cursorOptions, boundParams);
+		plan_nested_level++;
+#endif
+		PG_TRY();
+		{
+			if (planner_hook_next)
+				result = planner_hook_next(parse, query_string, cursorOptions, boundParams);
+			else
+				result = standard_planner(parse, query_string, cursorOptions, boundParams);
+		}
+		PG_FINALLY();
+		{
+#if PG_VERSION_NUM >= 170000
+		nesting_level--;
+#else
 		plan_nested_level--;
+#endif
+		}
+		PG_END_TRY();
 
 	}
 	return result;
@@ -998,6 +1028,7 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 {
 	Node	   *parsetree = pstmt->utilityStmt;
 	uint64		queryId = 0;
+	bool enabled = pgsm_track_utility && pgsm_enabled(nesting_level);
 
 #if PG_VERSION_NUM < 140000
 	int			len = strlen(queryString);
@@ -1015,7 +1046,7 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 * since we are already measuring the statement's costs at the utility
 	 * level.
 	 */
-	if (pgsm_track_utility && pgsm_enabled(exec_nested_level))
+	if (enabled)
 		pstmt->queryId = UINT64CONST(0);
 #endif
 
@@ -1028,13 +1059,16 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 * hash table entry for the PREPARE (with hash calculated from the query
 	 * string), and then a different one with the same query string (but hash
 	 * calculated from the query tree) would be used to accumulate costs of
-	 * ensuing EXECUTEs.  This would be confusing, and inconsistent with other
+	 * ensuing EXECUTEs.  This would be confusing.  Since PREPARE doesn't
-	 * cases where planning time is not included at all.
+	 * actually run the planner (only parse+rewrite), its costs are generally
+	 * pretty negligible and it seems okay to just ignore it.
 	 *
 	 * Likewise, we don't track execution of DEALLOCATE.
 	 */
-	if (pgsm_track_utility && pgsm_enabled(exec_nested_level) &&
+	if (enabled &&
-		PGSM_HANDLED_UTILITY(parsetree))
+		!IsA(parsetree, ExecuteStmt) && 
+ 		!IsA(parsetree, PrepareStmt) && 
+		!IsA(parsetree, DeallocateStmt))
 	{
 		pgsmEntry  *entry;
 		char	   *query_text;
@@ -1055,7 +1089,7 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 			elog(DEBUG1, "[pg_stat_monitor] pgsm_ProcessUtility: Failed to execute getrusage.");
 
 		INSTR_TIME_SET_CURRENT(start);
-		exec_nested_level++;
+		nesting_level++;
 
 		PG_TRY();
 		{
@@ -1095,11 +1129,11 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 										dest,
 										completionTag);
 #endif
-			exec_nested_level--;
+			nesting_level--;
 		}
 		PG_CATCH();
 		{
-			exec_nested_level--;
+			nesting_level--;
 			PG_RE_THROW();
 		}
 
@@ -1185,6 +1219,30 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	}
 	else
 	{
+		/*
+		 * Even though we're not tracking execution time for this statement,
+		 * we must still increment the nesting level, to ensure that functions
+		 * evaluated within it are not seen as top-level calls.  But don't do
+		 * so for EXECUTE; that way, when control reaches pgss_planner or
+		 * pgss_ExecutorStart, we will treat the costs as top-level if
+		 * appropriate.  Likewise, don't bump for PREPARE, so that parse
+		 * analysis will treat the statement as top-level if appropriate.
+		 * 
+		 * Likewise, we don't track execution of DEALLOCATE.
+		 *
+		 * To be absolutely certain we don't mess up the nesting level,
+		 * evaluate the bump_level condition just once.
+		 */
+		bool	bump_level =
+				!IsA(parsetree, ExecuteStmt) &&
+				!IsA(parsetree, PrepareStmt) &&
+				!IsA(parsetree, DeallocateStmt);
+
+		if (bump_level)
+			nesting_level++;
+
+		PG_TRY();
+		{
 #if PG_VERSION_NUM >= 140000
 		if (prev_ProcessUtility)
 			prev_ProcessUtility(pstmt, queryString,
@@ -1222,6 +1280,13 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 									completionTag);
 #endif
 	}
+	PG_FINALLY();
+	{
+		if (bump_level)
+			nesting_level--;	
+	}
+	PG_END_TRY();
+	}
 }
 
 #if PG_VERSION_NUM < 130000
@@ -1464,14 +1529,14 @@ pgsm_update_entry(pgsmEntry * entry,
 			e->counters.info.num_relations = num_relations;
 			_snprintf2(e->counters.info.relations, relations, num_relations, REL_LEN);
 
-			if (exec_nested_level > 0 && e->counters.info.parentid == 0 && pgsm_track == PGSM_TRACK_ALL)
+			if (nesting_level > 0 && e->counters.info.parentid == 0 && pgsm_track == PGSM_TRACK_ALL)
 			{
-				if (exec_nested_level >= 0 && exec_nested_level < max_stack_depth)
+				if (nesting_level >= 0 && nesting_level < max_stack_depth)
 				{
-					int			parent_query_len = nested_query_txts[exec_nested_level - 1] ?
+					int			parent_query_len = nested_query_txts[nesting_level - 1] ?
-					strlen(nested_query_txts[exec_nested_level - 1]) : 0;
+					strlen(nested_query_txts[nesting_level - 1]) : 0;
 
-					e->counters.info.parentid = nested_queryids[exec_nested_level - 1];
+					e->counters.info.parentid = nested_queryids[nesting_level - 1];
 					e->counters.info.parent_query = InvalidDsaPointer;
 					/* If we have a parent query, store it in the raw dsa area */
 					if (parent_query_len > 0)
@@ -1489,7 +1554,7 @@ pgsm_update_entry(pgsmEntry * entry,
 						if (DsaPointerIsValid(qry))
 						{
 							qry_buff = dsa_get_address(query_dsa_area, qry);
-							memcpy(qry_buff, nested_query_txts[exec_nested_level - 1], parent_query_len);
+							memcpy(qry_buff, nested_query_txts[nesting_level - 1], parent_query_len);
 							qry_buff[parent_query_len] = 0;
 							/* store the dsa pointer for parent query text */
 							e->counters.info.parent_query = qry;
@@ -1750,7 +1815,7 @@ pgsm_create_hash_entry(uint64 bucket_id, uint64 queryid, PlanInfo * plan_info)
 #if PG_VERSION_NUM < 140000
 	entry->key.toplevel = 1;
 #else
-	entry->key.toplevel = ((exec_nested_level + plan_nested_level) == 0);
+	entry->key.toplevel = ((nesting_level) == 0);
 #endif
 
 	if (IsTransactionState())

regression/expected/level_tracking.out

@@ -69,6 +69,61 @@ SELECT toplevel, calls, query FROM pg_stat_monitor
  t        |     1 | SET pg_stat_monitor.pgsm_track = 'all'
 (7 rows)
 
+-- DO block - top-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'top';
+SET pg_stat_monitor.pgsm_track_utility = FALSE;
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |             query              
+----------+-------+--------------------------------
+ t        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_monitor_reset()
+(2 rows)
+
+-- DO block - all-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'all';
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |             query              
+----------+-------+--------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | SELECT $1::TEXT
+ t        |     1 | SELECT pg_stat_monitor_reset()
+(4 rows)
+
 -- PL/pgSQL function - top-level tracking.
 SET pg_stat_monitor.pgsm_track = 'top';
 SET pg_stat_monitor.pgsm_track_utility = FALSE;
@@ -120,6 +175,31 @@ SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
      1 |    1 | SELECT pg_stat_monitor_reset()
 (3 rows)
 
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                   query                                   
+----------+-------+------+---------------------------------------------------------------------------
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    3 | SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C"
+ t        |     1 |    1 | SELECT pg_stat_monitor_reset()
+(5 rows)
+
 -- PL/pgSQL function - all-level tracking.
 SET pg_stat_monitor.pgsm_track = 'all';
 SELECT pg_stat_monitor_reset();
@@ -131,6 +211,7 @@ SELECT pg_stat_monitor_reset();
 -- we drop and recreate the functions to avoid any caching funnies
 DROP FUNCTION PLUS_ONE(INTEGER);
 DROP FUNCTION PLUS_TWO(INTEGER);
+DROP FUNCTION PLUS_THREE(INTEGER);
 -- PL/pgSQL function
 CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
 DECLARE
@@ -176,7 +257,34 @@ SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
      1 |    1 | SELECT pg_stat_monitor_reset()
 (5 rows)
 
-DROP FUNCTION PLUS_ONE(INTEGER);
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                   query                                   
+----------+-------+------+---------------------------------------------------------------------------
+ f        |     2 |    2 | SELECT (i + $2 + $3)::INTEGER
+ f        |     2 |    2 | SELECT (i + $2)::INTEGER LIMIT $3
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    5 | SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C"
+ f        |     2 |    2 | SELECT i + $2 LIMIT $3
+ t        |     1 |    1 | SELECT pg_stat_monitor_reset()
+(8 rows)
+
 --
 -- pg_stat_monitor.pgsm_track = none
 --

regression/expected/level_tracking_1.out

@@ -0,0 +1,322 @@
+--
+-- Statement level tracking
+--
+CREATE EXTENSION pg_stat_monitor;
+SET pg_stat_monitor.pgsm_track_utility = TRUE;
+SET pg_stat_monitor.pgsm_normalized_query = TRUE;
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+-- DO block - top-level tracking.
+CREATE TABLE stats_track_tab (x int);
+SET pg_stat_monitor.pgsm_track = 'top';
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END;
+$$ LANGUAGE plpgsql;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  WHERE query LIKE '%DELETE%' ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |             query              
+----------+-------+--------------------------------
+ t        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | DO $$                         +
+          |       | BEGIN                         +
+          |       |   DELETE FROM stats_track_tab;+
+          |       | END;                          +
+          |       | $$ LANGUAGE plpgsql
+(2 rows)
+
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+-- DO block - all-level tracking.
+SET pg_stat_monitor.pgsm_track = 'all';
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                 query                  
+----------+-------+----------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | DO $$                                 +
+          |       | BEGIN                                 +
+          |       |   DELETE FROM stats_track_tab;        +
+          |       | END; $$
+ t        |     1 | DO LANGUAGE plpgsql $$                +
+          |       | BEGIN                                 +
+          |       |   -- this is a SELECT                 +
+          |       |   PERFORM 'hello world'::TEXT;        +
+          |       | END; $$
+ f        |     1 | SELECT $1::TEXT
+ t        |     1 | SELECT pg_stat_monitor_reset()
+ t        |     1 | SET pg_stat_monitor.pgsm_track = 'all'
+(7 rows)
+
+-- DO block - top-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'top';
+SET pg_stat_monitor.pgsm_track_utility = FALSE;
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |             query              
+----------+-------+--------------------------------
+ t        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_monitor_reset()
+(2 rows)
+
+-- DO block - all-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'all';
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |             query              
+----------+-------+--------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ t        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | SELECT $1::TEXT
+ t        |     1 | SELECT pg_stat_monitor_reset()
+(4 rows)
+
+-- PL/pgSQL function - top-level tracking.
+SET pg_stat_monitor.pgsm_track = 'top';
+SET pg_stat_monitor.pgsm_track_utility = FALSE;
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i + 1 + 1.0)::INTEGER INTO r;
+  RETURN r;
+END; $$ LANGUAGE plpgsql;
+SELECT PLUS_TWO(3);
+ plus_two 
+----------
+        5
+(1 row)
+
+SELECT PLUS_TWO(7);
+ plus_two 
+----------
+        9
+(1 row)
+
+-- SQL function --- use LIMIT to keep it from being inlined
+CREATE FUNCTION PLUS_ONE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT (i + 1.0)::INTEGER LIMIT 1 $$ LANGUAGE SQL;
+SELECT PLUS_ONE(8);
+ plus_one 
+----------
+        9
+(1 row)
+
+SELECT PLUS_ONE(10);
+ plus_one 
+----------
+       11
+(1 row)
+
+SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ calls | rows |             query              
+-------+------+--------------------------------
+     2 |    2 | SELECT PLUS_ONE($1)
+     2 |    2 | SELECT PLUS_TWO($1)
+     1 |    1 | SELECT pg_stat_monitor_reset()
+(3 rows)
+
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                   query                                   
+----------+-------+------+---------------------------------------------------------------------------
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    3 | SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C"
+ t        |     2 |    2 | SELECT i + $2 LIMIT $3
+ t        |     1 |    1 | SELECT pg_stat_monitor_reset()
+(6 rows)
+
+-- PL/pgSQL function - all-level tracking.
+SET pg_stat_monitor.pgsm_track = 'all';
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+-- we drop and recreate the functions to avoid any caching funnies
+DROP FUNCTION PLUS_ONE(INTEGER);
+DROP FUNCTION PLUS_TWO(INTEGER);
+DROP FUNCTION PLUS_THREE(INTEGER);
+-- PL/pgSQL function
+CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
+DECLARE
+  r INTEGER;
+BEGIN
+  SELECT (i + 1 + 1.0)::INTEGER INTO r;
+  RETURN r;
+END; $$ LANGUAGE plpgsql;
+SELECT PLUS_TWO(-1);
+ plus_two 
+----------
+        1
+(1 row)
+
+SELECT PLUS_TWO(2);
+ plus_two 
+----------
+        4
+(1 row)
+
+-- SQL function --- use LIMIT to keep it from being inlined
+CREATE FUNCTION PLUS_ONE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT (i + 1.0)::INTEGER LIMIT 1 $$ LANGUAGE SQL;
+SELECT PLUS_ONE(3);
+ plus_one 
+----------
+        4
+(1 row)
+
+SELECT PLUS_ONE(1);
+ plus_one 
+----------
+        2
+(1 row)
+
+SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ calls | rows |               query               
+-------+------+-----------------------------------
+     2 |    2 | SELECT (i + $2 + $3)::INTEGER
+     2 |    2 | SELECT (i + $2)::INTEGER LIMIT $3
+     2 |    2 | SELECT PLUS_ONE($1)
+     2 |    2 | SELECT PLUS_TWO($1)
+     1 |    1 | SELECT pg_stat_monitor_reset()
+(5 rows)
+
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                   query                                   
+----------+-------+------+---------------------------------------------------------------------------
+ f        |     2 |    2 | SELECT (i + $2 + $3)::INTEGER
+ f        |     2 |    2 | SELECT (i + $2)::INTEGER LIMIT $3
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    5 | SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C"
+ t        |     2 |    2 | SELECT i + $2 LIMIT $3
+ t        |     1 |    1 | SELECT pg_stat_monitor_reset()
+(8 rows)
+
+--
+-- pg_stat_monitor.pgsm_track = none
+--
+SET pg_stat_monitor.pgsm_track = 'none';
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+SELECT 1 AS "one";
+ one 
+-----
+   1
+(1 row)
+
+SELECT 1 + 1 AS "two";
+ two 
+-----
+   2
+(1 row)
+
+SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+ calls | rows | query 
+-------+------+-------
+(0 rows)
+
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+DROP EXTENSION pg_stat_monitor;

regression/expected/pgsm_query_id_1.out

@@ -0,0 +1,115 @@
+CREATE EXTENSION pg_stat_monitor;
+CREATE DATABASE db1;
+CREATE DATABASE db2;
+\c db1
+CREATE TABLE t1 (a int);
+CREATE TABLE t2 (b int);
+CREATE FUNCTION add(integer, integer) RETURNS integer
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+\c db2
+CREATE TABLE t1 (a int);
+CREATE TABLE t3 (c int);
+CREATE FUNCTION add(integer, integer) RETURNS integer
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+\c contrib_regression
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+\c db1
+SELECT * FROM t1;
+ a 
+---
+(0 rows)
+
+SELECT *, ADD(1, 2) FROM t1;
+ a | add 
+---+-----
+(0 rows)
+
+SELECT * FROM t2;
+ b 
+---
+(0 rows)
+
+-- Check that spaces and comments do not generate a different pgsm_query_id
+SELECT     *     FROM t2 --WHATEVER;
+;
+ b 
+---
+(0 rows)
+
+SELECT     *     FROM t2 /* ...
+...
+More comments to check for spaces.
+*/
+     ;
+ b 
+---
+(0 rows)
+
+\c db2
+SELECT * FROM t1;
+ a 
+---
+(0 rows)
+
+SELECT *, ADD(1, 2) FROM t1;
+ a | add 
+---+-----
+(0 rows)
+
+set pg_stat_monitor.pgsm_enable_pgsm_query_id = off;
+SELECT * FROM t3;
+ c 
+---
+(0 rows)
+
+set pg_stat_monitor.pgsm_enable_pgsm_query_id = on;
+SELECT * FROM t3 where c = 20;
+ c 
+---
+(0 rows)
+
+\c contrib_regression
+SELECT datname, pgsm_query_id, query, calls FROM pg_stat_monitor ORDER BY pgsm_query_id, query, datname;
+      datname       |    pgsm_query_id    |                        query                        | calls 
+--------------------+---------------------+-----------------------------------------------------+-------
+ contrib_regression |  689150021118383254 | SELECT pg_stat_monitor_reset()                      |     1
+ db1                | 1897482803466821995 | SELECT * FROM t2                                    |     3
+ db1                | 1988437669671417938 | SELECT * FROM t1                                    |     1
+ db2                | 1988437669671417938 | SELECT * FROM t1                                    |     1
+ db2                | 6220142855706866455 | set pg_stat_monitor.pgsm_enable_pgsm_query_id = on  |     1
+ db2                | 6633979598391393345 | SELECT * FROM t3 where c = 20                       |     1
+ db1                | 8140395000078788481 | SELECT *, ADD(1, 2) FROM t1                         |     1
+ db2                | 8140395000078788481 | SELECT *, ADD(1, 2) FROM t1                         |     1
+ db2                |                     | SELECT * FROM t3                                    |     1
+ db2                |                     | set pg_stat_monitor.pgsm_enable_pgsm_query_id = off |     1
+(10 rows)
+
+SELECT pg_stat_monitor_reset();
+ pg_stat_monitor_reset 
+-----------------------
+ 
+(1 row)
+
+\c db1
+DROP TABLE t1;
+DROP TABLE t2;
+DROP FUNCTION ADD;
+\c db2
+DROP TABLE t1;
+DROP TABLE t3;
+DROP FUNCTION ADD;
+\c contrib_regression
+DROP DATABASE db1;
+DROP DATABASE db2;
+DROP EXTENSION pg_stat_monitor;

regression/sql/level_tracking.sql

@@ -34,6 +34,39 @@ END; $$;
 SELECT toplevel, calls, query FROM pg_stat_monitor
   ORDER BY query COLLATE "C", toplevel;
 
+-- DO block - top-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'top';
+SET pg_stat_monitor.pgsm_track_utility = FALSE;
+SELECT pg_stat_monitor_reset();
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+
+-- DO block - all-level tracking without utility.
+SET pg_stat_monitor.pgsm_track = 'all';
+SELECT pg_stat_monitor_reset();
+DELETE FROM stats_track_tab;
+DO $$
+BEGIN
+  DELETE FROM stats_track_tab;
+END; $$;
+DO LANGUAGE plpgsql $$
+BEGIN
+  -- this is a SELECT
+  PERFORM 'hello world'::TEXT;
+END; $$;
+SELECT toplevel, calls, query FROM pg_stat_monitor
+  ORDER BY query COLLATE "C", toplevel;
+
 -- PL/pgSQL function - top-level tracking.
 SET pg_stat_monitor.pgsm_track = 'top';
 SET pg_stat_monitor.pgsm_track_utility = FALSE;
@@ -58,6 +91,15 @@ SELECT PLUS_ONE(10);
 
 SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
 
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+
+SELECT PLUS_THREE(8);
+SELECT PLUS_THREE(10);
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
+
 -- PL/pgSQL function - all-level tracking.
 SET pg_stat_monitor.pgsm_track = 'all';
 SELECT pg_stat_monitor_reset();
@@ -65,6 +107,7 @@ SELECT pg_stat_monitor_reset();
 -- we drop and recreate the functions to avoid any caching funnies
 DROP FUNCTION PLUS_ONE(INTEGER);
 DROP FUNCTION PLUS_TWO(INTEGER);
+DROP FUNCTION PLUS_THREE(INTEGER);
 
 -- PL/pgSQL function
 CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
@@ -86,7 +129,15 @@ SELECT PLUS_ONE(3);
 SELECT PLUS_ONE(1);
 
 SELECT calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
-DROP FUNCTION PLUS_ONE(INTEGER);
+
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+
+SELECT PLUS_THREE(8);
+SELECT PLUS_THREE(10);
+
+SELECT toplevel, calls, rows, query FROM pg_stat_monitor ORDER BY query COLLATE "C";
 
 --
 -- pg_stat_monitor.pgsm_track = none