v5-0003-Use-view-s-definition-as-query-string-on-a-materi.patch

From 74d897c037b6dcdf642a93d4d8e2aeadf80beb97 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 38 +++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql | 14 +++++++
 src/backend/commands/matview.c                | 23 +++++++----
 src/backend/utils/adt/ruleutils.c             | 13 +++++++
 src/include/utils/ruleutils.h                 |  2 +-
 5 files changed, 82 insertions(+), 8 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index c6b011caa42..06e64dab346 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -271,6 +271,44 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series($1, $2) as id
+ f        |     1 | SELECT * FROM generate_series($1, $2) as id;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(4 rows)
+
 -- Create Table As, all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
diff --git a/contrib/pg_stat_statements/sql/level_tracking.sql b/contrib/pg_stat_statements/sql/level_tracking.sql
index accd4dfce26..f1af0cfd743 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -81,6 +81,20 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- Create Table As, all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..784fd9b0425 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ static bool transientrel_receive(TupleTableSlot *slot, DestReceiver *self);
 static void transientrel_shutdown(DestReceiver *self);
 static void transientrel_destroy(DestReceiver *self);
 static uint64 refresh_matview_datafill(DestReceiver *dest, Query *query,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, bool is_create);
 static char *make_temptable_name_n(char *tempname, int n);
 static void refresh_by_match_merge(Oid matviewOid, Oid tempOid, Oid relowner,
 								   int save_sec_context);
@@ -327,11 +327,20 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the queryString will be the refresh utility statement.
+		 * We need to fetch the query string from the view definition.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +415,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +423,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.3 (Apple Git-146)