Add pg_analyze_and_rewrite_varparams()

This new function extracts common code from PrepareQuery() and
exec_parse_message().  It is then exactly analogous to the existing
pg_analyze_and_rewrite_fixedparams() and
pg_analyze_and_rewrite_withcb().

To unify these two code paths, this makes PrepareQuery() now subject
to log_parser_stats.  Also, both paths now invoke
TRACE_POSTGRESQL_QUERY_REWRITE_START().  PrepareQuery() no longer
checks whether a utility statement was specified.  The grammar doesn't
allow that anyway, and exec_parse_message() supports it, so
restricting it doesn't seem necessary.

This also adds QueryEnvironment support to the *varparams functions,
for consistency with its cousins, even though it is not used right
now.

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

src/backend/commands/prepare.c

@@ -63,9 +63,7 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 	CachedPlanSource *plansource;
 	Oid		   *argtypes = NULL;
 	int			nargs;
-	Query	   *query;
 	List	   *query_list;
-	int			i;
 
 	/*
 	 * Disallow empty-string statement name (conflicts with protocol-level
@@ -97,6 +95,7 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 
 	if (nargs)
 	{
+		int			i;
 		ListCell   *l;
 
 		argtypes = (Oid *) palloc(nargs * sizeof(Oid));
@@ -115,44 +114,10 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 	 * Analyze the statement using these parameter types (any parameters
 	 * passed in from above us will not be visible to it), allowing
 	 * information about unknown parameters to be deduced from context.
+	 * Rewrite the query. The result could be 0, 1, or many queries.
 	 */
