Add some sanity checks in executor for query ID reporting

michaelpq · michaelpq · commit 24f520594809 · 2024-09-18T14:43:37.000+09:00
This commit adds three sanity checks in code paths of the executor where it is possible to use hooks, checking that a query ID is reported in pg_stat_activity if compute_query_id is enabled: - ExecutorRun() - ExecutorFinish() - ExecutorEnd() This causes the test in pg_stat_statements added in 933848d to complain immediately in ExecutorRun(). The idea behind this commit is to help extensions to detect if they are missing query ID reports when a query goes through the executor. Perhaps this will prove to be a bad idea, but let's see where this experience goes in v18 and newer versions. Reviewed-by: Sami Imseih Discussion: https://postgr.es/m/ZuJb5xCKHH0A9tMN@paquier.xyz
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
@@ -50,6 +50,7 @@
 #include "foreign/fdwapi.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
+#include "nodes/queryjumble.h"
 #include "parser/parse_relation.h"
 #include "rewrite/rewriteHandler.h"
 #include "tcop/utility.h"
@@ -295,6 +296,9 @@ ExecutorRun(QueryDesc *queryDesc,
 			ScanDirection direction, uint64 count,
 			bool execute_once)
 {
+	/* If enabled, the query ID should be set. */
+	Assert(!IsQueryIdEnabled() || pgstat_get_my_query_id() != 0);
+
 	if (ExecutorRun_hook)
 		(*ExecutorRun_hook) (queryDesc, direction, count, execute_once);
 	else
@@ -403,6 +407,9 @@ standard_ExecutorRun(QueryDesc *queryDesc,
 void
 ExecutorFinish(QueryDesc *queryDesc)
 {
+	/* If enabled, the query ID should be set. */
+	Assert(!IsQueryIdEnabled() || pgstat_get_my_query_id() != 0);
+
 	if (ExecutorFinish_hook)
 		(*ExecutorFinish_hook) (queryDesc);
 	else
@@ -463,6 +470,9 @@ standard_ExecutorFinish(QueryDesc *queryDesc)
 void
 ExecutorEnd(QueryDesc *queryDesc)
 {
+	/* If enabled, the query ID should be set. */
+	Assert(!IsQueryIdEnabled() || pgstat_get_my_query_id() != 0);
+
 	if (ExecutorEnd_hook)
 		(*ExecutorEnd_hook) (queryDesc);
 	else
