From 084205ee02f796ca9a3f6861c0f4866309da55f9 Mon Sep 17 00:00:00 2001
From: reshke kirill <reshke@double.cloud>
Date: Wed, 4 Dec 2024 14:45:03 +0000
Subject: [PATCH v1] Add pg_stat_progress_rewrite view.

---
 src/backend/catalog/system_views.sql |  8 ++++++++
 src/backend/commands/tablecmds.c     | 11 +++++++++++
 src/backend/utils/adt/pgstatfuncs.c  |  2 ++
 src/include/commands/progress.h      |  3 +++
 src/include/utils/backend_progress.h |  1 +
 5 files changed, 25 insertions(+)

diff --git a/src/backend/catalog/system_views.sql b/src/backend/catalog/system_views.sql
index da9a8fe99f2..5e892c13bce 100644
--- a/src/backend/catalog/system_views.sql
+++ b/src/backend/catalog/system_views.sql
@@ -1325,6 +1325,14 @@ CREATE VIEW pg_stat_progress_copy AS
     FROM pg_stat_get_progress_info('COPY') AS S
         LEFT JOIN pg_database D ON S.datid = D.oid;
 
+CREATE VIEW pg_stat_progress_rewrite AS
+    SELECT
+        S.pid AS pid, S.datid AS datid, D.datname AS datname,
+        S.relid AS relid,
+        S.param1 AS tuples_processed
+    FROM pg_stat_get_progress_info('REWRITE') AS S
+        LEFT JOIN pg_database D ON S.datid = D.oid;
+
 CREATE VIEW pg_user_mappings AS
     SELECT
         U.oid       AS umid,
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 6ccae4cb4a8..be37cf59bf3 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -59,6 +59,7 @@
 #include "commands/comment.h"
 #include "commands/defrem.h"
 #include "commands/event_trigger.h"
+#include "commands/progress.h"
 #include "commands/sequence.h"
 #include "commands/tablecmds.h"
 #include "commands/tablespace.h"
@@ -6148,6 +6149,7 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 		List	   *dropped_attrs = NIL;
 		ListCell   *lc;
 		Snapshot	snapshot;
+		int64		nprocessed;
 
 		if (newrel)
 			ereport(DEBUG1,
@@ -6221,6 +6223,10 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 		snapshot = RegisterSnapshot(GetLatestSnapshot());
 		scan = table_beginscan(oldrel, snapshot, 0, NULL);
 
+		/* initialize progress */
+		pgstat_progress_start_command(PROGRESS_COMMAND_REWRITE, RelationGetRelid(oldrel));
+		nprocessed = 0;
+
 		/*
 		 * Switch to per-tuple memory context and reset it for each tuple
 		 * produced, so we don't leak memory.
@@ -6378,9 +6384,14 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 
 			ResetExprContext(econtext);
 
+			++nprocessed;
+			pgstat_progress_update_param(PROGRESS_REWRITE_TUPLES_PROCESSED, nprocessed);
+
 			CHECK_FOR_INTERRUPTS();
 		}
 
+		pgstat_progress_end_command();
+
 		MemoryContextSwitchTo(oldCxt);
 		table_endscan(scan);
 		UnregisterSnapshot(snapshot);
diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 60a397dc561..e36a201aa3c 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -246,6 +246,8 @@ pg_stat_get_progress_info(PG_FUNCTION_ARGS)
 		cmdtype = PROGRESS_COMMAND_BASEBACKUP;
 	else if (pg_strcasecmp(cmd, "COPY") == 0)
 		cmdtype = PROGRESS_COMMAND_COPY;
+	else if (pg_strcasecmp(cmd, "REWRITE") == 0)
+		cmdtype = PROGRESS_COMMAND_REWRITE;
 	else
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
diff --git a/src/include/commands/progress.h b/src/include/commands/progress.h
index 5616d645230..fefbf9abd6d 100644
--- a/src/include/commands/progress.h
+++ b/src/include/commands/progress.h
@@ -155,4 +155,7 @@
 #define PROGRESS_COPY_TYPE_PIPE 3
 #define PROGRESS_COPY_TYPE_CALLBACK 4
 
+/* Progress parameters for PROGRESS_REWRITE  */
+#define PROGRESS_REWRITE_TUPLES_PROCESSED 0
+
 #endif
diff --git a/src/include/utils/backend_progress.h b/src/include/utils/backend_progress.h
index 7b63d38f975..f3e7067e549 100644
--- a/src/include/utils/backend_progress.h
+++ b/src/include/utils/backend_progress.h
@@ -28,6 +28,7 @@ typedef enum ProgressCommandType
 	PROGRESS_COMMAND_CREATE_INDEX,
 	PROGRESS_COMMAND_BASEBACKUP,
 	PROGRESS_COMMAND_COPY,
+	PROGRESS_COMMAND_REWRITE,
 } ProgressCommandType;
 
 #define PGSTAT_NUM_PROGRESS_PARAM	20
-- 
2.34.1