On 1/9/22 5:49 AM, Tom Lane wrote:
The idea I'd been vaguely thinking about is to allow attaching a list
of query-hash nodes to a Query, where each node would contain a "tag"
identifying the specific hash calculation method, and also the value
of the query's hash calculated according to that method. We could
probably get away with saying that all such hash values must be uint64.
The main difference from your function-OID idea, I think, is that
I'm envisioning the tags as being small integers with well-known
values, similarly to the way we manage stakind values in pg_statistic.
In this way, an extension that wants a hash that the core knows how
to calculate doesn't need its own copy of the code, and similarly
one extension could publish a calculation method for use by other
extensions.
To move forward, I have made a patch that implements this idea (see
attachment). It is a POC, but passes all regression tests.
Registration of an queryId generator implemented by analogy with
extensible methods machinery. Also, I switched queryId to int64 type and
renamed to 'label'.
Some lessons learned:
1. Single queryId implementation is deeply tangled with the core code
(stat reporting machinery and parallel workers as an example).
2. We need a custom queryId, that is based on a generated queryId
(according to the logic of pg_stat_statements).
3. We should think about safety of de-registering procedure.
4. We should reserve position of default in-core generator and think on
logic of enabling/disabling it.
5. We should add an EXPLAIN hook, to allow an extension to print this
custom queryId.
--
regards,
Andrey Lepikhov
Postgres Professional
>From f54bb60bd71b49ac8e1b85cd2ad86332a8a81e84 Mon Sep 17 00:00:00 2001
From: Andrey Lepikhov <a.lepik...@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
