v9-0004-Use-view-s-definition-as-query-string-on-a-materi.patch

From 4308e20f7b5e6be26d40c69f41ed52185fa86bae 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               | 10 ++++----
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 src/include/utils/ruleutils.h                 |  2 +-
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d1638a23159..c85a9deeec9 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -959,12 +959,14 @@ 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 | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
- f        |     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
-(3 rows)
+(4 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..92d5dd43864 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,21 @@ 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 pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		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 +416,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 +424,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.5 (Apple Git-154)