Fix citus_table_type column in citus_tables and citus_shards views for single shard tables (#6971)

`citus_table_type` column of `citus_tables` and `citus_shards` will show
"schema" for tenants schema tables and "distributed" for single shard
tables that are not in a tenant schema.

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

@@ -23,3 +23,6 @@ GRANT SELECT ON pg_catalog.pg_dist_tenant_schema TO public;
 -- udfs used to modify pg_dist_tenant_schema globally via drop trigger
 #include "udfs/citus_internal_unregister_tenant_schema_globally/12.0-1.sql"
 #include "udfs/citus_drop_trigger/12.0-1.sql"
+
+#include "udfs/citus_tables/12.0-1.sql"
+#include "udfs/citus_shards/12.0-1.sql"

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

@@ -40,4 +40,7 @@ DROP FUNCTION pg_catalog.citus_internal_unregister_tenant_schema_globally(Oid, t
 
 #include "../udfs/citus_drop_trigger/10.2-1.sql"
 
+#include "../udfs/citus_tables/11.1-1.sql"
+#include "../udfs/citus_shards/11.1-1.sql"
+
 DROP TABLE pg_catalog.pg_dist_tenant_schema;

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

@@ -0,0 +1,49 @@
+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 colocationid IN (SELECT colocationid FROM pg_dist_tenant_schema) THEN 'schema'
+      WHEN partkey IS NOT NULL THEN 'distributed'
+      WHEN repmodel = 't' THEN 'reference'
+      WHEN colocationid = 0 THEN 'local'
+      ELSE 'distributed' 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

@@ -3,7 +3,11 @@ 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,
+     CASE WHEN colocationid IN (SELECT colocationid FROM pg_dist_tenant_schema) THEN 'schema'
+      WHEN partkey IS NOT NULL THEN 'distributed'
+      WHEN repmodel = 't' THEN 'reference'
+      WHEN colocationid = 0 THEN 'local'
+      ELSE 'distributed' END AS citus_table_type,
      colocationid AS colocation_id,
      pg_dist_node.nodename,
      pg_dist_node.nodeport,

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

@@ -0,0 +1,51 @@
+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 colocationid IN (SELECT colocationid FROM pg_dist_tenant_schema) THEN 'schema'
+            WHEN partkey IS NOT NULL THEN 'distributed'
+            WHEN repmodel = 't' THEN 'reference'
+            WHEN colocationid = 0 THEN 'local'
+            ELSE 'distributed'
+        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,8 +6,11 @@ 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
+        CASE WHEN colocationid IN (SELECT colocationid FROM pg_dist_tenant_schema) THEN 'schema'
+            WHEN partkey IS NOT NULL THEN 'distributed'
+            WHEN repmodel = 't' THEN 'reference'
+            WHEN colocationid = 0 THEN 'local'
+            ELSE 'distributed'
         END AS citus_table_type,
         coalesce(column_to_column_name(logicalrelid, partkey), '<none>') AS distribution_column,
         colocationid AS colocation_id,

src/test/regress/expected/schema_based_sharding.out

@@ -1379,8 +1379,34 @@ SELECT pg_reload_conf();
 (1 row)
 
 \c - - - :master_port
+SET search_path TO regular_schema;
+CREATE TABLE type_sing(a INT);
+SELECT create_distributed_table('type_sing', NULL, colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SET citus.enable_schema_based_sharding TO ON;
+CREATE SCHEMA type_sch;
+CREATE TABLE type_sch.tbl (a INT);
+SELECT table_name, citus_table_type FROM public.citus_tables WHERE table_name::text LIKE 'type_%';
+  table_name  | citus_table_type
+---------------------------------------------------------------------
+ type_sch.tbl | schema
+ type_sing    | distributed
+(2 rows)
+
+SELECT table_name, citus_table_type FROM citus_shards WHERE table_name::text LIKE 'type_%' AND nodeport IN (:worker_1_port, :worker_2_port);
+  table_name  | citus_table_type
+---------------------------------------------------------------------
+ type_sch.tbl | schema
+ type_sing    | distributed
+(2 rows)
+
+RESET citus.enable_schema_based_sharding;
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6 CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch CASCADE;
 SELECT citus_remove_node('localhost', :master_port);
  citus_remove_node
 ---------------------------------------------------------------------

src/test/regress/sql/schema_based_sharding.sql

@@ -945,8 +945,21 @@ ALTER SYSTEM RESET citus.enable_schema_based_sharding;
 SELECT pg_reload_conf();
 
 \c - - - :master_port
+SET search_path TO regular_schema;
+
+CREATE TABLE type_sing(a INT);
+SELECT create_distributed_table('type_sing', NULL, colocate_with:='none');
+
+SET citus.enable_schema_based_sharding TO ON;
+CREATE SCHEMA type_sch;
+CREATE TABLE type_sch.tbl (a INT);
+
+SELECT table_name, citus_table_type FROM public.citus_tables WHERE table_name::text LIKE 'type_%';
+SELECT table_name, citus_table_type FROM citus_shards WHERE table_name::text LIKE 'type_%' AND nodeport IN (:worker_1_port, :worker_2_port);
+
+RESET citus.enable_schema_based_sharding;
 
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6 CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch CASCADE;
 
 SELECT citus_remove_node('localhost', :master_port);