From f54bb60bd71b49ac8e1b85cd2ad86332a8a81e84 Mon Sep 17 00:00:00 2001 From: Andrey Lepikhov 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