Set query ID for inner queries of CREATE TABLE AS and DECLARE

Some utility statements contain queries that can be planned and
executed: CREATE TABLE AS and DECLARE CURSOR.  This commit adds query ID
computation for the inner queries executed by these two utility
commands, with and without EXPLAIN.  This change leads to four new
callers of JumbleQuery() and post_parse_analyze_hook() so as extensions
can decide what to do with this new data.

Previously, extensions relying on the query ID, like pg_stat_statements,
were not able to track these nested queries as the query_id was 0.

For pg_stat_statements, this commit leads to additions under !toplevel
when pg_stat_statements.track is set to "all", as shown in its
regression tests.  The output of EXPLAIN for these two utilities gains a
"Query Identifier" if compute_query_id is enabled.

Author: Anthonin Bonnefoy
Reviewed-by: Michael Paquier, Jian He
Discussion: https://postgr.es/m/CAO6_XqqM6S9bQ2qd=75W+yKATwoazxSNhv5sjW06fjGAtHbTUA@mail.gmail.com

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 489dc7143f7aa821579489674c067ba557461ff0..9aee9f5010e0734ac2e55db2150e266abc9ea8ae 100644 (file)
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -924,8 +924,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
           |       |   DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
  t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
  f        |     1 | SELECT $1
+ f        |     1 | SELECT * FROM stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(4 rows)
+(5 rows)
 
 -- Explain analyze, top tracking.
 SET pg_stat_statements.track = 'top';
@@ -1047,9 +1048,10 @@ SELECT toplevel, calls, query FROM pg_stat_statements
 ----------+-------+-----------------------------------------------------------------
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
  f        |     1 | select generate_series($1, $2)
-(4 rows)
+(5 rows)
 
 -- CREATE TABLE AS, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1089,8 +1091,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                                   query                                   
 ----------+-------+---------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT $1
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- EXPLAIN with CREATE TABLE AS - top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1140,8 +1143,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | COMMIT
  t        |     1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab
  t        |     1 | FETCH FORWARD 1 FROM foocur
+ f        |     1 | SELECT * from stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(6 rows)
+(7 rows)
 
 -- DECLARE CURSOR, top-level tracking.
 SET pg_stat_statements.track = 'top';

diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
index aaeaaffe405cf8bbbe8b59a3c906b5a31386df1e..5c92e48a56c6e89e23a72345d974cba893c88ddb 100644 (file)
--- a/src/backend/commands/createas.c
+++ b/src/backend/commands/createas.c
@@ -37,6 +37,8 @@
 #include "commands/view.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
+#include "nodes/queryjumble.h"
+#include "parser/analyze.h"
 #include "rewrite/rewriteHandler.h"
 #include "tcop/tcopprot.h"
 #include "utils/builtins.h"
