From 9baad5def3d9a4b442a591b9f6286044c354f89d Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter@eisentraut.org>
Date: Tue, 28 Dec 2021 09:01:45 +0100
Subject: [PATCH v1 3/3] Add parse_analyze_withcb()

This extracts code from pg_analyze_and_rewrite_withcb() into a
separate function that mirrors the existing
parse_analyze_fixedparams() and parse_analyze_varparams().
---
 src/backend/parser/analyze.c | 38 ++++++++++++++++++++++++++++++++++++
 src/backend/tcop/postgres.c  | 22 ++-------------------
 src/include/parser/analyze.h |  5 +++++
 3 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 4a9d212b22..6f4cc1b9db 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -181,6 +181,44 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 	return query;
 }
 
+/*
+ * parse_analyze_withcb
+ *
+ * This variant is used when the caller supplies their own parser callback to
+ * resolve parameters and possibly other things.
+ */
+Query *
+parse_analyze_withcb(RawStmt *parseTree, const char *sourceText,
+					 ParserSetupHook parserSetup,
+					 void *parserSetupArg,
+					 QueryEnvironment *queryEnv)
+{
+	ParseState *pstate = make_parsestate(NULL);
+	Query	   *query;
+	JumbleState *jstate = NULL;
+
+	Assert(sourceText != NULL); /* required as of 8.4 */
+
+	pstate->p_sourcetext = sourceText;
+	pstate->p_queryEnv = queryEnv;
+	(*parserSetup) (pstate, parserSetupArg);
+
+	query = transformTopLevelStmt(pstate, parseTree);
+
+	if (IsQueryIdEnabled())
+		jstate = JumbleQuery(query, sourceText);
+
+	if (post_parse_analyze_hook)
+		(*post_parse_analyze_hook) (pstate, query, jstate);
+
+	free_parsestate(pstate);
+
+	pgstat_report_query_id(query->queryId, false);
+
+	return query;
+}
+
+
 /*
  * parse_sub_analyze
  *		Entry point for recursively analyzing a sub-statement.
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 8ade13994f..a646e34d4c 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -739,12 +739,8 @@ pg_analyze_and_rewrite_withcb(RawStmt *parsetree,
 							  void *parserSetupArg,
 							  QueryEnvironment *queryEnv)
 {
-	ParseState *pstate;
 	Query	   *query;
 	List	   *querytree_list;
-	JumbleState *jstate = NULL;
-
-	Assert(query_string != NULL);	/* required as of 8.4 */
 
 	TRACE_POSTGRESQL_QUERY_REWRITE_START(query_string);
 
@@ -754,22 +750,8 @@ pg_analyze_and_rewrite_withcb(RawStmt *parsetree,
 	if (log_parser_stats)
 		ResetUsage();
 
-	pstate = make_parsestate(NULL);
-	pstate->p_sourcetext = query_string;
-	pstate->p_queryEnv = queryEnv;
-	(*parserSetup) (pstate, parserSetupArg);
-
-	query = transformTopLevelStmt(pstate, parsetree);
-
-	if (IsQueryIdEnabled())
-		jstate = JumbleQuery(query, query_string);
-
-	if (post_parse_analyze_hook)
-		(*post_parse_analyze_hook) (pstate, query, jstate);
-
-	free_parsestate(pstate);
-
-	pgstat_report_query_id(query->queryId, false);
+	query = parse_analyze_withcb(parsetree, query_string, parserSetup, parserSetupArg,
+								 queryEnv);
 
 	if (log_parser_stats)
 		ShowUsage("PARSE ANALYSIS STATISTICS");
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index b1c592c6c6..423b956869 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -14,6 +14,7 @@
 #ifndef ANALYZE_H
 #define ANALYZE_H
 
+#include "nodes/params.h"
 #include "parser/parse_node.h"
 #include "utils/queryjumble.h"
 
@@ -28,6 +29,10 @@ extern Query *parse_analyze_fixedparams(RawStmt *parseTree, const char *sourceTe
 							const Oid *paramTypes, int numParams, QueryEnvironment *queryEnv);
 extern Query *parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 									  Oid **paramTypes, int *numParams, QueryEnvironment *queryEnv);
+extern Query *parse_analyze_withcb(RawStmt *parseTree, const char *sourceText,
+								   ParserSetupHook parserSetup,
+								   void *parserSetupArg,
+								   QueryEnvironment *queryEnv);
 
 extern Query *parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 								CommonTableExpr *parentCTE,
-- 
2.34.1