-	query = parse_analyze_varparams(rawstmt, pstate->p_sourcetext,
+	query_list = pg_analyze_and_rewrite_varparams(rawstmt, pstate->p_sourcetext,
-									&argtypes, &nargs);
+												  &argtypes, &nargs, NULL);
-
-	/*
-	 * Check that all parameter types were determined.
-	 */
-	for (i = 0; i < nargs; i++)
-	{
-		Oid			argtype = argtypes[i];
-
-		if (argtype == InvalidOid || argtype == UNKNOWNOID)
-			ereport(ERROR,
-					(errcode(ERRCODE_INDETERMINATE_DATATYPE),
-					 errmsg("could not determine data type of parameter $%d",
-							i + 1)));
-	}
-
-	/*
-	 * grammar only allows PreparableStmt, so this check should be redundant
-	 */
-	switch (query->commandType)
-	{
-		case CMD_SELECT:
-		case CMD_INSERT:
-		case CMD_UPDATE:
-		case CMD_DELETE:
-			/* OK */
-			break;
-		default:
-			ereport(ERROR,
-					(errcode(ERRCODE_INVALID_PSTATEMENT_DEFINITION),
-					 errmsg("utility statements cannot be prepared")));
-			break;
-	}
-
-	/* Rewrite the query. The result could be 0, 1, or many queries. */
-	query_list = QueryRewrite(query);
 
 	/* Finish filling in the CachedPlanSource */
 	CompleteCachedPlan(plansource,

src/backend/parser/analyze.c

@@ -148,7 +148,8 @@ parse_analyze_fixedparams(RawStmt *parseTree, const char *sourceText,
  */
 Query *
 parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
-						Oid **paramTypes, int *numParams)
+						Oid **paramTypes, int *numParams,
+						QueryEnvironment *queryEnv)
 {
 	ParseState *pstate = make_parsestate(NULL);
 	Query	   *query;
@@ -160,6 +161,8 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 
 	setup_parse_variable_parameters(pstate, paramTypes, numParams);
 
+	pstate->p_queryEnv = queryEnv;
+
 	query = transformTopLevelStmt(pstate, parseTree);
 
 	/* make sure all is well with parameter types */

src/backend/tcop/postgres.c

@@ -637,9 +637,11 @@ pg_parse_query(const char *query_string)
  * NOTE: for reasons mentioned above, this must be separate from raw parsing.
  */
 List *
-pg_analyze_and_rewrite_fixedparams(RawStmt *parsetree, const char *query_string,
+pg_analyze_and_rewrite_fixedparams(RawStmt *parsetree,
-					   const Oid *paramTypes, int numParams,
+								   const char *query_string,
-					   QueryEnvironment *queryEnv)
+								   const Oid *paramTypes,
+								   int numParams,
+								   QueryEnvironment *queryEnv)
 {
 	Query	   *query;
 	List	   *querytree_list;
@@ -668,6 +670,59 @@ pg_analyze_and_rewrite_fixedparams(RawStmt *parsetree, const char *query_string,
 	return querytree_list;
 }
 
+/*
+ * Do parse analysis and rewriting.  This is the same as
+ * pg_analyze_and_rewrite_fixedparams except that it's okay to deduce
+ * information about $n symbol datatypes from context.
+ */
+List *
+pg_analyze_and_rewrite_varparams(RawStmt *parsetree,
+								 const char *query_string,
+								 Oid **paramTypes,
+								 int *numParams,
+								 QueryEnvironment *queryEnv)
+{
+	Query	   *query;
+	List	   *querytree_list;
+
+	TRACE_POSTGRESQL_QUERY_REWRITE_START(query_string);
+
+	/*
+	 * (1) Perform parse analysis.
+	 */
+	if (log_parser_stats)
+		ResetUsage();
+
+	query = parse_analyze_varparams(parsetree, query_string, paramTypes, numParams,
+						  queryEnv);
+
+	/*
+	 * Check all parameter types got determined.
+	 */
+	for (int i = 0; i < *numParams; i++)
+	{
+		Oid			ptype = (*paramTypes)[i];
+
+		if (ptype == InvalidOid || ptype == UNKNOWNOID)
+			ereport(ERROR,
+					(errcode(ERRCODE_INDETERMINATE_DATATYPE),
+					 errmsg("could not determine data type of parameter $%d",
+							i + 1)));
+	}
+
+	if (log_parser_stats)
+		ShowUsage("PARSE ANALYSIS STATISTICS");
+
+	/*
+	 * (2) Rewrite the queries, as necessary
+	 */
+	querytree_list = pg_rewrite_query(query);
+
+	TRACE_POSTGRESQL_QUERY_REWRITE_DONE(query_string);
+
+	return querytree_list;
+}
+
 /*
  * Do parse analysis and rewriting.  This is the same as
  * pg_analyze_and_rewrite_fixedparams except that, instead of a fixed list of
@@ -1409,7 +1464,6 @@ exec_parse_message(const char *query_string,	/* string to execute */
 
 	if (parsetree_list != NIL)
 	{
-		Query	   *query;
 		bool		snapshot_set = false;
 
 		raw_parse_tree = linitial_node(RawStmt, parsetree_list);
@@ -1449,34 +1503,13 @@ exec_parse_message(const char *query_string,	/* string to execute */
 		/*
 		 * Analyze and rewrite the query.  Note that the originally specified
 		 * parameter set is not required to be complete, so we have to use
-		 * parse_analyze_varparams().
+		 * pg_analyze_and_rewrite_varparams().
-		 */
-		if (log_parser_stats)
-			ResetUsage();
-
-		query = parse_analyze_varparams(raw_parse_tree,
-										query_string,
-										&paramTypes,
-										&numParams);
-
-		/*
-		 * Check all parameter types got determined.
 		 */
-		for (int i = 0; i < numParams; i++)
+		querytree_list = pg_analyze_and_rewrite_varparams(raw_parse_tree,
-		{
+														  query_string,
-			Oid			ptype = paramTypes[i];
+														  &paramTypes,
-
+														  &numParams,
-			if (ptype == InvalidOid || ptype == UNKNOWNOID)
+														  NULL);
-				ereport(ERROR,
-						(errcode(ERRCODE_INDETERMINATE_DATATYPE),
-						 errmsg("could not determine data type of parameter $%d",
-								i + 1)));
-		}
-
-		if (log_parser_stats)
-			ShowUsage("PARSE ANALYSIS STATISTICS");
-
-		querytree_list = pg_rewrite_query(query);
 
 		/* Done with the snapshot used for parsing */
 		if (snapshot_set)

src/include/parser/analyze.h

@@ -27,7 +27,7 @@ extern PGDLLIMPORT post_parse_analyze_hook_type post_parse_analyze_hook;
 extern Query *parse_analyze_fixedparams(RawStmt *parseTree, const char *sourceText,
 							const Oid *paramTypes, int numParams, QueryEnvironment *queryEnv);
 extern Query *parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
-									  Oid **paramTypes, int *numParams);
+									  Oid **paramTypes, int *numParams, QueryEnvironment *queryEnv);
 
 extern Query *parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 								CommonTableExpr *parentCTE,

src/include/tcop/tcopprot.h

@@ -49,6 +49,11 @@ extern List *pg_analyze_and_rewrite_fixedparams(RawStmt *parsetree,
 									const char *query_string,
 									const Oid *paramTypes, int numParams,
 									QueryEnvironment *queryEnv);
+extern List *pg_analyze_and_rewrite_varparams(RawStmt *parsetree,
+											  const char *query_string,
+											  Oid **paramTypes,
+											  int *numParams,
+											  QueryEnvironment *queryEnv);
 extern List *pg_analyze_and_rewrite_withcb(RawStmt *parsetree,
 										   const char *query_string,
 										   ParserSetupHook parserSetup,