@@ -222,6 +224,7 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt,
 {
        Query      *query = castNode(Query, stmt->query);
        IntoClause *into = stmt->into;
+       JumbleState *jstate = NULL;
        bool            is_matview = (into->viewQuery != NULL);
        bool            do_refresh = false;
        DestReceiver *dest;
@@ -236,6 +239,13 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt,
         */
        dest = CreateIntoRelDestReceiver(into);
 
+       /* Query contained by CTAS needs to be jumbled if requested */
+       if (IsQueryIdEnabled())
+               jstate = JumbleQuery(query);
+
+       if (post_parse_analyze_hook)
+               (*post_parse_analyze_hook) (pstate, query, jstate);
+
        /*
         * The contained Query could be a SELECT, or an EXECUTE utility command.
         * If the latter, we just pass it off to ExecuteQuery.

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 18a5af6b919e717201cb3e8f28f76113ed393d58..7c0fd63b2f0b603522022cf07ffd7b29369484ef 100644 (file)
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -71,8 +71,7 @@ typedef struct SerializeMetrics
 
 static void ExplainOneQuery(Query *query, int cursorOptions,
                                                        IntoClause *into, ExplainState *es,
-                                                       const char *queryString, ParamListInfo params,
-                                                       QueryEnvironment *queryEnv);
+                                                       ParseState *pstate, ParamListInfo params);
 static void ExplainPrintJIT(ExplainState *es, int jit_flags,
                                                        JitInstrumentation *ji);
 static void ExplainPrintSerialize(ExplainState *es,
@@ -350,7 +349,7 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
                {
                        ExplainOneQuery(lfirst_node(Query, l),
                                                        CURSOR_OPT_PARALLEL_OK, NULL, es,
-                                                       pstate->p_sourcetext, params, pstate->p_queryEnv);
+                                                       pstate, params);
 
                        /* Separate plans with an appropriate separator */
                        if (lnext(rewritten, l) != NULL)
@@ -436,24 +435,22 @@ ExplainResultDesc(ExplainStmt *stmt)
 static void
 ExplainOneQuery(Query *query, int cursorOptions,
                                IntoClause *into, ExplainState *es,
-                               const char *queryString, ParamListInfo params,
-                               QueryEnvironment *queryEnv)
+                               ParseState *pstate, ParamListInfo params)
 {
        /* planner will not cope with utility statements */
        if (query->commandType == CMD_UTILITY)
        {
-               ExplainOneUtility(query->utilityStmt, into, es, queryString, params,
-                                                 queryEnv);
+               ExplainOneUtility(query->utilityStmt, into, es, pstate, params);
                return;
        }
 
        /* if an advisor plugin is present, let it manage things */
        if (ExplainOneQuery_hook)
                (*ExplainOneQuery_hook) (query, cursorOptions, into, es,
-                                                                queryString, params, queryEnv);
+                                                                pstate->p_sourcetext, params, pstate->p_queryEnv);
        else
                standard_ExplainOneQuery(query, cursorOptions, into, es,
-                                                                queryString, params, queryEnv);
+                                                                pstate->p_sourcetext, params, pstate->p_queryEnv);
 }
 
 /*
@@ -534,8 +531,7 @@ standard_ExplainOneQuery(Query *query, int cursorOptions,
  */
 void
 ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
