Show local managed tables in citus_tables and hide tables owned by extensions (#6321)

Co-authored-by: Marco Slot <marco.slot@gmail.com>

src/backend/distributed/sql/citus--11.0-4--11.1-1.sql

@@ -1,4 +1,6 @@
 #include "udfs/citus_locks/11.1-1.sql"
+#include "udfs/citus_tables/11.1-1.sql"
+#include "udfs/citus_shards/11.1-1.sql"
 #include "udfs/create_distributed_table_concurrently/11.1-1.sql"
 #include "udfs/citus_internal_delete_partition_metadata/11.1-1.sql"
 #include "udfs/citus_copy_shard_placement/11.1-1.sql"

src/backend/distributed/sql/downgrades/citus--11.1-1--11.0-4.sql

@@ -93,6 +93,9 @@ DROP FUNCTION pg_catalog.get_all_active_transactions(OUT datid oid, OUT process_
 DROP VIEW pg_catalog.citus_locks;
 DROP FUNCTION pg_catalog.citus_locks();
 
+#include "../udfs/citus_tables/10.0-4.sql"
+#include "../udfs/citus_shards/10.1-1.sql"
+
 DROP FUNCTION pg_catalog.replicate_reference_tables(citus.shard_transfer_mode);
 #include "../udfs/replicate_reference_tables/9.3-2.sql"
 

src/backend/distributed/sql/udfs/citus_shards/11.1-1.sql

@@ -0,0 +1,45 @@
+CREATE OR REPLACE VIEW pg_catalog.citus_shards AS
+SELECT
+     pg_dist_shard.logicalrelid AS table_name,
+     pg_dist_shard.shardid,
+     shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) as shard_name,
+     CASE WHEN partkey IS NOT NULL THEN 'distributed' WHEN repmodel = 't' THEN 'reference' ELSE 'local' END AS citus_table_type,
+     colocationid AS colocation_id,
+     pg_dist_node.nodename,
+     pg_dist_node.nodeport,
+     size as shard_size
+FROM
+   pg_dist_shard
+JOIN
+   pg_dist_placement
+ON
+   pg_dist_shard.shardid = pg_dist_placement.shardid
+JOIN
+   pg_dist_node
+ON
+   pg_dist_placement.groupid = pg_dist_node.groupid
+JOIN
+   pg_dist_partition
+ON
+   pg_dist_partition.logicalrelid = pg_dist_shard.logicalrelid
+LEFT JOIN
+   (SELECT (regexp_matches(table_name,'_(\d+)$'))[1]::int as shard_id, max(size) as size from citus_shard_sizes() GROUP BY shard_id) as shard_sizes
+ON
+    pg_dist_shard.shardid = shard_sizes.shard_id
+WHERE
+   pg_dist_placement.shardstate = 1
+AND
+   -- filter out tables owned by extensions
+   pg_dist_partition.logicalrelid NOT IN (
+      SELECT
+         objid
+      FROM
+         pg_depend
+      WHERE
+         classid = 'pg_class'::regclass AND refclassid = 'pg_extension'::regclass AND deptype = 'e'
+   )
+ORDER BY
+   pg_dist_shard.logicalrelid::text, shardid
+;
+
+GRANT SELECT ON pg_catalog.citus_shards TO public;

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

@@ -28,6 +28,16 @@ ON
     pg_dist_shard.shardid = shard_sizes.shard_id
 WHERE
    pg_dist_placement.shardstate = 1
+AND
+   -- filter out tables owned by extensions
+   pg_dist_partition.logicalrelid NOT IN (
+      SELECT
+         objid
+      FROM
+         pg_depend
+      WHERE
+         classid = 'pg_class'::regclass AND refclassid = 'pg_extension'::regclass AND deptype = 'e'
+   )
 ORDER BY
    pg_dist_shard.logicalrelid::text, shardid
 ;

src/backend/distributed/sql/udfs/citus_tables/11.1-1.sql

@@ -0,0 +1,48 @@
+DO $$
+declare
+citus_tables_create_query text;
+BEGIN
+citus_tables_create_query=$CTCQ$
+    CREATE OR REPLACE VIEW %I.citus_tables AS
+    SELECT
+        logicalrelid AS table_name,
+        CASE WHEN partkey IS NOT NULL THEN 'distributed' ELSE
+            CASE when repmodel = 't' THEN 'reference' ELSE 'local' END
+        END AS citus_table_type,
+        coalesce(column_to_column_name(logicalrelid, partkey), '<none>') AS distribution_column,
+        colocationid AS colocation_id,
+        pg_size_pretty(citus_total_relation_size(logicalrelid, fail_on_error := false)) AS table_size,
+        (select count(*) from pg_dist_shard where logicalrelid = p.logicalrelid) AS shard_count,
+        pg_get_userbyid(relowner) AS table_owner,
+        amname AS access_method
+    FROM
+        pg_dist_partition p
+    JOIN
+        pg_class c ON (p.logicalrelid = c.oid)
+    LEFT JOIN
+        pg_am a ON (a.oid = c.relam)
+    WHERE
+        -- filter out tables owned by extensions
+        logicalrelid NOT IN (
+            SELECT
+                objid
+            FROM
+                pg_depend
+            WHERE
+                classid = 'pg_class'::regclass AND refclassid = 'pg_extension'::regclass AND deptype = 'e'
+        )
+    ORDER BY
+        logicalrelid::text;
+$CTCQ$;
+
+IF EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = 'public') THEN
+    EXECUTE format(citus_tables_create_query, 'public');
+    GRANT SELECT ON public.citus_tables TO public;
+ELSE
+    EXECUTE format(citus_tables_create_query, 'citus');
+    ALTER VIEW citus.citus_tables SET SCHEMA pg_catalog;
+    GRANT SELECT ON pg_catalog.citus_tables TO public;
+END IF;
+
+END;
+$$;

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

