From 944ce61d7ff934727240d90ee7620bfb69ad3a5a Mon Sep 17 00:00:00 2001 From: Andrey Lepikhov Date: Wed, 22 Mar 2023 16:59:30 +0500 Subject: [PATCH] Add on more field into Query and PlannedStmt structures to allow an extension to pass some data across all the execution stages of specific instance of the query. --- .../pg_stat_statements/pg_stat_statements.c | 57 ++++++++-- src/backend/commands/extension.c | 103 ++++++++++++++++++ src/backend/executor/execParallel.c | 1 + src/backend/nodes/value.c | 12 ++ src/backend/optimizer/plan/planner.c | 1 + src/backend/rewrite/rewriteHandler.c | 6 + src/backend/tcop/postgres.c | 2 + src/include/commands/extension.h | 4 + src/include/nodes/parsenodes.h | 23 ++++ src/include/nodes/plannodes.h | 2 + src/include/nodes/value.h | 10 ++ 11 files changed, 209 insertions(+), 12 deletions(-) diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c index 5285c3f7fa..fc47661c86 100644 --- a/contrib/pg_stat_statements/pg_stat_statements.c +++ b/contrib/pg_stat_statements/pg_stat_statements.c @@ -49,6 +49,7 @@ #include "access/parallel.h" #include "catalog/pg_authid.h" +#include "commands/extension.h" #include "common/hashfn.h" #include "executor/instrument.h" #include "funcapi.h" @@ -107,6 +108,14 @@ static const uint32 PGSS_PG_MAJOR_VERSION = PG_VERSION_NUM / 100; !IsA(n, PrepareStmt) && \ !IsA(n, DeallocateStmt)) +#define GET_QUERYID(node) \ + (Bigint *) GetExtensionData(node->ext_field, "pg_stat_statements") + +#define INSERT_QUERYID(node, queryid) \ + castNode(Bigint, AddExtensionDataToNode((Node *) node, \ + "pg_stat_statements", \ + (Node *) makeBigint((int64) queryid), \ + true)) /* * Extension version number, for supporting older extension versions' objects */ @@ -821,6 +830,13 @@ error: static void pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate) { + Bigint *queryId; + + if ((queryId = GET_QUERYID(query)) == NULL) + queryId = INSERT_QUERYID(query, query->queryId); + + Assert(queryId); + if (prev_post_parse_analyze_hook) prev_post_parse_analyze_hook(pstate, query, jstate); @@ -837,7 +853,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate) { if (pgss_track_utility && !PGSS_HANDLED_UTILITY(query->utilityStmt)) { - query->queryId = UINT64CONST(0); + queryId->val = UINT64CONST(0); return; } } @@ -851,7 +867,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate) */ if (jstate && jstate->clocations_count > 0) pgss_store(pstate->p_sourcetext, - query->queryId, + queryId->val, query->stmt_location, query->stmt_len, PGSS_INVALID, @@ -873,7 +889,10 @@ pgss_planner(Query *parse, int cursorOptions, ParamListInfo boundParams) { - PlannedStmt *result; + PlannedStmt *result; + Bigint *queryId; + + queryId = GET_QUERYID(parse); /* * We can't process the query if no query_string is provided, as @@ -889,7 +908,7 @@ pgss_planner(Query *parse, */ if (pgss_enabled(plan_nested_level + exec_nested_level) && pgss_track_planning && query_string - && parse->queryId != UINT64CONST(0)) + && queryId && queryId->val != UINT64CONST(0)) { instr_time start; instr_time duration; @@ -936,7 +955,7 @@ pgss_planner(Query *parse, WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start); pgss_store(query_string, - parse->queryId, + queryId->val, parse->stmt_location, parse->stmt_len, PGSS_PLAN, @@ -966,17 +985,22 @@ pgss_planner(Query *parse, static void pgss_ExecutorStart(QueryDesc *queryDesc, int eflags) { + Bigint *queryId; + if (prev_ExecutorStart) prev_ExecutorStart(queryDesc, eflags); else standard_ExecutorStart(queryDesc, eflags); + queryId = GET_QUERYID(queryDesc->plannedstmt); + /* * 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 && queryId->val != UINT64CONST(0)) { /* * Set up to track total elapsed time in ExecutorRun. Make sure the @@ -1043,9 +1067,11 @@ pgss_ExecutorFinish(QueryDesc *queryDesc) static void pgss_ExecutorEnd(QueryDesc *queryDesc) { - uint64 queryId = queryDesc->plannedstmt->queryId; + Bigint *queryId; + + queryId = GET_QUERYID(queryDesc->plannedstmt); - if (queryId != UINT64CONST(0) && queryDesc->totaltime && + if (queryId && queryId->val != UINT64CONST(0) && queryDesc->totaltime && pgss_enabled(exec_nested_level)) { /* @@ -1055,7 +1081,7 @@ pgss_ExecutorEnd(QueryDesc *queryDesc) InstrEndLoop(queryDesc->totaltime); pgss_store(queryDesc->sourceText, - queryId, + queryId->val, queryDesc->plannedstmt->stmt_location, queryDesc->plannedstmt->stmt_len, PGSS_EXEC, @@ -1084,10 +1110,17 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString, DestReceiver *dest, QueryCompletion *qc) { Node *parsetree = pstmt->utilityStmt; - uint64 saved_queryId = pstmt->queryId; + uint64 saved_queryId; + Bigint *queryId; int saved_stmt_location = pstmt->stmt_location; int saved_stmt_len = pstmt->stmt_len; + if ((queryId = GET_QUERYID(pstmt)) == NULL) + queryId = INSERT_QUERYID(pstmt, pstmt->queryId); + + Assert(queryId); + saved_queryId = queryId->val; + /* * Force utility statements to get queryId zero. We do this even in cases * where the statement contains an optimizable statement for which a @@ -1103,7 +1136,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString, * only. */ if (pgss_enabled(exec_nested_level) && pgss_track_utility) - pstmt->queryId = UINT64CONST(0); + queryId->val = UINT64CONST(0); /* * If it's an EXECUTE statement, we don't track it and don't increment the @@ -1158,7 +1191,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString, * freed. We must copy everything we still need into local variables, * which we did above. * - * For the same reason, we can't risk restoring pstmt->queryId to its + * For the same reason, we can't risk restoring queryId to its * former value, which'd otherwise be a good idea. */ diff --git a/src/backend/commands/extension.c b/src/backend/commands/extension.c index 02ff4a9a7f..74f8a820cb 100644 --- a/src/backend/commands/extension.c +++ b/src/backend/commands/extension.c @@ -3395,3 +3395,106 @@ read_whole_file(const char *filename, int *length) buf[*length] = '\0'; return buf; } + +static ExtensionData * +_initial_actions(Node *node, const char *entryname) +{ + ExtensionData **ext_field; + ExtensionData *scanner; + ExtensionData *elem = NULL; + + switch (nodeTag(node)) + { + case T_Query: + ext_field = &((Query *) node)->ext_field; + break; + case T_PlannedStmt: + ext_field = &((PlannedStmt *) node)->ext_field; + break; + default: + /* TODO: use ereport */ + elog(ERROR, "Unexpected node type: %d", nodeTag(node)); + } + + Assert(strlen(entryname) > 0); +#ifdef USE_ASSERT_CHECKING + { + const char *c; + /* check name doesn't contain escapable symbols */ + for (c = entryname; *c; c++) + Assert(isgraph(*c) && strchr("(){}\\\"", *c) == NULL); + } +#endif + + /* Check existing entry */ + for (scanner = *ext_field; scanner != NULL; scanner = scanner->next) + { + if (strcmp(scanner->name, entryname) == 0) + { + elem = scanner; + break; + } + } + + if (elem == NULL) + { + elem = makeNode(ExtensionData); + elem->next = *ext_field; + elem->name = pstrdup(entryname); + *ext_field = elem; + } + + return elem; +} + +Node * +AddExtensionDataToNode(Node *node, const char *entryname, const Node *data, + bool noCopy) +{ + MemoryContext destCtx = GetMemoryChunkContext(node); + MemoryContext oldctx; + ExtensionData *elem; + + /* Safely store in the Query memory context */ + oldctx = MemoryContextSwitchTo(destCtx); + + elem = _initial_actions(node, entryname); + elem->node = (Node *) (noCopy ? data : copyObject(data)); + + MemoryContextSwitchTo(oldctx); + +#ifdef USE_ASSERT_CHECKING + if (noCopy) + { + /* brute check val is copyable */ + MemoryContext tmpctx; + + tmpctx = AllocSetContextCreate(CurrentMemoryContext, + "assert_checking", + ALLOCSET_DEFAULT_SIZES); + oldctx = MemoryContextSwitchTo(tmpctx); + (void) copyObject(data); + MemoryContextSwitchTo(oldctx); + MemoryContextDelete(tmpctx); + } +#endif + + return elem->node; +} + +/* + * Return the node, stored in extended area of Query or PlannedStmt or + * NULL, if not found. + */ +Node * +GetExtensionData(ExtensionData *extdata, const char *entryname) +{ + while (extdata != NULL) + { + if (strcmp(entryname, extdata->name) == 0) + return extdata->node; + + extdata = extdata->next; + } + return NULL; +} diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c index aa3f283453..ae0d67a419 100644 --- a/src/backend/executor/execParallel.c +++ b/src/backend/executor/execParallel.c @@ -176,6 +176,7 @@ ExecSerializePlan(Plan *plan, EState *estate) pstmt = makeNode(PlannedStmt); pstmt->commandType = CMD_SELECT; pstmt->queryId = pgstat_get_my_query_id(); + /* What about Extension nodes here? */ pstmt->hasReturning = false; pstmt->hasModifyingCTE = false; pstmt->canSetTag = true; diff --git a/src/backend/nodes/value.c b/src/backend/nodes/value.c index f11074970e..eb1a975e7d 100644 --- a/src/backend/nodes/value.c +++ b/src/backend/nodes/value.c @@ -28,6 +28,18 @@ makeInteger(int i) return v; } +/* + * makeInteger + */ +Bigint * +makeBigint(int64 i) +{ + Bigint *v = makeNode(Bigint); + + v->val = i; + return v; +} + /* * makeFloat * diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c index a1873ce26d..fb4ac7fd0a 100644 --- a/src/backend/optimizer/plan/planner.c +++ b/src/backend/optimizer/plan/planner.c @@ -517,6 +517,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions, result->commandType = parse->commandType; result->queryId = parse->queryId; + result->ext_field = copyObject(parse->ext_field); result->hasReturning = (parse->returningList != NIL); result->hasModifyingCTE = parse->hasModifyingCTE; result->canSetTag = parse->canSetTag; diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c index 980dc1816f..3403abfe42 100644 --- a/src/backend/rewrite/rewriteHandler.c +++ b/src/backend/rewrite/rewriteHandler.c @@ -4232,6 +4232,12 @@ QueryRewrite(Query *parsetree) query->queryId = input_query_id; + /* + * Should we copy extension fields here? Guess no, because it could be + * quite different query. Much simpler to skip these fields and allow + * extensions to realize that and re-generate data if needed. + */ + results = lappend(results, query); } diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c index cab709b07b..c6675b340c 100644 --- a/src/backend/tcop/postgres.c +++ b/src/backend/tcop/postgres.c @@ -823,6 +823,7 @@ pg_rewrite_query(Query *query) * here to avoid breaking pg_stat_statements. */ new_query->queryId = curr_query->queryId; + new_query->ext_field = copyObject(curr_query->ext_field); new_list = lappend(new_list, new_query); pfree(str); @@ -956,6 +957,7 @@ pg_plan_queries(List *querytrees, const char *query_string, int cursorOptions, stmt->stmt_location = query->stmt_location; stmt->stmt_len = query->stmt_len; stmt->queryId = query->queryId; + stmt->ext_field = copyObject(query->ext_field); } else { diff --git a/src/include/commands/extension.h b/src/include/commands/extension.h index 74ae391395..a652fdbf1b 100644 --- a/src/include/commands/extension.h +++ b/src/include/commands/extension.h @@ -53,4 +53,8 @@ extern bool extension_file_exists(const char *extensionName); extern ObjectAddress AlterExtensionNamespace(const char *extensionName, const char *newschema, Oid *oldschema); +extern Node *AddExtensionDataToNode(Node *node, const char *entryname, + const Node *data, bool noCopy); +extern Node *GetExtensionData(ExtensionData *extdata, const char *entryname); + #endif /* EXTENSION_H */ diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h index 028588fb33..59a41f398d 100644 --- a/src/include/nodes/parsenodes.h +++ b/src/include/nodes/parsenodes.h @@ -101,6 +101,27 @@ typedef uint64 AclMode; /* a bitmask of privilege bits */ #define ACL_SELECT_FOR_UPDATE ACL_UPDATE +/* + * Structure definition should be hidden from an user, but current serialization + * techniques make it impossible. + * XXX: Instead of name, maybe better to invent some unique id for each loaded + * library, according to the position in file_list, as an example? + */ +typedef struct ExtensionData +{ + /* + * Don't allow reading the field (remember, extension may be dynamically + * loaded into a backend. Parallel workers still don't need it. + * Allow writing it for debug purposes. + */ + pg_node_attr(no_equal, no_read, no_query_jumble) + + NodeTag type; + char *name; + Node *node; + struct ExtensionData *next; +} ExtensionData; + /***************************************************************************** * Query Tree *****************************************************************************/ @@ -236,6 +257,8 @@ typedef struct Query int stmt_location; /* length in bytes; 0 means "rest of string" */ int stmt_len pg_node_attr(query_jumble_ignore); + + ExtensionData *ext_field pg_node_attr(query_jumble_ignore, read_write_ignore, equal_ignore, read_as(NULL)); } Query; diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h index 659bd05c0c..038a77e57e 100644 --- a/src/include/nodes/plannodes.h +++ b/src/include/nodes/plannodes.h @@ -101,6 +101,8 @@ typedef struct PlannedStmt /* statement location in source string (copied from Query) */ int stmt_location; /* start location, or -1 if unknown */ int stmt_len; /* length in bytes; 0 means "rest of string" */ + + ExtensionData *ext_field pg_node_attr(query_jumble_ignore, read_write_ignore, equal_ignore, read_as(NULL)); } PlannedStmt; /* macro for fetching the Plan associated with a SubPlan node */ diff --git a/src/include/nodes/value.h b/src/include/nodes/value.h index b24c4c1afe..6774a0d9a7 100644 --- a/src/include/nodes/value.h +++ b/src/include/nodes/value.h @@ -33,6 +33,14 @@ typedef struct Integer int ival; } Integer; +typedef struct Bigint +{ + pg_node_attr(special_read_write) + + NodeTag type; + int64 val; +} Bigint; + /* * Float is internally represented as string. Using T_Float as the node type * simply indicates that the contents of the string look like a valid numeric @@ -77,11 +85,13 @@ typedef struct BitString } BitString; #define intVal(v) (castNode(Integer, v)->ival) +#define bigintVal(v) (castNode(Bigint, v)->val) #define floatVal(v) atof(castNode(Float, v)->fval) #define boolVal(v) (castNode(Boolean, v)->boolval) #define strVal(v) (castNode(String, v)->sval) extern Integer *makeInteger(int i); +extern Bigint *makeBigint(int64 i); extern Float *makeFloat(char *numericStr); extern Boolean *makeBoolean(bool val); extern String *makeString(char *str); -- 2.34.1