From 4773328e7a212167cfd7570066c37167de407165 Mon Sep 17 00:00:00 2001
From: naisila <nicypp@gmail.com>
Date: Mon, 7 Jul 2025 16:00:46 +0300
Subject: [PATCH] Add citus_stat_user_tables from pg_stat_user_tables

---
 .../distributed/sql/citus--13.0-1--13.1-1.sql |  1 +
 .../sql/downgrades/citus--13.1-1--13.0-1.sql  |  1 +
 .../udfs/citus_stat_user_tables/13.1-1.sql    | 56 +++++++++++++++++++
 .../udfs/citus_stat_user_tables/latest.sql    | 56 +++++++++++++++++++
 4 files changed, 114 insertions(+)
 create mode 100644 src/backend/distributed/sql/udfs/citus_stat_user_tables/13.1-1.sql
 create mode 100644 src/backend/distributed/sql/udfs/citus_stat_user_tables/latest.sql

diff --git a/src/backend/distributed/sql/citus--13.0-1--13.1-1.sql b/src/backend/distributed/sql/citus--13.0-1--13.1-1.sql
index b8a6522de..096691c7c 100644
--- a/src/backend/distributed/sql/citus--13.0-1--13.1-1.sql
+++ b/src/backend/distributed/sql/citus--13.0-1--13.1-1.sql
@@ -52,6 +52,7 @@ DROP VIEW IF EXISTS pg_catalog.citus_lock_waits;
 #include "udfs/citus_stat_counters_reset/13.1-1.sql"
 #include "udfs/citus_nodes/13.1-1.sql"
 #include "udfs/citus_column_stats/13.1-1.sql"
+#include "udfs/citus_stat_user_tables/13.1-1.sql"
 
 -- Since shard_name/13.1-1.sql first drops the function and then creates it, we first
 -- need to drop citus_shards view since that view depends on this function. And immediately
diff --git a/src/backend/distributed/sql/downgrades/citus--13.1-1--13.0-1.sql b/src/backend/distributed/sql/downgrades/citus--13.1-1--13.0-1.sql
index 2cda5e66f..a324e70c3 100644
--- a/src/backend/distributed/sql/downgrades/citus--13.1-1--13.0-1.sql
+++ b/src/backend/distributed/sql/downgrades/citus--13.1-1--13.0-1.sql
@@ -47,6 +47,7 @@ DROP FUNCTION pg_catalog.citus_stat_counters(oid);
 DROP FUNCTION pg_catalog.citus_stat_counters_reset(oid);
 DROP VIEW IF EXISTS pg_catalog.citus_nodes;
 DROP FUNCTION IF EXISTS pg_catalog.citus_column_stats;
+DROP FUNCTION IF EXISTS pg_catalog.citus_stat_user_tables;
 
 -- Definition of shard_name() prior to this release doesn't have a separate SQL file
 -- because it's quite an old UDF that its prior definition(s) was(were) squashed into
