Add tests for citus_stat_user_tables

src/backend/distributed/sql/udfs/citus_stat_user_tables/13.1-1.sql

@@ -1,6 +1,6 @@
-CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables(
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables()
 RETURNS TABLE (
-    dist_table regclass,
+    relname regclass,
     n_tup_ins bigint,
     n_tup_upd bigint,
     n_tup_del bigint,
@@ -17,8 +17,8 @@ BEGIN
         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,
+			            s.relname, s.n_tup_ins, s.n_tup_upd, s.n_tup_del,
-                        n_tup_hot_upd, n_tup_newpage_upd, n_live_tup, n_dead_tup
+                        s.n_tup_hot_upd, s.n_tup_newpage_upd, s.n_live_tup, s.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)
@@ -29,7 +29,7 @@ BEGIN
         FROM pg_dist_stats_double_json),
 
     pg_dist_stats_regular AS (
-        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS dist_table,
+        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS relname,
                (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,
@@ -41,16 +41,15 @@ BEGIN
         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,
+    SELECT s.relname, sum(s.n_tup_ins)::bigint AS n_tup_ins, sum(s.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(s.n_tup_del)::bigint AS n_tup_del, sum(s.n_tup_hot_upd)::bigint AS n_tup_hot_upd,
-           sum(n_tup_newpage_upd)::bigint AS n_tup_newpage_upd,
+           sum(s.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
+           sum(s.n_live_tup)::bigint AS n_live_tup, sum(s.n_dead_tup)::bigint AS n_dead_tup
-    FROM pg_dist_stats_regular
+    FROM pg_dist_stats_regular s
     GROUP BY 1 ORDER BY 1;
 
 END;
 $func$ LANGUAGE plpgsql;
 
-COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables(
+COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables()
-    qualified_table_name text)
     IS 'provides some pg_stat_user_tables entries for Citus tables';

src/backend/distributed/sql/udfs/citus_stat_user_tables/latest.sql

@@ -1,6 +1,6 @@
-CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables(
+CREATE OR REPLACE FUNCTION pg_catalog.citus_stat_user_tables()
 RETURNS TABLE (
-    dist_table regclass,
+    relname regclass,
     n_tup_ins bigint,
     n_tup_upd bigint,
     n_tup_del bigint,
@@ -17,8 +17,8 @@ BEGIN
         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,
+			            s.relname, s.n_tup_ins, s.n_tup_upd, s.n_tup_del,
-                        n_tup_hot_upd, n_tup_newpage_upd, n_live_tup, n_dead_tup
+                        s.n_tup_hot_upd, s.n_tup_newpage_upd, s.n_live_tup, s.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)
@@ -29,7 +29,7 @@ BEGIN
         FROM pg_dist_stats_double_json),
 
     pg_dist_stats_regular AS (
-        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS dist_table,
+        SELECT (json_array_elements(result::json)->>'dist_table')::regclass AS relname,
                (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,
@@ -41,16 +41,15 @@ BEGIN
         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,
+    SELECT s.relname, sum(s.n_tup_ins)::bigint AS n_tup_ins, sum(s.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(s.n_tup_del)::bigint AS n_tup_del, sum(s.n_tup_hot_upd)::bigint AS n_tup_hot_upd,
-           sum(n_tup_newpage_upd)::bigint AS n_tup_newpage_upd,
+           sum(s.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
+           sum(s.n_live_tup)::bigint AS n_live_tup, sum(s.n_dead_tup)::bigint AS n_dead_tup
-    FROM pg_dist_stats_regular
+    FROM pg_dist_stats_regular s
     GROUP BY 1 ORDER BY 1;
 
 END;
 $func$ LANGUAGE plpgsql;
 
-COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables(
+COMMENT ON FUNCTION pg_catalog.citus_stat_user_tables()
-    qualified_table_name text)
     IS 'provides some pg_stat_user_tables entries for Citus tables';

src/test/regress/expected/citus_aggregated_stats.out

@@ -60,5 +60,72 @@ SELECT * FROM citus_column_stats('dist_current_check');
 DROP TABLE current_check;
 DROP TABLE dist_current_check;
 RESET SESSION AUTHORIZATION;
+RESET row_security;
+-- compare pg_stat_user_tables with citus_stat_user_tables
+CREATE TABLE trunc_stats_test(id serial);
+CREATE TABLE trunc_stats_dist_test(id serial);
+SELECT create_distributed_table('trunc_stats_dist_test', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- rollback a savepoint: this should count 4 inserts and have 2
+-- live tuples after commit (and 2 dead ones due to aborted subxact)
+BEGIN;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+SAVEPOINT p1;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+TRUNCATE trunc_stats_test;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+TRUNCATE trunc_stats_dist_test;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+ROLLBACK TO SAVEPOINT p1;
+COMMIT;
+\c - - - :worker_1_port
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush
+---------------------------------------------------------------------
+
+(1 row)
+
+\c - - - :worker_2_port
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush
+---------------------------------------------------------------------
+
+(1 row)
+
+\c - - - :master_port
+SELECT pg_stat_force_next_flush();
+ pg_stat_force_next_flush
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT relname, n_tup_ins, n_live_tup, n_dead_tup
+FROM pg_stat_user_tables
+WHERE relname like 'trunc_stats%';
+        relname        | n_tup_ins | n_live_tup | n_dead_tup
+---------------------------------------------------------------------
+ trunc_stats_test      |         4 |          2 |          2
+ trunc_stats_dist_test |         0 |          0 |          0
+(2 rows)
+
+SELECT relname, n_tup_ins, n_live_tup, n_dead_tup
+FROM citus_stat_user_tables();
+        relname        | n_tup_ins | n_live_tup | n_dead_tup
+---------------------------------------------------------------------
+ trunc_stats_dist_test |         4 |          2 |          2
+(1 row)
+
 REVOKE ALL ON SCHEMA public FROM user1;
 DROP USER user1;
+DROP TABLE trunc_stats_test;
+DROP TABLE trunc_stats_dist_test;

src/test/regress/expected/multi_extension.out

@@ -1484,11 +1484,12 @@ SELECT * FROM multi_extension.print_extension_changes();
                                                                 | function citus_is_primary_node() boolean
                                                                 | function citus_stat_counters(oid) SETOF record
                                                                 | function citus_stat_counters_reset(oid) void
+                                                                | function citus_stat_user_tables() TABLE(relname 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)
                                                                 | function citus_unmark_object_distributed(oid,oid,integer,boolean) void
                                                                 | function shard_name(regclass,bigint,boolean) text
                                                                 | view citus_nodes
                                                                 | view citus_stat_counters
-(34 rows)
+(35 rows)
 
 -- Test downgrade to 13.1-1 from 13.2-1
 ALTER EXTENSION citus UPDATE TO '13.2-1';

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -189,6 +189,7 @@ ORDER BY 1;
  function citus_stat_tenants_local_internal(boolean)
  function citus_stat_tenants_local_reset()
  function citus_stat_tenants_reset()
+ function citus_stat_user_tables()
  function citus_table_is_visible(oid)
  function citus_table_size(regclass)
  function citus_task_wait(bigint,citus_task_status)
@@ -394,6 +395,6 @@ ORDER BY 1;
  view citus_stat_tenants_local
  view pg_dist_shard_placement
  view time_partitions
-(363 rows)
+(364 rows)
 
 DROP TABLE extension_basic_types;

src/test/regress/sql/citus_aggregated_stats.sql

@@ -51,5 +51,53 @@ DROP TABLE current_check;
 DROP TABLE dist_current_check;
 
 RESET SESSION AUTHORIZATION;
+RESET row_security;
+
+-- compare pg_stat_user_tables with citus_stat_user_tables
+CREATE TABLE trunc_stats_test(id serial);
+CREATE TABLE trunc_stats_dist_test(id serial);
+SELECT create_distributed_table('trunc_stats_dist_test', 'id');
+
+-- rollback a savepoint: this should count 4 inserts and have 2
+-- live tuples after commit (and 2 dead ones due to aborted subxact)
+BEGIN;
+
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+
+SAVEPOINT p1;
+
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+TRUNCATE trunc_stats_test;
+INSERT INTO trunc_stats_test DEFAULT VALUES;
+
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+TRUNCATE trunc_stats_dist_test;
+INSERT INTO trunc_stats_dist_test DEFAULT VALUES;
+
+ROLLBACK TO SAVEPOINT p1;
+COMMIT;
+
+\c - - - :worker_1_port
+SELECT pg_stat_force_next_flush();
+\c - - - :worker_2_port
+SELECT pg_stat_force_next_flush();
+\c - - - :master_port
+SELECT pg_stat_force_next_flush();
+
+SELECT relname, n_tup_ins, n_live_tup, n_dead_tup
+FROM pg_stat_user_tables
+WHERE relname like 'trunc_stats%';
+
+SELECT relname, n_tup_ins, n_live_tup, n_dead_tup
+FROM citus_stat_user_tables();
+
 REVOKE ALL ON SCHEMA public FROM user1;
 DROP USER user1;
+DROP TABLE trunc_stats_test;
+DROP TABLE trunc_stats_dist_test;