Add citus_schemas view (#6979)

DESCRIPTION: Adds citus_schemas view

The citus_schemas view will be created in public schema if it exists, if
not the view will be created in pg_catalog.

Need to:
- [x] Add tests
- [x] Fix tests

src/backend/distributed/sql/citus--11.3-1--12.0-1.sql

@@ -31,6 +31,8 @@ DROP FUNCTION citus_shard_sizes;
 #include "udfs/citus_tables/12.0-1.sql"
 #include "udfs/citus_shards/12.0-1.sql"
 
+#include "udfs/citus_schemas/12.0-1.sql"
+
 -- udfs used to include schema-based tenants in tenant monitoring
 #include "udfs/citus_stat_tenants_local/12.0-1.sql"
 

src/backend/distributed/sql/downgrades/citus--12.0-1--11.3-1.sql

@@ -43,6 +43,10 @@ DROP FUNCTION pg_catalog.citus_internal_unregister_tenant_schema_globally(Oid, t
 
 #include "../udfs/citus_drop_trigger/10.2-1.sql"
 
+-- citus_schemas might be created in either of the schemas
+DROP VIEW IF EXISTS public.citus_schemas;
+DROP VIEW IF EXISTS pg_catalog.citus_schemas;
+
 DROP VIEW pg_catalog.citus_shards;
 DROP FUNCTION pg_catalog.citus_shard_sizes;
 #include "../udfs/citus_shard_sizes/10.0-1.sql"

src/backend/distributed/sql/udfs/citus_schemas/12.0-1.sql

@@ -0,0 +1,42 @@
+DO $$
+declare
+citus_schemas_create_query text;
+BEGIN
+citus_schemas_create_query=$CSCQ$
+    CREATE OR REPLACE VIEW %I.citus_schemas AS
+    SELECT
+        ts.schemaid::regnamespace AS schema_name,
+        ts.colocationid AS colocation_id,
+        CASE
+            WHEN pg_catalog.has_schema_privilege(CURRENT_USER, ts.schemaid::regnamespace, 'USAGE')
+            THEN pg_size_pretty(coalesce(schema_sizes.schema_size, 0))
+            ELSE NULL
+        END AS schema_size,
+        pg_get_userbyid(n.nspowner) AS schema_owner
+    FROM
+        pg_dist_schema ts
+    JOIN
+        pg_namespace n ON (ts.schemaid = n.oid)
+    LEFT JOIN (
+        SELECT c.relnamespace::regnamespace schema_id, SUM(size) AS schema_size
+        FROM citus_shard_sizes() css, pg_dist_shard ds, pg_class c
+        WHERE css.shard_id = ds.shardid AND ds.logicalrelid = c.oid
+        GROUP BY schema_id
+    ) schema_sizes ON schema_sizes.schema_id = ts.schemaid
+    ORDER BY
+        schema_name;
+$CSCQ$;
+
+IF EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = 'public') THEN
+    EXECUTE format(citus_schemas_create_query, 'public');
+    REVOKE ALL ON public.citus_schemas FROM public;
+    GRANT SELECT ON public.citus_schemas TO public;
+ELSE
+    EXECUTE format(citus_schemas_create_query, 'citus');
+    ALTER VIEW citus.citus_schemas SET SCHEMA pg_catalog;
+    REVOKE ALL ON pg_catalog.citus_schemas FROM public;
+    GRANT SELECT ON pg_catalog.citus_schemas TO public;
+END IF;
+
+END;
+$$;

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

@@ -0,0 +1,42 @@
+DO $$
+declare
+citus_schemas_create_query text;
+BEGIN
+citus_schemas_create_query=$CSCQ$
+    CREATE OR REPLACE VIEW %I.citus_schemas AS
+    SELECT
+        ts.schemaid::regnamespace AS schema_name,
+        ts.colocationid AS colocation_id,
+        CASE
+            WHEN pg_catalog.has_schema_privilege(CURRENT_USER, ts.schemaid::regnamespace, 'USAGE')
+            THEN pg_size_pretty(coalesce(schema_sizes.schema_size, 0))
+            ELSE NULL
+        END AS schema_size,
+        pg_get_userbyid(n.nspowner) AS schema_owner
+    FROM
+        pg_dist_schema ts
+    JOIN
+        pg_namespace n ON (ts.schemaid = n.oid)
+    LEFT JOIN (
+        SELECT c.relnamespace::regnamespace schema_id, SUM(size) AS schema_size
+        FROM citus_shard_sizes() css, pg_dist_shard ds, pg_class c
+        WHERE css.shard_id = ds.shardid AND ds.logicalrelid = c.oid
+        GROUP BY schema_id
+    ) schema_sizes ON schema_sizes.schema_id = ts.schemaid
+    ORDER BY
+        schema_name;
+$CSCQ$;
+
+IF EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = 'public') THEN
+    EXECUTE format(citus_schemas_create_query, 'public');
+    REVOKE ALL ON public.citus_schemas FROM public;
+    GRANT SELECT ON public.citus_schemas TO public;
+ELSE
+    EXECUTE format(citus_schemas_create_query, 'citus');
+    ALTER VIEW citus.citus_schemas SET SCHEMA pg_catalog;
+    REVOKE ALL ON pg_catalog.citus_schemas FROM public;
+    GRANT SELECT ON pg_catalog.citus_schemas TO public;
+END IF;
+
+END;
+$$;