diff --git a/src/backend/distributed/sql/udfs/citus_stat_user_tables/13.1-1.sql b/src/backend/distributed/sql/udfs/citus_stat_user_tables/13.1-1.sql
new file mode 100644
index 000000000..445a9da04
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_stat_user_tables/13.1-1.sql
@@ -0,0 +1,56 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables(
+RETURNS TABLE (
+    dist_table regclass,
+    n_tup_ins bigint,
+    n_tup_upd bigint,
+    n_tup_del bigint,
+    n_tup_hot_upd bigint,
+    n_tup_newpage_upd bigint,
+    n_live_tup bigint,
+    n_dead_tup bigint
+)
+AS $func$
+BEGIN
+    RETURN QUERY
+
+    WITH pg_dist_stats_double_json AS (
+        SELECT ( SELECT json_agg(row_to_json(f)) FROM ( SELECT result FROM
+            run_command_on_shards(logicalrelid, $$ SELECT json_agg(row_to_json(d))
+                FROM ( SELECT '$$ || logicalrelid || $$' AS dist_table,
+			            s.relname, n_tup_ins, n_tup_upd, n_tup_del,
+                        n_tup_hot_upd, n_tup_newpage_upd, n_live_tup, n_dead_tup
+                        FROM pg_stat_user_tables s
+                        JOIN pg_class c ON s.relname = c.relname
+                        WHERE c.oid = '%s'::regclass::oid) d $$)) f)
+            FROM pg_dist_partition),
+
+    pg_dist_stats_single_json AS (
+        SELECT (json_array_elements(json_agg)->>'result') AS result
+        FROM pg_dist_stats_double_json),
+
+    pg_dist_stats_regular AS (
+        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS dist_table,
+               (json_array_elements(result::json)->>'relname')::name AS shardname,
+               (json_array_elements(result::json)->>'n_tup_ins')::bigint AS n_tup_ins,
+               (json_array_elements(result::json)->>'n_tup_upd')::bigint AS n_tup_upd,
+               (json_array_elements(result::json)->>'n_tup_del')::bigint AS n_tup_del,
+               (json_array_elements(result::json)->>'n_tup_hot_upd')::bigint AS n_tup_hot_upd,
+               (json_array_elements(result::json)->>'n_tup_newpage_upd')::bigint AS n_tup_newpage_upd,
+               (json_array_elements(result::json)->>'n_live_tup')::bigint AS n_live_tup,
+               (json_array_elements(result::json)->>'n_dead_tup')::bigint AS n_dead_tup
+        FROM pg_dist_stats_single_json
+        WHERE result != '')
+
+    SELECT dist_table, sum(n_tup_ins)::bigint AS n_tup_ins, sum(n_tup_upd)::bigint AS n_tup_upd,
+           sum(n_tup_del)::bigint AS n_tup_del, sum(n_tup_hot_upd)::bigint AS n_tup_hot_upd,
+           sum(n_tup_newpage_upd)::bigint AS n_tup_newpage_upd,
+           sum(n_live_tup)::bigint AS n_live_tup, sum(n_dead_tup)::bigint AS n_dead_tup
+    FROM pg_dist_stats_regular
+    GROUP BY 1 ORDER BY 1;
+
+END;
+$func$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables(
+    qualified_table_name text)
+    IS 'provides some pg_stat_user_tables entries for Citus tables';
diff --git a/src/backend/distributed/sql/udfs/citus_stat_user_tables/latest.sql b/src/backend/distributed/sql/udfs/citus_stat_user_tables/latest.sql
new file mode 100644
index 000000000..445a9da04
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_stat_user_tables/latest.sql
@@ -0,0 +1,56 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables(
+RETURNS TABLE (
+    dist_table regclass,
+    n_tup_ins bigint,
+    n_tup_upd bigint,
+    n_tup_del bigint,
+    n_tup_hot_upd bigint,
+    n_tup_newpage_upd bigint,
+    n_live_tup bigint,
+    n_dead_tup bigint
+)
+AS $func$
+BEGIN
+    RETURN QUERY
+
+    WITH pg_dist_stats_double_json AS (
+        SELECT ( SELECT json_agg(row_to_json(f)) FROM ( SELECT result FROM
+            run_command_on_shards(logicalrelid, $$ SELECT json_agg(row_to_json(d))
+                FROM ( SELECT '$$ || logicalrelid || $$' AS dist_table,
+			            s.relname, n_tup_ins, n_tup_upd, n_tup_del,
+                        n_tup_hot_upd, n_tup_newpage_upd, n_live_tup, n_dead_tup
+                        FROM pg_stat_user_tables s
+                        JOIN pg_class c ON s.relname = c.relname
+                        WHERE c.oid = '%s'::regclass::oid) d $$)) f)
+            FROM pg_dist_partition),
+
+    pg_dist_stats_single_json AS (
+        SELECT (json_array_elements(json_agg)->>'result') AS result
+        FROM pg_dist_stats_double_json),
+
+    pg_dist_stats_regular AS (
+        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS dist_table,
+               (json_array_elements(result::json)->>'relname')::name AS shardname,
+               (json_array_elements(result::json)->>'n_tup_ins')::bigint AS n_tup_ins,
+               (json_array_elements(result::json)->>'n_tup_upd')::bigint AS n_tup_upd,
+               (json_array_elements(result::json)->>'n_tup_del')::bigint AS n_tup_del,
+               (json_array_elements(result::json)->>'n_tup_hot_upd')::bigint AS n_tup_hot_upd,
+               (json_array_elements(result::json)->>'n_tup_newpage_upd')::bigint AS n_tup_newpage_upd,
+               (json_array_elements(result::json)->>'n_live_tup')::bigint AS n_live_tup,
+               (json_array_elements(result::json)->>'n_dead_tup')::bigint AS n_dead_tup
+        FROM pg_dist_stats_single_json
+        WHERE result != '')
+
+    SELECT dist_table, sum(n_tup_ins)::bigint AS n_tup_ins, sum(n_tup_upd)::bigint AS n_tup_upd,
+           sum(n_tup_del)::bigint AS n_tup_del, sum(n_tup_hot_upd)::bigint AS n_tup_hot_upd,
+           sum(n_tup_newpage_upd)::bigint AS n_tup_newpage_upd,
+           sum(n_live_tup)::bigint AS n_live_tup, sum(n_dead_tup)::bigint AS n_dead_tup
+    FROM pg_dist_stats_regular
+    GROUP BY 1 ORDER BY 1;
+
+END;
+$func$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables(
+    qualified_table_name text)
+    IS 'provides some pg_stat_user_tables entries for Citus tables';