add own_node to citus_shards

src/backend/distributed/sql/citus--12.1-1--12.2-1.sql

@@ -7,7 +7,8 @@
 ALTER TABLE pg_dist_shard ADD COLUMN needsisolatednode boolean NOT NULL DEFAULT false;
 
 #include "udfs/citus_internal_add_shard_metadata/12.2-1.sql"
-
 #include "udfs/citus_internal_shard_group_set_needsisolatednode/12.2-1.sql"
-
 #include "udfs/citus_shard_property_set/12.2-1.sql"
+
+DROP VIEW citus_shards;
+#include "udfs/citus_shards/12.2-1.sql"

src/backend/distributed/sql/downgrades/citus--12.2-1--12.1-1.sql

@@ -2,7 +2,8 @@
 
 #include "../udfs/citus_add_rebalance_strategy/10.1-1.sql"
 
-ALTER TABLE pg_dist_shard DROP COLUMN needsisolatednode;
+DROP VIEW pg_catalog.citus_shards;
+#include "../udfs/citus_shards/12.0-1.sql"
 
 DROP FUNCTION pg_catalog.citus_internal_add_shard_metadata(regclass, bigint, "char", text, text, boolean);
 #include "../udfs/citus_internal_add_shard_metadata/10.2-1.sql"
@@ -12,3 +13,5 @@ DROP FUNCTION pg_catalog.citus_shard_property_set(shard_id bigint, anti_affinity
 DROP FUNCTION pg_catalog.citus_internal_shard_group_set_needsisolatednode(
                             shard_id bigint,
                             enabled boolean);
+
+ALTER TABLE pg_dist_shard DROP COLUMN needsisolatednode;

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

@@ -0,0 +1,71 @@
+CREATE OR REPLACE VIEW citus.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_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,
+     CASE
+      WHEN NOT pg_dist_shard.needsisolatednode THEN false
+      ELSE
+        -- own_node = true if the node doesn't have any other shards except the ones that are colocated with this shard
+        NOT EXISTS (
+            -- get all the distributed table shards that are placed on the same node as this shard
+            SELECT pds1.shardid
+            FROM pg_dist_shard pds1
+            JOIN pg_dist_placement pdp1 USING (shardid)
+            JOIN pg_dist_partition pdp2 USING (logicalrelid)
+            WHERE pdp1.groupid = pg_dist_placement.groupid AND
+                  (pdp2.partkey IS NOT NULL OR (pdp2.repmodel != 't' AND pdp2.colocationid != 0))
+            EXCEPT
+            -- get all the shards that are colocated with this shard
+            SELECT pds1.shardid
+            FROM pg_dist_shard pds1
+            JOIN pg_dist_partition pdp1 USING (logicalrelid)
+            WHERE pdp1.colocationid = pg_dist_partition.colocationid AND
+                  ((pds1.shardminvalue IS NULL AND pg_dist_shard.shardminvalue IS NULL) OR (pds1.shardminvalue = pg_dist_shard.shardminvalue))
+        )
+     END AS own_node
+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 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
+;
+
+ALTER VIEW citus.citus_shards SET SCHEMA pg_catalog;
+GRANT SELECT ON pg_catalog.citus_shards TO public;

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

@@ -11,7 +11,28 @@ SELECT
      colocationid AS colocation_id,
      pg_dist_node.nodename,
      pg_dist_node.nodeport,
-     size as shard_size
+     size as shard_size,
+     CASE
+      WHEN NOT pg_dist_shard.needsisolatednode THEN false
+      ELSE
+        -- own_node = true if the node doesn't have any other shards except the ones that are colocated with this shard
+        NOT EXISTS (
+            -- get all the distributed table shards that are placed on the same node as this shard
+            SELECT pds1.shardid
+            FROM pg_dist_shard pds1
+            JOIN pg_dist_placement pdp1 USING (shardid)
+            JOIN pg_dist_partition pdp2 USING (logicalrelid)
+            WHERE pdp1.groupid = pg_dist_placement.groupid AND
+                  (pdp2.partkey IS NOT NULL OR (pdp2.repmodel != 't' AND pdp2.colocationid != 0))
+            EXCEPT
+            -- get all the shards that are colocated with this shard
+            SELECT pds1.shardid
+            FROM pg_dist_shard pds1
+            JOIN pg_dist_partition pdp1 USING (logicalrelid)
+            WHERE pdp1.colocationid = pg_dist_partition.colocationid AND
+                  ((pds1.shardminvalue IS NULL AND pg_dist_shard.shardminvalue IS NULL) OR (pds1.shardminvalue = pg_dist_shard.shardminvalue))
+        )
+     END AS own_node
 FROM
    pg_dist_shard
 JOIN

src/test/regress/expected/isolate_placement.out

@@ -7,6 +7,26 @@
 SET client_min_messages TO WARNING;
 CALL citus_cleanup_orphaned_resources();
 RESET client_min_messages;
+-- Returns true if all placement groups within given shard group are isolated.
+--
+-- Not created in isolate_placement schema because it's dropped a few times during the test.
+CREATE OR REPLACE FUNCTION verify_placements_in_shard_group_isolated(
+    qualified_table_name text,
+    shard_group_index bigint)
+RETURNS boolean
+AS $func$
+DECLARE
+    v_result boolean;
+  BEGIN
+    SELECT bool_and(own_node) INTO v_result
+    FROM citus_shards
+    JOIN (
+        SELECT shardids FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index
+    ) q
+    ON (shardid = ANY(q.shardids));
+    RETURN v_result;
+  END;
+$func$ LANGUAGE plpgsql;
 CREATE SCHEMA isolate_placement;
 SET search_path TO isolate_placement;
 -- test null input
@@ -1162,6 +1182,32 @@ SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_
  t
 (1 row)
 
+SELECT DISTINCT(table_name::regclass::text)
+FROM citus_shards
+JOIN pg_class ON (oid = table_name)
+WHERE relnamespace = 'isolate_placement'::regnamespace AND own_node
+ORDER BY 1;
+            table_name
+---------------------------------------------------------------------
+ dist_1
+ dist_2
+ dist_3
+ dist_post_colocated
+ dist_post_concurrently_colocated
+(5 rows)
+
+SELECT bool_or(own_node) = false
+FROM citus_shards
+JOIN (
+    SELECT unnest(shardids) shardid
+    FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
+    WHERE shardgroupindex != 1
+) shards_except_group_1 USING (shardid);
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
 DROP TABLE dist_to_be_replicated;
 SELECT citus_drain_node('localhost', :master_port, shard_transfer_mode=>'block_writes');
  citus_drain_node
@@ -1757,6 +1803,7 @@ DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_isolate_place
 DROP TABLE single_shard_1, single_shard_3;
 SET client_min_messages TO WARNING;
 DROP SCHEMA isolate_placement CASCADE;
+DROP FUNCTION public.verify_placements_in_shard_group_isolated(text, bigint);
 SELECT citus_remove_node('localhost', :master_port);
  citus_remove_node
 ---------------------------------------------------------------------

src/test/regress/expected/multi_test_helpers.out

@@ -595,33 +595,3 @@ RETURNS SETOF jsonb AS $func$
     WHERE needsisolatednodejson::text LIKE '%true%';
   END;
 $func$ LANGUAGE plpgsql;
--- Returns true if all placement groups within given shard group are isolated.
-CREATE OR REPLACE FUNCTION verify_placements_in_shard_group_isolated(
-    qualified_table_name text,
-    shard_group_index bigint)
-RETURNS boolean
-AS $func$
-DECLARE
-    v_result boolean;
-  BEGIN
-    SELECT bool_and(ok_for_nodegroup) INTO v_result FROM (
-        SELECT array_agg(shardid ORDER BY shardid) =
-               (SELECT shardids FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index)
-               AS ok_for_nodegroup -- check whether each of those nodes only contain placements of given shard group
-        FROM citus_shards
-        JOIN pg_dist_node USING (nodename, nodeport)
-        WHERE citus_table_type = 'distributed' AND -- only interested in distributed table shards on the nodes we're interested in
-              groupid IN ( -- only interested in the nodes that contain placements of given shard group
-                SELECT DISTINCT(pdn.groupid)
-                FROM citus_shards cs
-                JOIN pg_dist_node pdn USING (nodename, nodeport)
-                WHERE cs.shardid IN (
-                    SELECT unnest(shardids) FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index
-                )
-              )
-        GROUP BY groupid
-    ) q;
-
-    RETURN v_result;
-  END;
-$func$ LANGUAGE plpgsql;

src/test/regress/sql/isolate_placement.sql

@@ -9,6 +9,27 @@ SET client_min_messages TO WARNING;
 CALL citus_cleanup_orphaned_resources();
 RESET client_min_messages;
 
+-- Returns true if all placement groups within given shard group are isolated.
+--
+-- Not created in isolate_placement schema because it's dropped a few times during the test.
+CREATE OR REPLACE FUNCTION verify_placements_in_shard_group_isolated(
+    qualified_table_name text,
+    shard_group_index bigint)
+RETURNS boolean
+AS $func$
+DECLARE
+    v_result boolean;
+  BEGIN
+    SELECT bool_and(own_node) INTO v_result
+    FROM citus_shards
+    JOIN (
+        SELECT shardids FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index
+    ) q
+    ON (shardid = ANY(q.shardids));
+    RETURN v_result;
+  END;
+$func$ LANGUAGE plpgsql;
+
 CREATE SCHEMA isolate_placement;
 SET search_path TO isolate_placement;
 
@@ -634,6 +655,20 @@ SET client_min_messages TO NOTICE;
 
 SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 
+SELECT DISTINCT(table_name::regclass::text)
+FROM citus_shards
+JOIN pg_class ON (oid = table_name)
+WHERE relnamespace = 'isolate_placement'::regnamespace AND own_node
+ORDER BY 1;
+
+SELECT bool_or(own_node) = false
+FROM citus_shards
+JOIN (
+    SELECT unnest(shardids) shardid
+    FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
+    WHERE shardgroupindex != 1
+) shards_except_group_1 USING (shardid);
+
 DROP TABLE dist_to_be_replicated;
 
 SELECT citus_drain_node('localhost', :master_port, shard_transfer_mode=>'block_writes');
@@ -980,5 +1015,6 @@ DROP TABLE single_shard_1, single_shard_3;
 
 SET client_min_messages TO WARNING;
 DROP SCHEMA isolate_placement CASCADE;
+DROP FUNCTION public.verify_placements_in_shard_group_isolated(text, bigint);
 
 SELECT citus_remove_node('localhost', :master_port);

src/test/regress/sql/multi_test_helpers.sql

@@ -623,34 +623,3 @@ RETURNS SETOF jsonb AS $func$
     WHERE needsisolatednodejson::text LIKE '%true%';
   END;
 $func$ LANGUAGE plpgsql;
-
--- Returns true if all placement groups within given shard group are isolated.
-CREATE OR REPLACE FUNCTION verify_placements_in_shard_group_isolated(
-    qualified_table_name text,
-    shard_group_index bigint)
-RETURNS boolean
-AS $func$
-DECLARE
-    v_result boolean;
-  BEGIN
-    SELECT bool_and(ok_for_nodegroup) INTO v_result FROM (
-        SELECT array_agg(shardid ORDER BY shardid) =
-               (SELECT shardids FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index)
-               AS ok_for_nodegroup -- check whether each of those nodes only contain placements of given shard group
-        FROM citus_shards
-        JOIN pg_dist_node USING (nodename, nodeport)
-        WHERE citus_table_type = 'distributed' AND -- only interested in distributed table shards on the nodes we're interested in
-              groupid IN ( -- only interested in the nodes that contain placements of given shard group
-                SELECT DISTINCT(pdn.groupid)
-                FROM citus_shards cs
-                JOIN pg_dist_node pdn USING (nodename, nodeport)
-                WHERE cs.shardid IN (
-                    SELECT unnest(shardids) FROM public.get_enumerated_shard_groups(qualified_table_name) WHERE shardgroupindex = shard_group_index
-                )
-              )
-        GROUP BY groupid
-    ) q;
-
-    RETURN v_result;
-  END;
-$func$ LANGUAGE plpgsql;