-                                 const char *queryString, ParamListInfo params,
-                                 QueryEnvironment *queryEnv)
+                                 ParseState *pstate, ParamListInfo params)
 {
        if (utilityStmt == NULL)
                return;
@@ -547,7 +543,9 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
                 * ExplainOneQuery.  Copy to be safe in the EXPLAIN EXECUTE case.
                 */
                CreateTableAsStmt *ctas = (CreateTableAsStmt *) utilityStmt;
+               Query      *ctas_query;
                List       *rewritten;
+               JumbleState *jstate = NULL;
 
                /*
                 * Check if the relation exists or not.  This is done at this stage to
@@ -565,11 +563,16 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
                        return;
                }
 
-               rewritten = QueryRewrite(castNode(Query, copyObject(ctas->query)));
+               ctas_query = castNode(Query, copyObject(ctas->query));
+               if (IsQueryIdEnabled())
+                       jstate = JumbleQuery(ctas_query);
+               if (post_parse_analyze_hook)
+                       (*post_parse_analyze_hook) (pstate, ctas_query, jstate);
+               rewritten = QueryRewrite(ctas_query);
                Assert(list_length(rewritten) == 1);
                ExplainOneQuery(linitial_node(Query, rewritten),
                                                CURSOR_OPT_PARALLEL_OK, ctas->into, es,
-                                               queryString, params, queryEnv);
+                                               pstate, params);
        }
        else if (IsA(utilityStmt, DeclareCursorStmt))
        {
@@ -582,17 +585,25 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
                 * be created, however.
                 */
                DeclareCursorStmt *dcs = (DeclareCursorStmt *) utilityStmt;
+               Query      *dcs_query;
                List       *rewritten;
+               JumbleState *jstate = NULL;
 
-               rewritten = QueryRewrite(castNode(Query, copyObject(dcs->query)));
+               dcs_query = castNode(Query, copyObject(dcs->query));
+               if (IsQueryIdEnabled())
+                       jstate = JumbleQuery(dcs_query);
+               if (post_parse_analyze_hook)
+                       (*post_parse_analyze_hook) (pstate, dcs_query, jstate);
+
+               rewritten = QueryRewrite(dcs_query);
                Assert(list_length(rewritten) == 1);
                ExplainOneQuery(linitial_node(Query, rewritten),
                                                dcs->options, NULL, es,
-                                               queryString, params, queryEnv);
+                                               pstate, params);
        }
        else if (IsA(utilityStmt, ExecuteStmt))
                ExplainExecuteQuery((ExecuteStmt *) utilityStmt, into, es,
-                                                       queryString, params, queryEnv);
+                                                       pstate, params);
        else if (IsA(utilityStmt, NotifyStmt))
        {
                if (es->format == EXPLAIN_FORMAT_TEXT)

diff --git a/src/backend/commands/portalcmds.c b/src/backend/commands/portalcmds.c
index 4f6acf67198dd8ce94e6e774c04228f75eaf71b6..ac52ca25e99398ac34e43d4f8c34b31a5347d77b 100644 (file)
--- a/src/backend/commands/portalcmds.c
+++ b/src/backend/commands/portalcmds.c
@@ -28,6 +28,8 @@
 #include "executor/executor.h"
 #include "executor/tstoreReceiver.h"
 #include "miscadmin.h"
+#include "nodes/queryjumble.h"
+#include "parser/analyze.h"
 #include "rewrite/rewriteHandler.h"
 #include "tcop/pquery.h"
 #include "tcop/tcopprot.h"
@@ -44,6 +46,7 @@ PerformCursorOpen(ParseState *pstate, DeclareCursorStmt *cstmt, ParamListInfo pa
                                  bool isTopLevel)
 {
        Query      *query = castNode(Query, cstmt->query);
+       JumbleState *jstate = NULL;
        List       *rewritten;
        PlannedStmt *plan;
        Portal          portal;
@@ -71,6 +74,13 @@ PerformCursorOpen(ParseState *pstate, DeclareCursorStmt *cstmt, ParamListInfo pa
                                (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
                                 errmsg("cannot create a cursor WITH HOLD within security-restricted operation")));
 
+       /* Query contained by DeclareCursor needs to be jumbled if requested */
+       if (IsQueryIdEnabled())
+               jstate = JumbleQuery(query);
+
+       if (post_parse_analyze_hook)
+               (*post_parse_analyze_hook) (pstate, query, jstate);
+
        /*
         * Parse analysis was done already, but we still have to run the rule
         * rewriter.  We do not do AcquireRewriteLocks: we assume the query either

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 07257d4db94773ef47d0683f9c1cf9b320c4551a..a93f970a292ece3354e16505d653584ed6ecfbe0 100644 (file)
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -561,13 +561,12 @@ DropAllPreparedStatements(void)
  * "into" is NULL unless we are doing EXPLAIN CREATE TABLE AS EXECUTE,
  * in which case executing the query should result in creating that table.
  *
- * Note: the passed-in queryString is that of the EXPLAIN EXECUTE,
+ * Note: the passed-in pstate's queryString is that of the EXPLAIN EXECUTE,
  * not the original PREPARE; we get the latter string from the plancache.
  */
 void
 ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
-                                       const char *queryString, ParamListInfo params,
-                                       QueryEnvironment *queryEnv)
+                                       ParseState *pstate, ParamListInfo params)
 {
        PreparedStatement *entry;
        const char *query_string;
@@ -610,10 +609,10 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
        /* Evaluate parameters, if any */
        if (entry->plansource->num_params)
        {
-               ParseState *pstate;
+               ParseState *pstate_params;
 
-               pstate = make_parsestate(NULL);
-               pstate->p_sourcetext = queryString;
+               pstate_params = make_parsestate(NULL);
+               pstate_params->p_sourcetext = pstate->p_sourcetext;
 
                /*
                 * Need an EState to evaluate parameters; must not delete it till end
@@ -624,12 +623,12 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
                estate = CreateExecutorState();
                estate->es_param_list_info = params;
 
-               paramLI = EvaluateParams(pstate, entry, execstmt->params, estate);
+               paramLI = EvaluateParams(pstate_params, entry, execstmt->params, estate);
        }
 
        /* Replan if needed, and acquire a transient refcount */
        cplan = GetCachedPlan(entry->plansource, paramLI,
-                                                 CurrentResourceOwner, queryEnv);
+                                                 CurrentResourceOwner, pstate->p_queryEnv);
 
        INSTR_TIME_SET_CURRENT(planduration);
        INSTR_TIME_SUBTRACT(planduration, planstart);
@@ -655,12 +654,11 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
                PlannedStmt *pstmt = lfirst_node(PlannedStmt, p);
 
                if (pstmt->commandType != CMD_UTILITY)
-                       ExplainOnePlan(pstmt, into, es, query_string, paramLI, queryEnv,
+                       ExplainOnePlan(pstmt, into, es, query_string, paramLI, pstate->p_queryEnv,
                                                   &planduration, (es->buffers ? &bufusage : NULL),
                                                   es->memory ? &mem_counters : NULL);
                else
-                       ExplainOneUtility(pstmt->utilityStmt, into, es, query_string,
-                                                         paramLI, queryEnv);
+                       ExplainOneUtility(pstmt->utilityStmt, into, es, pstate, paramLI);
 
                /* No need for CommandCounterIncrement, as ExplainOnePlan did it */
 

diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 3ab0aae78f73fa573ba89cd47b6ed75ea38d9d93..aa5872bc154528b828475ae81e2b3c45a431c3eb 100644 (file)
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -100,8 +100,8 @@ extern ExplainState *NewExplainState(void);
 extern TupleDesc ExplainResultDesc(ExplainStmt *stmt);
 
 extern void ExplainOneUtility(Node *utilityStmt, IntoClause *into,
-                                                         ExplainState *es, const char *queryString,
-                                                         ParamListInfo params, QueryEnvironment *queryEnv);
+                                                         ExplainState *es, ParseState *pstate,
+                                                         ParamListInfo params);
 
 extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into,
                                                   ExplainState *es, const char *queryString,

diff --git a/src/include/commands/prepare.h b/src/include/commands/prepare.h
index 61472c111d6ed0930a20a272df8f49a8c1c2cb84..e6fd400e027b2ad3af1e6561bb813b8742c762cb 100644 (file)
--- a/src/include/commands/prepare.h
+++ b/src/include/commands/prepare.h
@@ -43,8 +43,8 @@ extern void ExecuteQuery(ParseState *pstate,
                                                 DestReceiver *dest, QueryCompletion *qc);
 extern void DeallocateQuery(DeallocateStmt *stmt);
 extern void ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into,
-                                                               ExplainState *es, const char *queryString,
-                                                               ParamListInfo params, QueryEnvironment *queryEnv);
+                                                               ExplainState *es, ParseState *pstate,
+                                                               ParamListInfo params);
 
 /* Low-level access to stored prepared statements */
 extern void StorePreparedStatement(const char *stmt_name,

diff --git a/src/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index dcbdaa038850af4eba1027837a2147201264a822..d2eef8097cfd02a5dd236e5d6733bb62a5c47192 100644 (file)
--- a/src/test/regress/expected/explain.out
+++ b/src/test/regress/expected/explain.out
@@ -662,6 +662,23 @@ select explain_filter('explain (verbose) select * from int8_tbl i8');
  Query Identifier: N
 (3 rows)
 
+-- Test compute_query_id with utility statements containing plannable query
+select explain_filter('explain (verbose) declare test_cur cursor for select * from int8_tbl');
+                       explain_filter                        
+-------------------------------------------------------------
+ Seq Scan on public.int8_tbl  (cost=N.N..N.N rows=N width=N)
+   Output: q1, q2
+ Query Identifier: N
+(3 rows)
+
+select explain_filter('explain (verbose) create table test_ctas as select 1');
+             explain_filter             
+----------------------------------------
+ Result  (cost=N.N..N.N rows=N width=N)
+   Output: N
+ Query Identifier: N
+(3 rows)
+
 -- Test SERIALIZE option
 select explain_filter('explain (analyze,serialize) select * from int8_tbl i8');
                                         explain_filter                                         

diff --git a/src/test/regress/sql/explain.sql b/src/test/regress/sql/explain.sql
index b861e2b53d57354f0caf527d649dbe5199cf577c..3ca285a1d7e5b7a5e053255f443b11e163cf4b45 100644 (file)
--- a/src/test/regress/sql/explain.sql
+++ b/src/test/regress/sql/explain.sql
@@ -163,6 +163,10 @@ select explain_filter('explain (verbose) select * from t1 where pg_temp.mysin(f1
 set compute_query_id = on;
 select explain_filter('explain (verbose) select * from int8_tbl i8');
 
+-- Test compute_query_id with utility statements containing plannable query
+select explain_filter('explain (verbose) declare test_cur cursor for select * from int8_tbl');
+select explain_filter('explain (verbose) create table test_ctas as select 1');
+
 -- Test SERIALIZE option
 select explain_filter('explain (analyze,serialize) select * from int8_tbl i8');
 select explain_filter('explain (analyze,serialize text,buffers,timing off) select * from int8_tbl i8');