From baf6ab05908d75e8da88ad7affac8c28b5ead248 Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Thu, 3 Oct 2024 08:52:02 +0200 Subject: Track location to extract relevant part in nested statement Previously, Query generated through transform would have unset stmt_location. Extensions relying on the statement location to extract the relevant part of the statement would fallback to use the whole statement instead, thus showing the same string in the top and nested level which was a source of confusion. This patch fixes the issue by keeping track of the statement locations and propagate it to Query during transform, allowing pgss to only show the relevant part of the query for nested query. --- .../expected/level_tracking.out | 76 +++++++++---------- src/backend/parser/analyze.c | 29 +++++++ src/backend/parser/gram.y | 34 +++++++-- src/backend/parser/parse_merge.c | 2 + src/include/nodes/parsenodes.h | 6 ++ src/include/parser/parse_node.h | 2 + 6 files changed, 104 insertions(+), 45 deletions(-) diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out index 21a6c3ba7bd..74df1c3457d 100644 --- a/contrib/pg_stat_statements/expected/level_tracking.out +++ b/contrib/pg_stat_statements/expected/level_tracking.out @@ -192,15 +192,15 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | explain (costs off) SELECT $1 t | 1 | explain (costs off) SELECT $1 UNION SELECT $2 t | 1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2 - f | 1 | explain (costs off) (SELECT $1, $2); - f | 1 | explain (costs off) DELETE FROM stats_track_tab; - f | 1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1)); - f | 1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id+ + f | 1 | DELETE FROM stats_track_tab + f | 1 | INSERT INTO stats_track_tab VALUES (($1)) + f | 1 | MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id + | | WHEN MATCHED THEN UPDATE SET x = id + - | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id); - f | 1 | explain (costs off) SELECT $1 UNION SELECT $2; - f | 1 | explain (costs off) SELECT $1; - f | 1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2; + | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id) + f | 1 | SELECT $1 + f | 1 | SELECT $1 UNION SELECT $2 + f | 1 | SELECT $1, $2 + f | 1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2 (15 rows) -- Explain - top-level tracking. @@ -397,8 +397,8 @@ explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT SELECT toplevel, calls, query FROM pg_stat_statements ORDER BY toplevel desc, query COLLATE "C"; - toplevel | calls | query -----------+-------+---------------------------------------------------------------------------------------------------------------------------------- + toplevel | calls | query +----------+-------+------------------------------------------------------------------------------------------------------------ t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t t | 1 | explain (costs off) (SELECT $1, $2, $3) t | 1 | explain (costs off) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6 @@ -407,8 +407,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | explain (costs off) DELETE FROM stats_track_tab WHERE x=$1 t | 1 | explain (costs off) INSERT INTO stats_track_tab VALUES ($1), ($2) t | 1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1)) - t | 1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id + - | | WHEN MATCHED THEN UPDATE SET x = id + + t | 1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id+ + | | WHEN MATCHED THEN UPDATE SET x = id + | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id) t | 1 | explain (costs off) SELECT $1 t | 1 | explain (costs off) SELECT $1, $2 @@ -416,24 +416,22 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | explain (costs off) SELECT $1, $2, $3, $4, $5 t | 1 | explain (costs off) UPDATE stats_track_tab SET x=$1 t | 1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2 - f | 1 | explain (costs off) (SELECT $1, $2, $3); explain (costs off) (SELECT 1, 2, 3, 4); - f | 1 | explain (costs off) (SELECT 1, 2, 3); explain (costs off) (SELECT $1, $2, $3, $4); - f | 1 | explain (costs off) DELETE FROM stats_track_tab; explain (costs off) DELETE FROM stats_track_tab WHERE x=$1; - f | 1 | explain (costs off) DELETE FROM stats_track_tab; explain (costs off) DELETE FROM stats_track_tab WHERE x=1; - f | 1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1)); explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2); - f | 1 | explain (costs off) INSERT INTO stats_track_tab VALUES ((1)); explain (costs off) INSERT INTO stats_track_tab VALUES ($1), ($2); - f | 1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id + - | | WHEN MATCHED THEN UPDATE SET x = id + - | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id); explain (costs off) SELECT 1, 2, 3, 4, 5; - f | 1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series(1, 10) id) ON x = id + - | | WHEN MATCHED THEN UPDATE SET x = id + - | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id); explain (costs off) SELECT $1, $2, $3, $4, $5; - f | 1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5; - f | 1 | explain (costs off) SELECT $1; explain (costs off) SELECT 1, 2; - f | 1 | explain (costs off) SELECT 1, 2 UNION SELECT 3, 4; explain (costs off) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6; - f | 1 | explain (costs off) SELECT 1; explain (costs off) SELECT $1, $2; - f | 1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2; explain (costs off) UPDATE stats_track_tab SET x=1; - f | 1 | explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1; explain (costs off) UPDATE stats_track_tab SET x=$1; + f | 1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6 + f | 1 | DELETE FROM stats_track_tab + f | 1 | DELETE FROM stats_track_tab WHERE x=$1 + f | 1 | INSERT INTO stats_track_tab VALUES ($1), ($2) + f | 1 | INSERT INTO stats_track_tab VALUES (($1)) + f | 1 | MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id + + | | WHEN MATCHED THEN UPDATE SET x = id + + | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id) + f | 1 | SELECT $1 + f | 1 | SELECT $1, $2 + f | 1 | SELECT $1, $2 UNION SELECT $3, $4 + f | 1 | SELECT $1, $2, $3 + f | 1 | SELECT $1, $2, $3, $4 + f | 1 | SELECT $1, $2, $3, $4, $5 + f | 1 | UPDATE stats_track_tab SET x=$1 + f | 1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2 (29 rows) -- Explain - top-level tracking with multi statement. @@ -645,15 +643,15 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | explain (costs off) WITH a AS (select $1) SELECT $2 t | 1 | explain (costs off) WITH a AS (select $1) SELECT $2 UNION SELECT $3 t | 1 | explain (costs off) WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3 - f | 1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3)); - f | 1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab; - f | 1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2)); - f | 1 | explain (costs off) WITH a AS (select $1) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($2, $3) id) ON x = id+ + f | 1 | WITH a AS (select $1) (SELECT $2, $3) + f | 1 | WITH a AS (select $1) DELETE FROM stats_track_tab + f | 1 | WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2)) + f | 1 | WITH a AS (select $1) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($2, $3) id) ON x = id + | | WHEN MATCHED THEN UPDATE SET x = id + - | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id); - f | 1 | explain (costs off) WITH a AS (select $1) SELECT $2 UNION SELECT $3; - f | 1 | explain (costs off) WITH a AS (select $1) SELECT $2; - f | 1 | explain (costs off) WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3; + | | WHEN NOT MATCHED THEN INSERT (x) VALUES (id) + f | 1 | WITH a AS (select $1) SELECT $2 + f | 1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3 + f | 1 | WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3 (15 rows) -- Explain with CTE - top-level tracking @@ -765,7 +763,7 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab t | 1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1 t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t - f | 1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1; + f | 1 | SELECT $1 (4 rows) -- Explain analyze, top tracking. diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c index e901203424d..98ff3cd2ea5 100644 --- a/src/backend/parser/analyze.c +++ b/src/backend/parser/analyze.c @@ -113,6 +113,8 @@ parse_analyze_fixedparams(RawStmt *parseTree, const char *sourceText, Assert(sourceText != NULL); /* required as of 8.4 */ pstate->p_sourcetext = sourceText; + pstate->p_stmt_len = parseTree->stmt_len; + pstate->p_stmt_location = parseTree->stmt_location; if (numParams > 0) setup_parse_fixed_parameters(pstate, paramTypes, numParams); @@ -153,6 +155,8 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText, Assert(sourceText != NULL); /* required as of 8.4 */ pstate->p_sourcetext = sourceText; + pstate->p_stmt_len = parseTree->stmt_len; + pstate->p_stmt_location = parseTree->stmt_location; setup_parse_variable_parameters(pstate, paramTypes, numParams); @@ -195,6 +199,8 @@ parse_analyze_withcb(RawStmt *parseTree, const char *sourceText, Assert(sourceText != NULL); /* required as of 8.4 */ pstate->p_sourcetext = sourceText; + pstate->p_stmt_len = parseTree->stmt_len; + pstate->p_stmt_location = parseTree->stmt_location; pstate->p_queryEnv = queryEnv; (*parserSetup) (pstate, parserSetupArg); @@ -518,6 +524,8 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt) Node *qual; qry->commandType = CMD_DELETE; + qry->stmt_location = stmt->location; + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); /* process the WITH clause independently of all else */ if (stmt->withClause) @@ -606,6 +614,8 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt) Assert(pstate->p_ctenamespace == NIL); qry->commandType = CMD_INSERT; + qry->stmt_location = stmt->location; + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); pstate->p_is_insert = true; /* process the WITH clause independently of all else */ @@ -1347,6 +1357,21 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt) ListCell *l; qry->commandType = CMD_SELECT; + qry->stmt_location = stmt->location; + if (stmt->stmt_len > 0) + + /* + * If the select statement is within parentheses, stmt_len will be set + * and represent the length of the select within parentheses + */ + qry->stmt_len = stmt->stmt_len; + else + + /* + * Otherwise, we fallback to computing the length from the + * ParseState's length and location + */ + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); /* process the WITH clause independently of all else */ if (stmt->withClause) @@ -1730,6 +1755,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt) int tllen; qry->commandType = CMD_SELECT; + qry->stmt_location = stmt->location; + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); /* * Find leftmost leaf SelectStmt. We currently only need to do this in @@ -2429,6 +2456,8 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt) Node *qual; qry->commandType = CMD_UPDATE; + qry->stmt_location = stmt->location; + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); pstate->p_is_insert = false; /* process the WITH clause independently of all else */ diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y index 4aa8646af7b..0133187bb91 100644 --- a/src/backend/parser/gram.y +++ b/src/backend/parser/gram.y @@ -190,7 +190,7 @@ static void insertSelectOptions(SelectStmt *stmt, SelectLimit *limitClause, WithClause *withClause, core_yyscan_t yyscanner); -static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg); +static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location); static Node *doNegate(Node *n, int location); static void doNegateFloat(Float *v); static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location); @@ -12170,6 +12170,9 @@ InsertStmt: $5->onConflictClause = $6; $5->returningList = $7; $5->withClause = $1; + if (@$ < 0) /* see comments for YYLLOC_DEFAULT */ + @$ = @2; + $5->location = @$; $$ = (Node *) $5; } ; @@ -12323,6 +12326,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias n->whereClause = $6; n->returningList = $7; n->withClause = $1; + if (@$ < 0) /* see comments for YYLLOC_DEFAULT */ + @$ = @2; + n->location = @$; $$ = (Node *) n; } ; @@ -12397,6 +12403,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias n->whereClause = $7; n->returningList = $8; n->withClause = $1; + if (@$ < 0) /* see comments for YYLLOC_DEFAULT */ + @$ = @2; + n->location = @$; $$ = (Node *) n; } ; @@ -12474,6 +12483,9 @@ MergeStmt: m->joinCondition = $8; m->mergeWhenClauses = $9; m->returningList = $10; + if (@$ < 0) /* see comments for YYLLOC_DEFAULT */ + @$ = @2; + m->location = @$; $$ = (Node *) m; } @@ -12714,7 +12726,12 @@ SelectStmt: select_no_parens %prec UMINUS ; select_with_parens: - '(' select_no_parens ')' { $$ = $2; } + '(' select_no_parens ')' + { + SelectStmt *n = (SelectStmt *) $2; + n->stmt_len = @3 - @2; + $$ = $2; + } | '(' select_with_parens ')' { $$ = $2; } ; @@ -12836,6 +12853,7 @@ simple_select: n->groupDistinct = ($7)->distinct; n->havingClause = $8; n->windowClause = $9; + n->location = @1; $$ = (Node *) n; } | SELECT distinct_clause target_list @@ -12853,6 +12871,7 @@ simple_select: n->groupDistinct = ($7)->distinct; n->havingClause = $8; n->windowClause = $9; + n->location = @1; $$ = (Node *) n; } | values_clause { $$ = $1; } @@ -12877,15 +12896,15 @@ simple_select: } | select_clause UNION set_quantifier select_clause { - $$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4); + $$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1); } | select_clause INTERSECT set_quantifier select_clause { - $$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4); + $$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1); } | select_clause EXCEPT set_quantifier select_clause { - $$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4); + $$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1); } ; @@ -18963,11 +18982,13 @@ insertSelectOptions(SelectStmt *stmt, errmsg("multiple WITH clauses not allowed"), parser_errposition(exprLocation((Node *) withClause)))); stmt->withClause = withClause; + /* Update SelectStmt's location to the start of the with clause */ + stmt->location = withClause->location; } } static Node * -makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg) +makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location) { SelectStmt *n = makeNode(SelectStmt); @@ -18975,6 +18996,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg) n->all = all; n->larg = (SelectStmt *) larg; n->rarg = (SelectStmt *) rarg; + n->location = location; return (Node *) n; } diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c index 87df79027d7..f5f34116a61 100644 --- a/src/backend/parser/parse_merge.c +++ b/src/backend/parser/parse_merge.c @@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt) Assert(pstate->p_ctenamespace == NIL); qry->commandType = CMD_MERGE; + qry->stmt_location = stmt->location; + qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location); qry->hasRecursive = false; /* process the WITH clause independently of all else */ diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h index 1c314cd9074..b47fe03cdea 100644 --- a/src/include/nodes/parsenodes.h +++ b/src/include/nodes/parsenodes.h @@ -2045,6 +2045,7 @@ typedef struct InsertStmt List *returningList; /* list of expressions to return */ WithClause *withClause; /* WITH clause */ OverridingKind override; /* OVERRIDING clause */ + ParseLoc location; /* token location, or -1 if unknown */ } InsertStmt; /* ---------------------- @@ -2059,6 +2060,7 @@ typedef struct DeleteStmt Node *whereClause; /* qualifications */ List *returningList; /* list of expressions to return */ WithClause *withClause; /* WITH clause */ + ParseLoc location; /* token location, or -1 if unknown */ } DeleteStmt; /* ---------------------- @@ -2074,6 +2076,7 @@ typedef struct UpdateStmt List *fromClause; /* optional from clause for more tables */ List *returningList; /* list of expressions to return */ WithClause *withClause; /* WITH clause */ + ParseLoc location; /* token location, or -1 if unknown */ } UpdateStmt; /* ---------------------- @@ -2089,6 +2092,7 @@ typedef struct MergeStmt List *mergeWhenClauses; /* list of MergeWhenClause(es) */ List *returningList; /* list of expressions to return */ WithClause *withClause; /* WITH clause */ + ParseLoc location; /* token location, or -1 if unknown */ } MergeStmt; /* ---------------------- @@ -2158,6 +2162,8 @@ typedef struct SelectStmt bool all; /* ALL specified? */ struct SelectStmt *larg; /* left child */ struct SelectStmt *rarg; /* right child */ + ParseLoc location; /* token location, or -1 if unknown */ + ParseLoc stmt_len; /* Eventually add fields for CORRESPONDING spec here */ } SelectStmt; diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h index 543df568147..ba572b3aea2 100644 --- a/src/include/parser/parse_node.h +++ b/src/include/parser/parse_node.h @@ -193,6 +193,8 @@ struct ParseState { ParseState *parentParseState; /* stack link */ const char *p_sourcetext; /* source text, or NULL if not available */ + ParseLoc p_stmt_location; /* start location, or -1 if unknown */ + ParseLoc p_stmt_len; /* length in bytes; 0 means "rest of string" */ List *p_rtable; /* range table so far */ List *p_rteperminfos; /* list of RTEPermissionInfo nodes for each * RTE_RELATION entry in rtable */ -- 2.39.5 (Apple Git-154)