@@ -6,7 +6,9 @@ citus_tables_create_query=$CTCQ$
     CREATE OR REPLACE VIEW %I.citus_tables AS
     SELECT
         logicalrelid AS table_name,
-        CASE WHEN partkey IS NOT NULL THEN 'distributed' ELSE 'reference' END AS citus_table_type,
+        CASE WHEN partkey IS NOT NULL THEN 'distributed' ELSE
+            CASE when repmodel = 't' THEN 'reference' ELSE 'local' END
+        END AS citus_table_type,
         coalesce(column_to_column_name(logicalrelid, partkey), '<none>') AS distribution_column,
         colocationid AS colocation_id,
         pg_size_pretty(citus_total_relation_size(logicalrelid, fail_on_error := false)) AS table_size,
@@ -20,7 +22,15 @@ citus_tables_create_query=$CTCQ$
     LEFT JOIN
         pg_am a ON (a.oid = c.relam)
     WHERE
-        partkey IS NOT NULL OR repmodel = 't'
+        -- filter out tables owned by extensions
+        logicalrelid NOT IN (
+            SELECT
+                objid
+            FROM
+                pg_depend
+            WHERE
+                classid = 'pg_class'::regclass AND refclassid = 'pg_extension'::regclass AND deptype = 'e'
+        )
     ORDER BY
         logicalrelid::text;
 $CTCQ$;

src/test/regress/expected/alter_table_set_access_method.out

@@ -456,7 +456,8 @@ SELECT table_name, citus_table_type, distribution_column, shard_count, access_me
 ---------------------------------------------------------------------
  table_type_dist        | distributed      | a                   |           4 | heap
  table_type_ref         | reference        | <none>              |           1 | heap
-(2 rows)
+ table_type_citus_local | local            | <none>              |           1 | heap
+(3 rows)
 
 SELECT c.relname, a.amname FROM pg_class c, pg_am a where c.relname SIMILAR TO 'table_type\D*' AND c.relnamespace = 'alter_table_set_access_method'::regnamespace AND c.relam = a.oid;
         relname         | amname
@@ -512,7 +513,8 @@ SELECT table_name, citus_table_type, distribution_column, shard_count, access_me
 ---------------------------------------------------------------------
  table_type_dist        | distributed      | a                   |           4 | columnar
  table_type_ref         | reference        | <none>              |           1 | columnar
-(2 rows)
+ table_type_citus_local | local            | <none>              |           1 | columnar
+(3 rows)
 
 SELECT c.relname, a.amname FROM pg_class c, pg_am a where c.relname SIMILAR TO 'table_type\D*' AND c.relnamespace = 'alter_table_set_access_method'::regnamespace AND c.relam = a.oid;
         relname         |  amname

src/test/regress/expected/multi_mx_ddl.out

@@ -278,6 +278,19 @@ HINT:  To remove the local data, run: SELECT truncate_local_data_after_distribut
 
 (1 row)
 
+-- table should not show up in citus_tables or citus_shards
+SELECT count(*) FROM citus_tables WHERE table_name = 'seg_test'::regclass;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT count(*) FROM citus_shards WHERE table_name = 'seg_test'::regclass;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
 \c - - - :worker_1_port
 -- should be able to see contents from worker
 SELECT * FROM seg_test;
@@ -286,6 +299,19 @@ SELECT * FROM seg_test;
  42
 (1 row)
 
+-- table should not show up in citus_tables or citus_shards
+SELECT count(*) FROM citus_tables WHERE table_name = 'seg_test'::regclass;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+SELECT count(*) FROM citus_shards WHERE table_name = 'seg_test'::regclass;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
 \c - - - :master_port
 -- test metadata sync in the presence of an extension-owned table
 SELECT start_metadata_sync_to_node('localhost', :worker_1_port);

src/test/regress/sql/multi_mx_ddl.sql

@@ -155,11 +155,19 @@ SET citus.shard_replication_factor TO 1;
 SET citus.shard_count TO 4;
 SELECT create_distributed_table('seg_test', 'x');
 
+-- table should not show up in citus_tables or citus_shards
+SELECT count(*) FROM citus_tables WHERE table_name = 'seg_test'::regclass;
+SELECT count(*) FROM citus_shards WHERE table_name = 'seg_test'::regclass;
+
 \c - - - :worker_1_port
 
 -- should be able to see contents from worker
 SELECT * FROM seg_test;
 
+-- table should not show up in citus_tables or citus_shards
+SELECT count(*) FROM citus_tables WHERE table_name = 'seg_test'::regclass;
+SELECT count(*) FROM citus_shards WHERE table_name = 'seg_test'::regclass;
+
 \c - - - :master_port
 
 -- test metadata sync in the presence of an extension-owned table