From eaab4d7c088ff3ee9b0e6ec3de96677bafd184c0 Mon Sep 17 00:00:00 2001
From: Mats Kindahl <mats@timescale.com>
Date: Mon, 4 Mar 2024 12:38:05 +0100
Subject: Improve support for ExplainOneQuery hook

There is a hook ExplainOneQuery_hook, but there is no corresponding
standard_ExplainOneQuery and the code that belongs there is written
in-line in ExplainOneQuery(). As a result, anybody adding a hook for
ExplainOneQuery_hook has to copy the code from ExplainOneQuery() to run
the standard ExplainOneQuery before adding their own messages.

This commit fixes this by refactoring ExplainOneQuery() and factor out
the standard code into standard_ExplainOneQuery() and call it from
ExplainOneQuery() in a similar manner to how it is done for other
hooks.
---
 src/backend/commands/explain.c | 106 ++++++++++++++++++---------------
 src/include/commands/explain.h |   4 ++
 2 files changed, 61 insertions(+), 49 deletions(-)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 78754bc6ba..3b7bed3ca2 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -391,63 +391,71 @@ ExplainOneQuery(Query *query, int cursorOptions,
 
 	/* if an advisor plugin is present, let it manage things */
 	if (ExplainOneQuery_hook)
-		(*ExplainOneQuery_hook) (query, cursorOptions, into, es,
-								 queryString, params, queryEnv);
+		(*ExplainOneQuery_hook)(query, cursorOptions, into, es,
+							 queryString, params, queryEnv);
 	else
-	{
-		PlannedStmt *plan;
-		instr_time	planstart,
-					planduration;
-		BufferUsage bufusage_start,
-					bufusage;
-		MemoryContextCounters mem_counters;
-		MemoryContext planner_ctx = NULL;
-		MemoryContext saved_ctx = NULL;
-
-		if (es->memory)
-		{
-			/*
-			 * Create a new memory context to measure planner's memory
-			 * consumption accurately.  Note that if the planner were to be
-			 * modified to use a different memory context type, here we would
-			 * be changing that to AllocSet, which might be undesirable.
-			 * However, we don't have a way to create a context of the same
-			 * type as another, so we pray and hope that this is OK.
-			 */
-			planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
-												"explain analyze planner context",
-												ALLOCSET_DEFAULT_SIZES);
-			saved_ctx = MemoryContextSwitchTo(planner_ctx);
-		}
+		standard_ExplainOneQuery(query, cursorOptions, into, es,
+								 queryString, params, queryEnv);
+}
 
-		if (es->buffers)
-			bufusage_start = pgBufferUsage;
-		INSTR_TIME_SET_CURRENT(planstart);
+void
+standard_ExplainOneQuery(Query *query, int cursorOptions,
+						 IntoClause *into, ExplainState *es,
+						 const char *queryString, ParamListInfo params,
+						 QueryEnvironment *queryEnv)
+{
+	PlannedStmt *plan;
+	instr_time	planstart,
+				planduration;
+	BufferUsage bufusage_start,
+				bufusage;
+	MemoryContextCounters mem_counters;
+	MemoryContext planner_ctx = NULL;
+	MemoryContext saved_ctx = NULL;
+
+	if (es->memory)
+	{
+		/*
+		 * Create a new memory context to measure planner's memory consumption
+		 * accurately.  Note that if the planner were to be modified to use a
+		 * different memory context type, here we would be changing that to
+		 * AllocSet, which might be undesirable.  However, we don't have a way
+		 * to create a context of the same type as another, so we pray and
+		 * hope that this is OK.
+		 */
+		planner_ctx = AllocSetContextCreate(CurrentMemoryContext,
+											"explain analyze planner context",
+											ALLOCSET_DEFAULT_SIZES);
+		saved_ctx = MemoryContextSwitchTo(planner_ctx);
+	}
 
-		/* plan the query */
-		plan = pg_plan_query(query, queryString, cursorOptions, params);
+	if (es->buffers)
+		bufusage_start = pgBufferUsage;
+	INSTR_TIME_SET_CURRENT(planstart);
 
-		INSTR_TIME_SET_CURRENT(planduration);
-		INSTR_TIME_SUBTRACT(planduration, planstart);
+	/* plan the query */
+	plan = pg_plan_query(query, queryString, cursorOptions, params);
 
-		if (es->memory)
-		{
-			MemoryContextSwitchTo(saved_ctx);
-			MemoryContextMemConsumed(planner_ctx, &mem_counters);
-		}
+	INSTR_TIME_SET_CURRENT(planduration);
+	INSTR_TIME_SUBTRACT(planduration, planstart);
 
-		/* calc differences of buffer counters. */
-		if (es->buffers)
-		{
-			memset(&bufusage, 0, sizeof(BufferUsage));
-			BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
-		}
+	if (es->memory)
+	{
+		MemoryContextSwitchTo(saved_ctx);
+		MemoryContextMemConsumed(planner_ctx, &mem_counters);
+	}
 
-		/* run it (if needed) and produce output */
-		ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
-					   &planduration, (es->buffers ? &bufusage : NULL),
-					   es->memory ? &mem_counters : NULL);
+	/* calc differences of buffer counters. */
+	if (es->buffers)
+	{
+		memset(&bufusage, 0, sizeof(BufferUsage));
+		BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start);
 	}
+
+	/* run it (if needed) and produce output */
+	ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
+				   &planduration, (es->buffers ? &bufusage : NULL),
+				   es->memory ? &mem_counters : NULL);
 }
 
 /*
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 7c0f0b5636..cf195f1359 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -80,6 +80,10 @@ extern PGDLLIMPORT explain_get_index_name_hook_type explain_get_index_name_hook;
 
 extern void ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 						 ParamListInfo params, DestReceiver *dest);
+extern void standard_ExplainOneQuery(Query *query, int cursorOptions,
+									 IntoClause *into, ExplainState *es,
+									 const char *queryString, ParamListInfo params,
+									 QueryEnvironment *queryEnv);
 
 extern ExplainState *NewExplainState(void);
 
-- 
2.34.1