src/test/regress/expected/multi_extension.out

@@ -88,10 +88,11 @@ WHERE pgd.refclassid = 'pg_extension'::regclass AND
 	  pge.extname    = 'citus' AND
 	  pgio.schema    NOT IN ('pg_catalog', 'citus', 'citus_internal', 'test', 'columnar', 'columnar_internal')
 ORDER BY 1, 2;
- type |      identity
+ type |       identity
 ---------------------------------------------------------------------
+ view | public.citus_schemas
  view | public.citus_tables
-(1 row)
+(2 rows)
 
 -- DROP EXTENSION pre-created by the regression suite
 DROP EXTENSION citus;
@@ -1367,7 +1368,8 @@ SELECT * FROM multi_extension.print_extension_changes();
                  | function citus_schema_undistribute(regnamespace) void
 				 | function citus_stat_tenants_local_internal(boolean) SETOF record
                  | table pg_dist_schema
-(7 rows)
+                 | view public.citus_schemas
+(8 rows)
 
 DROP TABLE multi_extension.prev_objects, multi_extension.extension_diff;
 -- show running version
@@ -1387,10 +1389,11 @@ WHERE pgd.refclassid = 'pg_extension'::regclass AND
 	  pge.extname    = 'citus' AND
 	  pgio.schema    NOT IN ('pg_catalog', 'citus', 'citus_internal', 'test', 'columnar', 'columnar_internal')
 ORDER BY 1, 2;
- type |      identity
+ type |       identity
 ---------------------------------------------------------------------
+ view | public.citus_schemas
  view | public.citus_tables
-(1 row)
+(2 rows)
 
 -- see incompatible version errors out
 RESET citus.enable_version_checks;
@@ -1498,10 +1501,11 @@ ALTER EXTENSION citus UPDATE;
 -- if cache is invalidated succesfull, this \d should work without any problem
 \d
             List of relations
- Schema |     Name     | Type |  Owner
+ Schema |     Name      | Type |  Owner
 ---------------------------------------------------------------------
- public | citus_tables | view | postgres
-(1 row)
+ public | citus_schemas | view | postgres
+ public | citus_tables  | view | postgres
+(2 rows)
 
 \c - - - :master_port
 -- test https://github.com/citusdata/citus/issues/3409

src/test/regress/expected/schema_based_sharding.out

@@ -1474,8 +1474,86 @@ SELECT table_name, citus_table_type FROM citus_shards WHERE table_name::text LIK
 (2 rows)
 
 RESET citus.enable_schema_based_sharding;
