Add some infrastructure for contrib/pg_stat_statements.

Add a queryId field to Query and PlannedStmt.  This is not used by the
core backend, except for being copied around at appropriate times.
It's meant to allow plug-ins to track a particular query forward from
parse analysis to execution.

The queryId is intentionally not dumped into stored rules (and hence this
commit doesn't bump catversion).  You could argue that choice either way,
but it seems better that stored rule strings not have any dependency
on plug-ins that might or might not be present.

Also, add a post_parse_analyze_hook that gets invoked at the end of
parse analysis (but only for top-level analysis of complete queries,
not cases such as analyzing a domain's default-value expression).
This is mainly meant to be used to compute and assign a queryId,
but it could have other applications.

Peter Geoghegan

11 files changed, 34 insertions, 0 deletions

diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index cf23b088724..33ee62f40d0 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -78,6 +78,7 @@ _copyPlannedStmt(const PlannedStmt *from)
 	PlannedStmt *newnode = makeNode(PlannedStmt);
 
 	COPY_SCALAR_FIELD(commandType);
+	COPY_SCALAR_FIELD(queryId);
 	COPY_SCALAR_FIELD(hasReturning);
 	COPY_SCALAR_FIELD(hasModifyingCTE);
 	COPY_SCALAR_FIELD(canSetTag);
@@ -2402,6 +2403,7 @@ _copyQuery(const Query *from)
 
 	COPY_SCALAR_FIELD(commandType);
 	COPY_SCALAR_FIELD(querySource);
+	COPY_SCALAR_FIELD(queryId);
 	COPY_SCALAR_FIELD(canSetTag);
 	COPY_NODE_FIELD(utilityStmt);
 	COPY_SCALAR_FIELD(resultRelation);
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
index 5e691f96f79..b749e9bbe30 100644
--- a/src/backend/nodes/equalfuncs.c
+++ b/src/backend/nodes/equalfuncs.c
@@ -897,6 +897,7 @@ _equalQuery(const Query *a, const Query *b)
 {
 	COMPARE_SCALAR_FIELD(commandType);
 	COMPARE_SCALAR_FIELD(querySource);
+	/* we intentionally ignore queryId, since it might not be set */
 	COMPARE_SCALAR_FIELD(canSetTag);
 	COMPARE_NODE_FIELD(utilityStmt);
 	COMPARE_SCALAR_FIELD(resultRelation);
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index e925434eb39..594b3fdea85 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -242,6 +242,7 @@ _outPlannedStmt(StringInfo str, const PlannedStmt *node)
 	WRITE_NODE_TYPE("PLANNEDSTMT");
 
 	WRITE_ENUM_FIELD(commandType, CmdType);
+	WRITE_UINT_FIELD(queryId);
 	WRITE_BOOL_FIELD(hasReturning);
 	WRITE_BOOL_FIELD(hasModifyingCTE);
 	WRITE_BOOL_FIELD(canSetTag);
@@ -2152,6 +2153,7 @@ _outQuery(StringInfo str, const Query *node)
 
 	WRITE_ENUM_FIELD(commandType, CmdType);
 	WRITE_ENUM_FIELD(querySource, QuerySource);
+	/* we intentionally do not print the queryId field */
 	WRITE_BOOL_FIELD(canSetTag);
 
 	/*
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 9b579560c5e..7960793641c 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -195,6 +195,7 @@ _readQuery(void)
 
 	READ_ENUM_FIELD(commandType, CmdType);
 	READ_ENUM_FIELD(querySource, QuerySource);
+	local_node->queryId = 0;			/* not saved in output format */
 	READ_BOOL_FIELD(canSetTag);
 	READ_NODE_FIELD(utilityStmt);
 	READ_INT_FIELD(resultRelation);
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 6b0541b9b59..dcf32c07444 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -225,6 +225,7 @@ standard_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	result = makeNode(PlannedStmt);
 
 	result->commandType = parse->commandType;
+	result->queryId = parse->queryId;
 	result->hasReturning = (parse->returningList != NIL);
 	result->hasModifyingCTE = parse->hasModifyingCTE;
 	result->canSetTag = parse->canSetTag;
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 485d686b058..15d848ff4fe 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -44,6 +44,9 @@
 #include "utils/rel.h"
 
 
+/* Hook for plugins to get control at end of parse analysis */
+post_parse_analyze_hook_type post_parse_analyze_hook = NULL;
+
 static Query *transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt);
 static Query *transformInsertStmt(ParseState *pstate, InsertStmt *stmt);
 static List *transformInsertRow(ParseState *pstate, List *exprlist,
@@ -95,6 +98,9 @@ parse_analyze(Node *parseTree, const char *sourceText,
 
 	query = transformTopLevelStmt(pstate, parseTree);
 
+	if (post_parse_analyze_hook)
+		(*post_parse_analyze_hook) (pstate, query);
+
 	free_parsestate(pstate);
 
 	return query;
@@ -125,6 +131,9 @@ parse_analyze_varparams(Node *parseTree, const char *sourceText,
 	/* make sure all is well with parameter types */
 	check_variable_parameters(pstate, query);
 
+	if (post_parse_analyze_hook)
+		(*post_parse_analyze_hook) (pstate, query);
+
 	free_parsestate(pstate);
 
 	return query;
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 04f9622f788..8f75948d0dd 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -2157,6 +2157,7 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 List *
 QueryRewrite(Query *parsetree)
 {
+	uint32		input_query_id = parsetree->queryId;
 	List	   *querylist;
 	List	   *results;
 	ListCell   *l;
@@ -2181,6 +2182,8 @@ QueryRewrite(Query *parsetree)
 	 * Step 2
 	 *
 	 * Apply all the RIR rules on each query
+	 *
+	 * This is also a handy place to mark each query with the original queryId
 	 */
 	results = NIL;
 	foreach(l, querylist)
@@ -2188,6 +2191,9 @@ QueryRewrite(Query *parsetree)
 		Query	   *query = (Query *) lfirst(l);
 
 		query = fireRIRrules(query, NIL, false);
+
+		query->queryId = input_query_id;
+
 		results = lappend(results, query);
 	}
 
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 14ca7681392..02be36362ce 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -626,6 +626,9 @@ pg_analyze_and_rewrite_params(Node *parsetree,
 
 	query = transformTopLevelStmt(pstate, parsetree);
 
+	if (post_parse_analyze_hook)
+		(*post_parse_analyze_hook) (pstate, query);
+
 	free_parsestate(pstate);
 
 	if (log_parser_stats)
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 07a1ab75502..bc9b6bd774c 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -103,6 +103,8 @@ typedef struct Query
 
 	QuerySource querySource;	/* where did I come from? */
 
+	uint32		queryId;		/* query identifier (can be set by plugins) */
+
 	bool		canSetTag;		/* do I set the command result tag? */
 
 	Node	   *utilityStmt;	/* non-null if this is DECLARE CURSOR or a
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index c7c1a154fc4..fb9a863e15c 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -37,6 +37,8 @@ typedef struct PlannedStmt
 
 	CmdType		commandType;	/* select|insert|update|delete */
 
+	uint32		queryId;		/* query identifier (copied from Query) */
+
 	bool		hasReturning;	/* is it insert|update|delete RETURNING? */
 
 	bool		hasModifyingCTE;	/* has insert|update|delete in WITH? */
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 8367db8b8c5..fe7f80a5aaa 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -16,6 +16,11 @@
 
 #include "parser/parse_node.h"
 
+/* Hook for plugins to get control at end of parse analysis */
+typedef void (*post_parse_analyze_hook_type) (ParseState *pstate,
+											  Query *query);
+extern PGDLLIMPORT post_parse_analyze_hook_type post_parse_analyze_hook;
+
 
 extern Query *parse_analyze(Node *parseTree, const char *sourceText,
 			  Oid *paramTypes, int numParams);