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().

Reviewed-by: Nathan Bossart <bossartn@amazon.com>
Discussion: https://www.postgresql.org/message-id/flat/c67ce276-52b4-0239-dc0e-39875bf81840@enterprisedb.com

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.

src/backend/tcop/postgres.c

@@ -736,12 +736,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);
 
@@ -751,22 +747,8 @@ pg_analyze_and_rewrite_withcb(RawStmt *parsetree,
 	if (log_parser_stats)
 		ResetUsage();
 
-	pstate = make_parsestate(NULL);
+	query = parse_analyze_withcb(parsetree, query_string, parserSetup, parserSetupArg,
-	pstate->p_sourcetext = query_string;
+								 queryEnv);
-	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);
 
 	if (log_parser_stats)
 		ShowUsage("PARSE ANALYSIS STATISTICS");

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,