+-- test citus_schemas
+SET citus.enable_schema_based_sharding TO ON;
+CREATE USER citus_schema_role SUPERUSER;
+SET ROLE citus_schema_role;
+CREATE SCHEMA citus_sch1;
+CREATE TABLE citus_sch1.tbl1(a INT);
+CREATE TABLE citus_sch1.tbl2(a INT);
+RESET ROLE;
+CREATE SCHEMA citus_sch2;
+CREATE TABLE citus_sch2.tbl1(a INT);
+SET citus.enable_schema_based_sharding TO OFF;
+INSERT INTO citus_sch1.tbl1 SELECT * FROM generate_series(1, 10000);
+INSERT INTO citus_sch1.tbl2 SELECT * FROM generate_series(1, 5000);
+INSERT INTO citus_sch2.tbl1 SELECT * FROM generate_series(1, 12000);
+SELECT
+    cs.schema_name,
+    cs.colocation_id = ctc.colocation_id AS correct_colocation_id,
+    cs.schema_size = ctc.calculated_size AS correct_size,
+    cs.schema_owner
+FROM public.citus_schemas cs
+JOIN
+(
+    SELECT
+        c.relnamespace, ct.colocation_id,
+        pg_size_pretty(sum(citus_total_relation_size(ct.table_name))) AS calculated_size
+    FROM public.citus_tables ct, pg_class c
+    WHERE ct.table_name::oid = c.oid
+    GROUP BY 1, 2
+) ctc ON cs.schema_name = ctc.relnamespace
+WHERE cs.schema_name::text LIKE 'citus\_sch_'
+ORDER BY cs.schema_name::text;
+ schema_name | correct_colocation_id | correct_size |   schema_owner
+---------------------------------------------------------------------
+ citus_sch1  | t                     | t            | citus_schema_role
+ citus_sch2  | t                     | t            | postgres
+(2 rows)
+
+-- test empty schema and empty tables
+SET citus.enable_schema_based_sharding TO ON;
+CREATE SCHEMA citus_empty_sch1;
+CREATE SCHEMA citus_empty_sch2;
+CREATE TABLE citus_empty_sch2.tbl1(a INT);
+SET citus.enable_schema_based_sharding TO OFF;
+SELECT schema_name, schema_size FROM public.citus_schemas
+WHERE schema_name::text LIKE 'citus\_empty\_sch_' ORDER BY schema_name::text;
+   schema_name    | schema_size
+---------------------------------------------------------------------
+ citus_empty_sch1 | 0 bytes
+ citus_empty_sch2 | 0 bytes
+(2 rows)
+
+-- test with non-privileged role
+CREATE USER citus_schema_nonpri;
+SET ROLE citus_schema_nonpri;
+SET client_min_messages TO ERROR;
+SELECT schema_name, colocation_id > 0 AS colocation_id_visible, schema_size IS NOT NULL AS schema_size_visible, schema_owner
+FROM public.citus_schemas WHERE schema_name::text LIKE 'citus\_sch_' ORDER BY schema_name::text;
+ schema_name | colocation_id_visible | schema_size_visible |   schema_owner
+---------------------------------------------------------------------
+ citus_sch1  | t                     | f                   | citus_schema_role
+ citus_sch2  | t                     | f                   | postgres
+(2 rows)
+
+RESET client_min_messages;
+RESET ROLE;
+-- test using citus_tables from workers
+\c - - - :worker_1_port
+SELECT schema_name, colocation_id > 0 AS colocation_id_visible, schema_size IS NOT NULL AS schema_size_visible, schema_owner
+FROM public.citus_schemas WHERE schema_name::text LIKE 'citus\_sch_' ORDER BY schema_name::text;
+ schema_name | colocation_id_visible | schema_size_visible |   schema_owner
+---------------------------------------------------------------------
+ citus_sch1  | t                     | t                   | citus_schema_role
+ citus_sch2  | t                     | t                   | postgres
+(2 rows)
+
+\c - - - :master_port
+SET search_path TO regular_schema;
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2 CASCADE;
+DROP ROLE citus_schema_role, citus_schema_nonpri;
 SELECT citus_remove_node('localhost', :master_port);
  citus_remove_node
 ---------------------------------------------------------------------

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -327,6 +327,7 @@ ORDER BY 1;
  view citus_dist_stat_activity
  view citus_lock_waits
  view citus_locks
