From f54bb60bd71b49ac8e1b85cd2ad86332a8a81e84 Mon Sep 17 00:00:00 2001
From: Andrey Lepikhov <a.lepikhov@postgrespro.ru>
Date: Thu, 20 Jan 2022 16:05:17 +0500
Subject: [PATCH] Initial commit

---
 .../pg_stat_statements/pg_stat_statements.c   |  57 ++++--
 src/backend/commands/explain.c                |  11 +-
 src/backend/executor/execMain.c               |   4 +-
 src/backend/executor/execParallel.c           |   3 +-
 src/backend/nodes/copyfuncs.c                 |  21 ++-
 src/backend/nodes/outfuncs.c                  |  17 +-
 src/backend/nodes/readfuncs.c                 |  17 +-
 src/backend/optimizer/plan/planner.c          |   2 +-
 src/backend/parser/analyze.c                  |  15 +-
 src/backend/rewrite/rewriteHandler.c          |   4 +-
 src/backend/tcop/postgres.c                   |  11 +-
 src/backend/utils/misc/guc.c                  |   2 +-
 src/backend/utils/misc/queryjumble.c          | 164 ++++++++++++++++--
 src/include/nodes/nodes.h                     |   1 +
 src/include/nodes/parsenodes.h                |  10 +-
 src/include/nodes/plannodes.h                 |   2 +-
 src/include/parser/analyze.h                  |   3 +-
 src/include/utils/queryjumble.h               |  12 +-
 18 files changed, 294 insertions(+), 62 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 082bfa8f77..ebfc3331df 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -307,8 +307,7 @@ PG_FUNCTION_INFO_V1(pg_stat_statements_info);
 
 static void pgss_shmem_startup(void);
 static void pgss_shmem_shutdown(int code, Datum arg);
-static void pgss_post_parse_analyze(ParseState *pstate, Query *query,
-									JumbleState *jstate);
+static void pgss_post_parse_analyze(ParseState *pstate, Query *query);
 static PlannedStmt *pgss_planner(Query *parse,
 								 const char *query_string,
 								 int cursorOptions,
@@ -813,13 +812,29 @@ error:
 }
 
 /*
- * Post-parse-analysis hook: mark query with a queryId
+ * Post-parse-analysis hook: create a label for the query.
  */
 static void
-pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
+pgss_post_parse_analyze(ParseState *pstate, Query *query)
 {
+	JumbleState	   *jstate;
+	QueryLabel	   *label = get_query_label(query->queryIds, 0);
+
 	if (prev_post_parse_analyze_hook)
-		prev_post_parse_analyze_hook(pstate, query, jstate);
+		prev_post_parse_analyze_hook(pstate, query);
+
+	if (label)
+	{
+		add_custom_query_label(&query->queryIds, -1, label->hash);
+		jstate = (JumbleState *) label->context;
+	}
+	else
+	{
+		add_custom_query_label(&query->queryIds, -1, UINT64CONST(0));
+		jstate = NULL;
+	}
+
+	label = get_query_label(query->queryIds, -1);
 
 	/* Safety check... */
 	if (!pgss || !pgss_hash || !pgss_enabled(exec_nested_level))
@@ -833,7 +848,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 	if (query->utilityStmt)
 	{
 		if (pgss_track_utility && !PGSS_HANDLED_UTILITY(query->utilityStmt))
-			query->queryId = UINT64CONST(0);
+			label->hash = UINT64CONST(0);
 		return;
 	}
 
@@ -846,7 +861,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 	 */
 	if (jstate && jstate->clocations_count > 0)
 		pgss_store(pstate->p_sourcetext,
-				   query->queryId,
+				   label->hash,
 				   query->stmt_location,
 				   query->stmt_len,
 				   PGSS_INVALID,
@@ -868,6 +883,9 @@ pgss_planner(Query *parse,
 			 ParamListInfo boundParams)
 {
 	PlannedStmt *result;
+	int64 queryId;
+
+	queryId = get_query_label_hash(parse->queryIds, -1);
 
 	/*
 	 * We can't process the query if no query_string is provided, as
@@ -883,7 +901,7 @@ pgss_planner(Query *parse,
 	 */
 	if (pgss_enabled(plan_nested_level + exec_nested_level)
 		&& pgss_track_planning && query_string
-		&& parse->queryId != UINT64CONST(0))
+		&& queryId != UINT64CONST(0))
 	{
 		instr_time	start;
 		instr_time	duration;
@@ -930,7 +948,7 @@ pgss_planner(Query *parse,
 		WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start);
 
 		pgss_store(query_string,
-				   parse->queryId,
+				   queryId,
 				   parse->stmt_location,
 				   parse->stmt_len,
 				   PGSS_PLAN,
@@ -959,17 +977,21 @@ pgss_planner(Query *parse,
 static void
 pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 {
+	int64 queryId;
+
 	if (prev_ExecutorStart)
 		prev_ExecutorStart(queryDesc, eflags);
 	else
 		standard_ExecutorStart(queryDesc, eflags);
 
+	queryId = get_query_label_hash(queryDesc->plannedstmt->queryIds, -1);
+
 	/*
 	 * If query has queryId zero, don't track it.  This prevents double
 	 * counting of optimizable statements that are directly contained in
 	 * utility statements.
 	 */
-	if (pgss_enabled(exec_nested_level) && queryDesc->plannedstmt->queryId != UINT64CONST(0))
+	if (pgss_enabled(exec_nested_level) && queryId != UINT64CONST(0))
 	{
 		/*
 		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
@@ -1036,7 +1058,7 @@ pgss_ExecutorFinish(QueryDesc *queryDesc)
 static void
 pgss_ExecutorEnd(QueryDesc *queryDesc)
 {
-	uint64		queryId = queryDesc->plannedstmt->queryId;
+	uint64		queryId = get_query_label_hash(queryDesc->plannedstmt->queryIds, -1);
 
 	if (queryId != UINT64CONST(0) && queryDesc->totaltime &&
 		pgss_enabled(exec_nested_level))
@@ -1076,7 +1098,16 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 					DestReceiver *dest, QueryCompletion *qc)
 {
 	Node	   *parsetree = pstmt->utilityStmt;
-	uint64		saved_queryId = pstmt->queryId;
+	QueryLabel *label = get_query_label(pstmt->queryIds, -1);
+	uint64		saved_queryId;
+
+	if (!label)
+	{
+		add_custom_query_label(&pstmt->queryIds, -1, UINT64CONST(0));
+		label = get_query_label(pstmt->queryIds, -1);
+	}
+
+	saved_queryId = label->hash;
 
 	/*
 	 * Force utility statements to get queryId zero.  We do this even in cases
@@ -1093,7 +1124,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	 * only.
 	 */
 	if (pgss_enabled(exec_nested_level) && pgss_track_utility)
-		pstmt->queryId = UINT64CONST(0);
+		label->hash = UINT64CONST(0);
 
 	/*
 	 * If it's an EXECUTE statement, we don't track it and don't increment the
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index b970997c34..b31a870156 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -166,7 +166,6 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 {
 	ExplainState *es = NewExplainState();
 	TupOutputState *tstate;
-	JumbleState *jstate = NULL;
 	Query	   *query;
 	List	   *rewritten;
 	ListCell   *lc;
@@ -246,10 +245,10 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 
 	query = castNode(Query, stmt->query);
 	if (IsQueryIdEnabled())
-		jstate = JumbleQuery(query, pstate->p_sourcetext);
+		GenerateQueryLabels(query, pstate->p_sourcetext);
 
 	if (post_parse_analyze_hook)
-		(*post_parse_analyze_hook) (pstate, query, jstate);
+		(*post_parse_analyze_hook) (pstate, query);
 
 	/*
 	 * Parse analysis was done already, but we still have to run the rule
@@ -604,14 +603,14 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	/* Create textual dump of plan tree */
 	ExplainPrintPlan(es, queryDesc);
 
-	if (es->verbose && plannedstmt->queryId != UINT64CONST(0))
+	if (es->verbose && plannedstmt->queryIds != NIL)
 	{
 		/*
 		 * Output the queryid as an int64 rather than a uint64 so we match
 		 * what would be seen in the BIGINT pg_stat_statements.queryid column.
 		 */
-		ExplainPropertyInteger("Query Identifier", NULL, (int64)
-							   plannedstmt->queryId, es);
+		ExplainPropertyInteger("Query Identifier", NULL,
+							   get_query_label_hash(plannedstmt->queryIds, 0), es);
 	}
 
 	/* Show buffer usage in planning */
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 549d9eb696..cbb97cf91f 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -62,6 +62,7 @@
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/partcache.h"
+#include "utils/queryjumble.h"
 #include "utils/rls.h"
 #include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
@@ -135,7 +136,8 @@ ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 * that it's harmless to report the query_id multiple time, as the call
 	 * will be ignored if the top level query_id has already been reported.
 	 */
-	pgstat_report_query_id(queryDesc->plannedstmt->queryId, false);
+	pgstat_report_query_id(
+				get_query_label_hash(queryDesc->plannedstmt->queryIds, 0), false);
 
 	if (ExecutorStart_hook)
 		(*ExecutorStart_hook) (queryDesc, eflags);
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 5dd8ab7db2..d025860f45 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -49,6 +49,7 @@
 #include "utils/dsa.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
+#include "utils/queryjumble.h"
 #include "utils/snapmgr.h"
 
 /*
@@ -175,7 +176,7 @@ ExecSerializePlan(Plan *plan, EState *estate)
 	 */
 	pstmt = makeNode(PlannedStmt);
 	pstmt->commandType = CMD_SELECT;
-	pstmt->queryId = pgstat_get_my_query_id();
+	add_custom_query_label(&pstmt->queryIds, 0, pgstat_get_my_query_id());
 	pstmt->hasReturning = false;
 	pstmt->hasModifyingCTE = false;
 	pstmt->canSetTag = true;
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 90b5da51c9..6f159e0cbd 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -87,7 +87,7 @@ _copyPlannedStmt(const PlannedStmt *from)
 	PlannedStmt *newnode = makeNode(PlannedStmt);
 
 	COPY_SCALAR_FIELD(commandType);
-	COPY_SCALAR_FIELD(queryId);
+	COPY_NODE_FIELD(queryIds);
 	COPY_SCALAR_FIELD(hasReturning);
 	COPY_SCALAR_FIELD(hasModifyingCTE);
 	COPY_SCALAR_FIELD(canSetTag);
@@ -3159,6 +3159,20 @@ _copyTriggerTransition(const TriggerTransition *from)
 	return newnode;
 }
 
+static QueryLabel *
+_copyQueryLabel(const QueryLabel *from)
+{
+	QueryLabel *newnode = makeNode(QueryLabel);
+
+	COPY_SCALAR_FIELD(kind);
+	COPY_SCALAR_FIELD(hash);
+
+	/* Caller should re-generate labels if it want to use it. */
+	newnode->context = NULL;
+
+	return newnode;
+}
+
 static Query *
 _copyQuery(const Query *from)
 {
@@ -3166,7 +3180,7 @@ _copyQuery(const Query *from)
 
 	COPY_SCALAR_FIELD(commandType);
 	COPY_SCALAR_FIELD(querySource);
-	COPY_SCALAR_FIELD(queryId);
+	COPY_NODE_FIELD(queryIds);
 	COPY_SCALAR_FIELD(canSetTag);
 	COPY_NODE_FIELD(utilityStmt);
 	COPY_SCALAR_FIELD(resultRelation);
@@ -5405,6 +5419,9 @@ copyObjectImpl(const void *from)
 			/*
 			 * PARSE NODES
 			 */
+		case T_QueryLabel:
+			retval = _copyQueryLabel(from);
+			break;
 		case T_Query:
 			retval = _copyQuery(from);
 			break;
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index 2b0236937a..c86eb636b6 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -305,7 +305,7 @@ _outPlannedStmt(StringInfo str, const PlannedStmt *node)
 	WRITE_NODE_TYPE("PLANNEDSTMT");
 
 	WRITE_ENUM_FIELD(commandType, CmdType);
-	WRITE_UINT64_FIELD(queryId);
+	WRITE_NODE_FIELD(queryIds);
 	WRITE_BOOL_FIELD(hasReturning);
 	WRITE_BOOL_FIELD(hasModifyingCTE);
 	WRITE_BOOL_FIELD(canSetTag);
@@ -3030,6 +3030,16 @@ _outStatsElem(StringInfo str, const StatsElem *node)
 	WRITE_NODE_FIELD(expr);
 }
 
+/* Needed for dump of a PlannedStmt node */
+static void
+_outQueryLabel(StringInfo str, const QueryLabel *node)
+{
+	WRITE_NODE_TYPE("QUERYLABEL");
+
+	WRITE_INT_FIELD(kind);
+	WRITE_UINT64_FIELD(hash);
+}
+
 static void
 _outQuery(StringInfo str, const Query *node)
 {
@@ -3037,7 +3047,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 */
+	/* we intentionally do not print the queryId fields */
 	WRITE_BOOL_FIELD(canSetTag);
 
 	/*
@@ -4403,6 +4413,9 @@ outNode(StringInfo str, const void *obj)
 			case T_StatsElem:
 				_outStatsElem(str, obj);
 				break;
+			case T_QueryLabel:
+				_outQueryLabel(str, obj);
+				break;
 			case T_Query:
 				_outQuery(str, obj);
 				break;
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 3f68f7c18d..47c35707d4 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -238,6 +238,17 @@ readBitmapset(void)
 	return _readBitmapset();
 }
 
+static QueryLabel *
+_readQueryLabel(void)
+{
+	READ_LOCALS(QueryLabel);
+	READ_INT_FIELD(kind);
+	READ_UINT64_FIELD(hash);
+	local_node->context = NULL;
+
+	READ_DONE();
+}
+
 /*
  * _readQuery
  */
@@ -248,7 +259,7 @@ _readQuery(void)
 
 	READ_ENUM_FIELD(commandType, CmdType);
 	READ_ENUM_FIELD(querySource, QuerySource);
-	local_node->queryId = UINT64CONST(0);	/* not saved in output format */
+	local_node->queryIds = NIL;	/* not saved in output format */
 	READ_BOOL_FIELD(canSetTag);
 	READ_NODE_FIELD(utilityStmt);
 	READ_INT_FIELD(resultRelation);
@@ -1578,7 +1589,7 @@ _readPlannedStmt(void)
 	READ_LOCALS(PlannedStmt);
 
 	READ_ENUM_FIELD(commandType, CmdType);
-	READ_UINT64_FIELD(queryId);
+	READ_NODE_FIELD(queryIds);
 	READ_BOOL_FIELD(hasReturning);
 	READ_BOOL_FIELD(hasModifyingCTE);
 	READ_BOOL_FIELD(canSetTag);
@@ -2728,6 +2739,8 @@ parseNodeString(void)
 
 	if (MATCH("QUERY", 5))
 		return_value = _readQuery();
+	else if (MATCH("QUERYLABEL", 10))
+		return_value = _readQueryLabel();
 	else if (MATCH("WITHCHECKOPTION", 15))
 		return_value = _readWithCheckOption();
 	else if (MATCH("SORTGROUPCLAUSE", 15))
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index bd09f85aea..2ef8a97997 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -510,7 +510,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	result = makeNode(PlannedStmt);
 
 	result->commandType = parse->commandType;
-	result->queryId = parse->queryId;
+	result->queryIds = list_copy(parse->queryIds);
 	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 6ac2e9ce23..5925641795 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -113,7 +113,6 @@ parse_analyze(RawStmt *parseTree, const char *sourceText,
 {
 	ParseState *pstate = make_parsestate(NULL);
 	Query	   *query;
-	JumbleState *jstate = NULL;
 
 	Assert(sourceText != NULL); /* required as of 8.4 */
 
@@ -127,14 +126,15 @@ parse_analyze(RawStmt *parseTree, const char *sourceText,
 	query = transformTopLevelStmt(pstate, parseTree);
 
 	if (IsQueryIdEnabled())
-		jstate = JumbleQuery(query, sourceText);
+		GenerateQueryLabels(query, sourceText);
 
 	if (post_parse_analyze_hook)
-		(*post_parse_analyze_hook) (pstate, query, jstate);
+		(*post_parse_analyze_hook) (pstate, query);
 
 	free_parsestate(pstate);
 
-	pgstat_report_query_id(query->queryId, false);
+	/* Report id of default generator. */
+	pgstat_report_query_id(get_query_label_hash(query->queryIds, 0), false);
 
 	return query;
 }
@@ -152,7 +152,6 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 {
 	ParseState *pstate = make_parsestate(NULL);
 	Query	   *query;
-	JumbleState *jstate = NULL;
 
 	Assert(sourceText != NULL); /* required as of 8.4 */
 
@@ -166,14 +165,14 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 	check_variable_parameters(pstate, query);
 
 	if (IsQueryIdEnabled())
-		jstate = JumbleQuery(query, sourceText);
+		GenerateQueryLabels(query, sourceText);
 
 	if (post_parse_analyze_hook)
-		(*post_parse_analyze_hook) (pstate, query, jstate);
+		(*post_parse_analyze_hook) (pstate, query);
 
 	free_parsestate(pstate);
 
-	pgstat_report_query_id(query->queryId, false);
+	pgstat_report_query_id(get_query_label_hash(query->queryIds, 0), false);
 
 	return query;
 }
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 3d82138cb3..1cce0949b9 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -4079,7 +4079,7 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 List *
 QueryRewrite(Query *parsetree)
 {
-	uint64		input_query_id = parsetree->queryId;
+	List	   *input_query_ids = parsetree->queryIds;
 	List	   *querylist;
 	List	   *results;
 	ListCell   *l;
@@ -4114,7 +4114,7 @@ QueryRewrite(Query *parsetree)
 
 		query = fireRIRrules(query, NIL);
 
-		query->queryId = input_query_id;
+		query->queryIds = input_query_ids;
 
 		results = lappend(results, query);
 	}
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index fda2e9360e..76a8312b47 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -686,7 +686,6 @@ pg_analyze_and_rewrite_params(RawStmt *parsetree,
 	ParseState *pstate;
 	Query	   *query;
 	List	   *querytree_list;
-	JumbleState *jstate = NULL;
 
 	Assert(query_string != NULL);	/* required as of 8.4 */
 
@@ -706,14 +705,14 @@ pg_analyze_and_rewrite_params(RawStmt *parsetree,
 	query = transformTopLevelStmt(pstate, parsetree);
 
 	if (IsQueryIdEnabled())
-		jstate = JumbleQuery(query, query_string);
+		GenerateQueryLabels(query, query_string);
 
 	if (post_parse_analyze_hook)
-		(*post_parse_analyze_hook) (pstate, query, jstate);
+		(*post_parse_analyze_hook) (pstate, query);
 
 	free_parsestate(pstate);
 
-	pgstat_report_query_id(query->queryId, false);
+	pgstat_report_query_id(get_query_label_hash(query->queryIds, 0), false);
 
 	if (log_parser_stats)
 		ShowUsage("PARSE ANALYSIS STATISTICS");
@@ -797,7 +796,7 @@ pg_rewrite_query(Query *query)
 				 * queryId is not saved in stored rules, but we must preserve
 				 * it here to avoid breaking pg_stat_statements.
 				 */
-				new_query->queryId = query->queryId;
+				new_query->queryIds = list_copy(query->queryIds);
 
 				new_list = lappend(new_list, new_query);
 				pfree(str);
@@ -933,7 +932,7 @@ pg_plan_queries(List *querytrees, const char *query_string, int cursorOptions,
 			stmt->utilityStmt = query->utilityStmt;
 			stmt->stmt_location = query->stmt_location;
 			stmt->stmt_len = query->stmt_len;
-			stmt->queryId = query->queryId;
+			stmt->queryIds = list_copy(query->queryIds);
 		}
 		else
 		{
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index 4c94f09c64..fb5dd1cdce 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -4695,7 +4695,7 @@ static struct config_enum ConfigureNamesEnum[] =
 		},
 		&compute_query_id,
 		COMPUTE_QUERY_ID_AUTO, compute_query_id_options,
-		NULL, NULL, NULL
+		NULL, assign_query_id, NULL
 	},
 
 	{
diff --git a/src/backend/utils/misc/queryjumble.c b/src/backend/utils/misc/queryjumble.c
index a67487e5fe..59bf53bc46 100644
--- a/src/backend/utils/misc/queryjumble.c
+++ b/src/backend/utils/misc/queryjumble.c
@@ -97,18 +97,18 @@ CleanQuerytext(const char *query, int *location, int *len)
 	return query;
 }
 
-JumbleState *
-JumbleQuery(Query *query, const char *querytext)
+int64
+JumbleQuery(Query *query, const char *querytext, void **context)
 {
-	JumbleState *jstate = NULL;
+	JumbleState	   *jstate = NULL;
+	int64			queryId;
 
 	Assert(IsQueryIdEnabled());
 
 	if (query->utilityStmt)
 	{
-		query->queryId = compute_utility_query_id(querytext,
-												  query->stmt_location,
-												  query->stmt_len);
+		queryId = compute_utility_query_id(querytext, query->stmt_location,
+										   query->stmt_len);
 	}
 	else
 	{
@@ -125,19 +125,20 @@ JumbleQuery(Query *query, const char *querytext)
 
 		/* Compute query ID and mark the Query node with it */
 		JumbleQueryInternal(jstate, query);
-		query->queryId = DatumGetUInt64(hash_any_extended(jstate->jumble,
-														  jstate->jumble_len,
-														  0));
+		queryId = DatumGetUInt64(hash_any_extended(jstate->jumble,
+												   jstate->jumble_len,
+												   0));
 
 		/*
 		 * If we are unlucky enough to get a hash of zero, use 1 instead, to
 		 * prevent confusion with the utility-statement case.
 		 */
-		if (query->queryId == UINT64CONST(0))
-			query->queryId = UINT64CONST(1);
+		if (queryId == UINT64CONST(0))
+			queryId = UINT64CONST(1);
 	}
 
-	return jstate;
+	*context = jstate;
+	return queryId;
 }
 
 /*
@@ -150,7 +151,29 @@ void
 EnableQueryId(void)
 {
 	if (compute_query_id != COMPUTE_QUERY_ID_OFF)
+	{
+		(void) RegisterQueryIdGen(0, JumbleQuery);
 		query_id_enabled = true;
+	}
+}
+
+void
+assign_query_id(int newval, void *extra)
+{
+	switch (newval)
+	{
+		case COMPUTE_QUERY_ID_OFF:
+			/* De-register routine should be implemented */
+			break;
+		case COMPUTE_QUERY_ID_AUTO:
+			break;
+		case COMPUTE_QUERY_ID_ON:
+			(void) RegisterQueryIdGen(0, JumbleQuery);
+			break;
+		default:
+			elog(ERROR, "Unknown value");
+			break;
+	}
 }
 
 /*
@@ -856,3 +879,120 @@ RecordConstLocation(JumbleState *jstate, int location)
 		jstate->clocations_count++;
 	}
 }
+
+/* *******************************************************************************
+ *
+ * Query label custom generation machinery.
+ *
+ * ******************************************************************************/
+
+#include "utils/hsearch.h"
+
+static HTAB *QueryLabelGens = NULL;
+
+typedef struct QLGeneratorEntry
+{
+	int16 kind;
+	query_generator_type callback;
+} QLGeneratorEntry;
+
+/*
+ * An internal function to register a new callback structure
+ * Use a positive value for the 'kind' field to register an query label generator.
+ * Caller should worry about intersections by this value with other extensions.
+ */
+bool
+RegisterQueryIdGen(const int16 kind, query_generator_type callback)
+{
+	QLGeneratorEntry   *entry;
+	bool				found;
+
+	if (QueryLabelGens == NULL)
+	{
+		HASHCTL		ctl;
+
+		ctl.keysize = sizeof(int16);
+		ctl.entrysize = sizeof(QLGeneratorEntry);
+
+		QueryLabelGens = hash_create("Query Label Generators", 8, &ctl,
+									 HASH_ELEM | HASH_BLOBS);
+	}
+
+	if (kind < 0)
+		elog(ERROR, "Can't use value %d as an query generator kind.", kind);
+
+	entry = (QLGeneratorEntry *) hash_search(QueryLabelGens, (void *) &kind,
+											 HASH_ENTER, &found);
+	if (found)
+		/* Give caller a chance to process the problem. */
+		return false;
+
+	entry->callback = callback;
+	return true;
+}
+
+/*
+ * Create QueryLabel entry for each registered generator.
+ * All memory allocations is made in current memory context.
+ */
+void
+GenerateQueryLabels(Query *query, const char *querytext)
+{
+	HASH_SEQ_STATUS		hash_seq;
+	QLGeneratorEntry   *entry;
+
+	/* Newly generated query haven't any labels. */
+	Assert(query->queryIds == NIL);
+
+	if (QueryLabelGens == NULL || hash_get_num_entries(QueryLabelGens) == 0)
+		return;
+
+	hash_seq_init(&hash_seq, QueryLabelGens);
+	while ((entry = hash_seq_search(&hash_seq)) != NULL)
+	{
+		QueryLabel *qlabel = makeNode(QueryLabel);
+
+		qlabel->kind = entry->kind;
+		qlabel->hash = entry->callback(query, querytext, &qlabel->context);
+		query->queryIds = lappend(query->queryIds, qlabel);
+	}
+}
+
+int64
+get_query_label_hash(List *queryIds, const int16 kind)
+{
+	QueryLabel *label;
+
+	label = get_query_label(queryIds, kind);
+	return (label) ? label->hash : 0;
+}
+
+QueryLabel *
+get_query_label(List *queryIds, const int16 kind)
+{
+	ListCell *lc;
+
+	foreach(lc, queryIds)
+	{
+		QueryLabel *label = lfirst_node(QueryLabel, lc);
+		if (label->kind == kind)
+			return label;
+	}
+	return NULL;
+}
+
+bool
+add_custom_query_label(List **queryIds, int16 kind, int64 hash)
+{
+	QueryLabel *label;
+
+	if (get_query_label_hash(*queryIds, kind) != 0)
+		elog(ERROR, "Duplicated custom label %ld for the kind %d.", hash, kind);
+
+	label = makeNode(QueryLabel);
+	label->kind = kind;
+	label->hash = hash;
+	label->context = NULL;
+	*queryIds = lappend(*queryIds, label);
+	return true;
+}
diff --git a/src/include/nodes/nodes.h b/src/include/nodes/nodes.h
index f9ddafd345..0835034c2d 100644
--- a/src/include/nodes/nodes.h
+++ b/src/include/nodes/nodes.h
@@ -314,6 +314,7 @@ typedef enum NodeTag
 	T_RawStmt,
 	T_Query,
 	T_PlannedStmt,
+	T_QueryLabel,
 	T_InsertStmt,
 	T_DeleteStmt,
 	T_UpdateStmt,
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 3e9bdc781f..bc9ca32273 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -70,6 +70,14 @@ typedef enum SetQuantifier
 	SET_QUANTIFIER_DISTINCT
 } SetQuantifier;
 
+typedef struct QueryLabel
+{
+	NodeTag type;
+	int16	kind; /* unique ID of generator. 0 reserved for in-core jumbling routine. */
+	int64	hash; /* generated stamp. */
+	void   *context; /* internal generator-related data. */
+} QueryLabel;
+
 /*
  * Grantable rights are encoded so that we can OR them together in a bitmask.
  * The present representation of AclItem limits us to 16 distinct rights,
@@ -121,7 +129,7 @@ typedef struct Query
 
 	QuerySource querySource;	/* where did I come from? */
 
-	uint64		queryId;		/* query identifier (can be set by plugins) */
+	List	   *queryIds;		/* query identifiers (can be added by plugins) */
 
 	bool		canSetTag;		/* do I set the command result tag? */
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 0b518ce6b2..f588311de2 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -45,7 +45,7 @@ typedef struct PlannedStmt
 
 	CmdType		commandType;	/* select|insert|update|delete|utility */
 
-	uint64		queryId;		/* query identifier (copied from Query) */
+	List	   *queryIds;		/* query identifiers (copied from Query) */
 
 	bool		hasReturning;	/* is it insert|update|delete RETURNING? */
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 0022184de0..455396c054 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -19,8 +19,7 @@
 
 /* Hook for plugins to get control at end of parse analysis */
 typedef void (*post_parse_analyze_hook_type) (ParseState *pstate,
-											  Query *query,
-											  JumbleState *jstate);
+											  Query *query);
 extern PGDLLIMPORT post_parse_analyze_hook_type post_parse_analyze_hook;
 
 
diff --git a/src/include/utils/queryjumble.h b/src/include/utils/queryjumble.h
index a4c277269e..f09b88e3ec 100644
--- a/src/include/utils/queryjumble.h
+++ b/src/include/utils/queryjumble.h
@@ -65,7 +65,7 @@ extern int	compute_query_id;
 
 
 extern const char *CleanQuerytext(const char *query, int *location, int *len);
-extern JumbleState *JumbleQuery(Query *query, const char *querytext);
+extern int64 JumbleQuery(Query *query, const char *querytext, void **context);
 extern void EnableQueryId(void);
 
 extern bool query_id_enabled;
@@ -84,4 +84,14 @@ IsQueryIdEnabled(void)
 	return query_id_enabled;
 }
 
+typedef int64 (*query_generator_type) (Query *query, const char *querytext,
+									   void **context);
+
+extern void assign_query_id(int newval, void *extra);
+extern bool RegisterQueryIdGen(const int16 kind, query_generator_type callback);
+extern void GenerateQueryLabels(Query *query, const char *querytext);
+extern int64 get_query_label_hash(List *queryIds, const int16 kind);
+extern QueryLabel *get_query_label(List *queryIds, const int16 kind);
+extern bool add_custom_query_label(List **queryIds, int16 kind, int64 hash);
+
 #endif							/* QUERYJUMBLE_H */
-- 
2.25.1