+ view citus_schema.citus_schemas
  view citus_schema.citus_tables
  view citus_shard_indexes_on_worker
  view citus_shards
@@ -337,5 +338,5 @@ ORDER BY 1;
  view citus_stat_tenants_local
  view pg_dist_shard_placement
  view time_partitions
-(329 rows)
+(330 rows)
 

src/test/regress/sql/schema_based_sharding.sql

@@ -994,7 +994,73 @@ SELECT table_name, citus_table_type FROM citus_shards WHERE table_name::text LIK
 
 RESET citus.enable_schema_based_sharding;
 
+-- test citus_schemas
+SET citus.enable_schema_based_sharding TO ON;
+CREATE USER citus_schema_role SUPERUSER;
+SET ROLE citus_schema_role;
+CREATE SCHEMA citus_sch1;
+CREATE TABLE citus_sch1.tbl1(a INT);
+CREATE TABLE citus_sch1.tbl2(a INT);
+RESET ROLE;
+
+CREATE SCHEMA citus_sch2;
+CREATE TABLE citus_sch2.tbl1(a INT);
+SET citus.enable_schema_based_sharding TO OFF;
+
+INSERT INTO citus_sch1.tbl1 SELECT * FROM generate_series(1, 10000);
+INSERT INTO citus_sch1.tbl2 SELECT * FROM generate_series(1, 5000);
+
+INSERT INTO citus_sch2.tbl1 SELECT * FROM generate_series(1, 12000);
+
+SELECT
+    cs.schema_name,
+    cs.colocation_id = ctc.colocation_id AS correct_colocation_id,
+    cs.schema_size = ctc.calculated_size AS correct_size,
+    cs.schema_owner
+FROM public.citus_schemas cs
+JOIN
+(
+    SELECT
+        c.relnamespace, ct.colocation_id,
+        pg_size_pretty(sum(citus_total_relation_size(ct.table_name))) AS calculated_size
+    FROM public.citus_tables ct, pg_class c
+    WHERE ct.table_name::oid = c.oid
+    GROUP BY 1, 2
+) ctc ON cs.schema_name = ctc.relnamespace
+WHERE cs.schema_name::text LIKE 'citus\_sch_'
+ORDER BY cs.schema_name::text;
+
+-- test empty schema and empty tables
+SET citus.enable_schema_based_sharding TO ON;
+CREATE SCHEMA citus_empty_sch1;
+
+CREATE SCHEMA citus_empty_sch2;
+CREATE TABLE citus_empty_sch2.tbl1(a INT);
+SET citus.enable_schema_based_sharding TO OFF;
+
+SELECT schema_name, schema_size FROM public.citus_schemas
+WHERE schema_name::text LIKE 'citus\_empty\_sch_' ORDER BY schema_name::text;
+
+-- test with non-privileged role
+CREATE USER citus_schema_nonpri;
+SET ROLE citus_schema_nonpri;
+
+SET client_min_messages TO ERROR;
+SELECT schema_name, colocation_id > 0 AS colocation_id_visible, schema_size IS NOT NULL AS schema_size_visible, schema_owner
+FROM public.citus_schemas WHERE schema_name::text LIKE 'citus\_sch_' ORDER BY schema_name::text;
+
+RESET client_min_messages;
+RESET ROLE;
+
+-- test using citus_tables from workers
+\c - - - :worker_1_port
+SELECT schema_name, colocation_id > 0 AS colocation_id_visible, schema_size IS NOT NULL AS schema_size_visible, schema_owner
+FROM public.citus_schemas WHERE schema_name::text LIKE 'citus\_sch_' ORDER BY schema_name::text;
+\c - - - :master_port
+SET search_path TO regular_schema;
+
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2 CASCADE;
+DROP ROLE citus_schema_role, citus_schema_nonpri;
 
 SELECT citus_remove_node('localhost', :master_port);