-- Due to a race condition that happens in TransferShards() when the same shard id
-- is used to create the same shard on a different worker node, need to call
-- citus_cleanup_orphaned_resources() to clean up any orphaned resources before
-- running the tests.
--
-- See https://github.com/citusdata/citus/pull/7180#issuecomment-1706786615.
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
SELECT citus_internal_shard_group_set_needsisolatednode(0, NULL);
ERROR:  enabled cannot be NULL
SELECT citus_internal_shard_group_set_needsisolatednode(NULL, false);
ERROR:  shard_id cannot be NULL
SET citus.shard_replication_factor TO 1;
SET citus.next_shard_id TO 2000000;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- test with user that doesn't have permission to execute the function
SELECT citus_internal_shard_group_set_needsisolatednode(shardid, true) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
ERROR:  This is an internal Citus function can only be used in a distributed transaction
DROP TABLE single_shard_1;
CREATE ROLE test_user_isolate_placement WITH LOGIN;
GRANT ALL ON SCHEMA isolate_placement TO test_user_isolate_placement;
ALTER SYSTEM SET citus.enable_manual_metadata_changes_for_user TO 'test_user_isolate_placement';
SELECT pg_reload_conf();
 pg_reload_conf
---------------------------------------------------------------------
 t
(1 row)

SELECT pg_sleep(0.1);
 pg_sleep
---------------------------------------------------------------------

(1 row)

SET ROLE test_user_isolate_placement;
-- test invalid shard id
SELECT citus_internal_shard_group_set_needsisolatednode(0, true);
ERROR:  could not find valid entry for shard xxxxx
-- test null needs_isolated_node
SELECT citus_internal_add_shard_metadata(
    relation_id=>0,
    shard_id=>0,
    storage_type=>'0',
    shard_min_value=>'0',
    shard_max_value=>'0',
    needs_isolated_node=>null);
ERROR:  needs isolated node cannot be NULL
RESET ROLE;
REVOKE ALL ON SCHEMA isolate_placement FROM test_user_isolate_placement;
DROP USER test_user_isolate_placement;
ALTER SYSTEM RESET citus.enable_manual_metadata_changes_for_user;
SELECT pg_reload_conf();
 pg_reload_conf
---------------------------------------------------------------------
 t
(1 row)

SELECT pg_sleep(0.1);
 pg_sleep
---------------------------------------------------------------------

(1 row)

SET search_path TO isolate_placement;
SET citus.shard_replication_factor TO 1;
SET citus.next_shard_id TO 2001000;
CREATE USER mysuperuser superuser;
SET ROLE mysuperuser;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE USER regularuser;
GRANT USAGE ON SCHEMA isolate_placement TO regularuser;
ALTER SYSTEM SET citus.enable_manual_metadata_changes_for_user TO 'regularuser';
SELECT pg_reload_conf();
 pg_reload_conf
---------------------------------------------------------------------
 t
(1 row)

SELECT pg_sleep(0.1);
 pg_sleep
---------------------------------------------------------------------

(1 row)

SET ROLE regularuser;
-- throws an error as the user is not the owner of the table
SELECT citus_shard_property_set(shardid) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
ERROR:  must be owner of table single_shard_1
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
ERROR:  must be owner of table single_shard_1
SELECT citus_internal_shard_group_set_needsisolatednode(shardid, true) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
ERROR:  must be owner of table single_shard_1
-- assign all tables to regularuser
RESET ROLE;
SELECT result FROM run_command_on_all_nodes($$ REASSIGN OWNED BY mysuperuser TO regularuser; $$);
     result
---------------------------------------------------------------------
 REASSIGN OWNED
 REASSIGN OWNED
 REASSIGN OWNED
(3 rows)

SET ROLE regularuser;
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.single_shard_1')
$$)
ORDER BY result;
                        result
---------------------------------------------------------------------
 [{"1": [{"isolate_placement.single_shard_1": true}]}]
 [{"1": [{"isolate_placement.single_shard_1": true}]}]
 [{"1": [{"isolate_placement.single_shard_1": true}]}]
(3 rows)

SELECT citus_internal_shard_group_set_needsisolatednode(shardid, false) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
 citus_internal_shard_group_set_needsisolatednode
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.single_shard_1')
$$)
ORDER BY result;
                        result
---------------------------------------------------------------------
 [{"1": [{"isolate_placement.single_shard_1": true}]}]
 [{"1": [{"isolate_placement.single_shard_1": true}]}]
 {}
(3 rows)

SELECT citus_internal_shard_group_set_needsisolatednode(shardid, true) FROM pg_dist_shard WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
 citus_internal_shard_group_set_needsisolatednode
---------------------------------------------------------------------

(1 row)

DROP TABLE single_shard_1;
RESET ROLE;
REVOKE USAGE ON SCHEMA isolate_placement FROM regularuser;
ALTER SYSTEM RESET citus.enable_manual_metadata_changes_for_user;
SELECT pg_reload_conf();
 pg_reload_conf
---------------------------------------------------------------------
 t
(1 row)

SELECT pg_sleep(0.1);
 pg_sleep
---------------------------------------------------------------------

(1 row)

DROP ROLE regularuser, mysuperuser;
SET search_path TO isolate_placement;
SET citus.next_shard_id TO 2002000;
SET citus.shard_count TO 32;
SET citus.shard_replication_factor TO 1;
SET client_min_messages TO WARNING;
SELECT 1 FROM citus_add_node('localhost', :master_port, groupid => 0);
 ?column?
---------------------------------------------------------------------
        1
(1 row)

SET client_min_messages TO NOTICE;
SET citus.shard_replication_factor TO 2;
CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
CREATE TABLE dist_3(a int);
SELECT create_distributed_table('dist_1', 'a');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_3', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SET citus.shard_replication_factor TO 1;
-- none of the placements have been marked as needsisolatednode yet
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
 result
---------------------------------------------------------------------
 {}
 {}
 {}
(3 rows)

SELECT shardids[2] AS shardgroup_5_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 5 \gset
SELECT citus_shard_property_set(:shardgroup_5_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT shardids[3] AS shardgroup_10_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 10 \gset
SELECT citus_shard_property_set(:shardgroup_10_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                                                                                   result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

SELECT shardids[1] AS shardgroup_3_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 3 \gset
SELECT citus_shard_property_set(:shardgroup_3_shardid, anti_affinity=>false);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                                                                                   result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}, {"10": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

SELECT shardids[1] AS shardgroup_10_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 10 \gset
SELECT citus_shard_property_set(:shardgroup_10_shardid, anti_affinity=>false);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                        result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

SELECT shardids[1] AS shardgroup_5_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 5 \gset
SELECT citus_shard_property_set(:shardgroup_5_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                        result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

-- test metadata sync
-- first, need to re-create them with shard_replication_factor = 1 because we will first remove worker_2
DROP TABLE dist_1, dist_2, dist_3;
SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);
 ?column?
---------------------------------------------------------------------
        1
(1 row)

CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
CREATE TABLE dist_3(a int);
SELECT create_distributed_table('dist_1', 'a');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_3', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
 ?column?
---------------------------------------------------------------------
        1
(1 row)

SELECT shardids[1] AS shardgroup_5_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 5 \gset
SELECT citus_shard_property_set(:shardgroup_5_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                        result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

CREATE TABLE dist_4(a int);
SELECT create_distributed_table('dist_4', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_4_concurrently(a int);
SELECT create_distributed_table_concurrently('dist_4_concurrently', 'a', colocate_with=>'dist_1');
NOTICE:  relation dist_4_concurrently does not have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the relation will error out during create_distributed_table_concurrently unless there is a REPLICA IDENTITY or PRIMARY KEY. INSERT commands will still work.
 create_distributed_table_concurrently
---------------------------------------------------------------------

(1 row)

-- Placements of a new distributed table created within the same colocated
-- group inherit needsisolatednode from the colocated placements too.
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                                                                   result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}, {"isolate_placement.dist_4": true}, {"isolate_placement.dist_4_concurrently": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}, {"isolate_placement.dist_4": true}, {"isolate_placement.dist_4_concurrently": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}, {"isolate_placement.dist_4": true}, {"isolate_placement.dist_4_concurrently": true}]}]
(3 rows)

DROP TABLE dist_4, dist_4_concurrently;
-- Returns source and target node ids that can be used to perform a
-- shard transfer for one of the placements of given shard.
CREATE OR REPLACE FUNCTION get_candidate_node_for_shard_transfer(
    p_shardid bigint)
RETURNS TABLE (source_nodeid integer, target_nodeid integer)
SET search_path TO 'pg_catalog, public'
AS $func$
DECLARE
    v_source_nodeids integer[];
    v_target_nodeid integer;
BEGIN
    SELECT array_agg(nodeid) INTO v_source_nodeids
    FROM pg_dist_shard
    JOIN pg_dist_placement USING (shardid)
    JOIN pg_dist_node USING (groupid)
    WHERE noderole = 'primary' AND shardid = p_shardid;

    IF v_source_nodeids IS NULL
    THEN
        RAISE EXCEPTION 'could not determine the source node of shard %', p_shardid;
    END IF;

    SELECT nodeid INTO v_target_nodeid
    FROM pg_dist_node
    WHERE isactive AND shouldhaveshards AND noderole='primary' AND
          nodeid NOT IN (SELECT unnest(v_source_nodeids))
    LIMIT 1;

    IF v_target_nodeid IS NULL
    THEN
        RAISE EXCEPTION 'could not determine a node to transfer the placement to';
    END IF;

    RETURN QUERY SELECT v_source_nodeids[1], v_target_nodeid;
END;
$func$ LANGUAGE plpgsql;
SELECT shardids[1] AS shardgroup_15_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 15 \gset
SELECT citus_move_shard_placement(:shardgroup_5_shardid, source_nodeid, target_nodeid, 'block_writes')
FROM get_candidate_node_for_shard_transfer(:shardgroup_5_shardid);
 citus_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT citus_move_shard_placement(:shardgroup_15_shardid, source_nodeid, target_nodeid, 'block_writes')
FROM get_candidate_node_for_shard_transfer(:shardgroup_15_shardid);
 citus_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- so that citus_copy_shard_placement works
UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid = 'isolate_placement.dist_1'::regclass;
SELECT citus_copy_shard_placement(:shardgroup_5_shardid, source_nodeid, target_nodeid, 'block_writes')
FROM get_candidate_node_for_shard_transfer(:shardgroup_5_shardid);
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT citus_copy_shard_placement(:shardgroup_15_shardid, source_nodeid, target_nodeid, 'block_writes')
FROM get_candidate_node_for_shard_transfer(:shardgroup_15_shardid);
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                        result
---------------------------------------------------------------------
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
 [{"5": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_3": true}]}]
(3 rows)

DROP TABLE dist_1, dist_2, dist_3;
CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
SELECT create_distributed_table('dist_1', 'a', shard_count=>3);
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT shardids[1] AS shardgroup_3_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 3 \gset
SELECT citus_shard_property_set(:shardgroup_3_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                      result
---------------------------------------------------------------------
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
(3 rows)

-- so that replicate_table_shards works
UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid = 'isolate_placement.dist_1'::regclass;
SET client_min_messages TO WARNING;
SELECT replicate_table_shards('isolate_placement.dist_1', shard_replication_factor=>2, shard_transfer_mode=>'block_writes');
 replicate_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                      result
---------------------------------------------------------------------
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"3": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
(3 rows)

DROP TABLE dist_1, dist_2;
CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
SELECT create_distributed_table('dist_1', 'a');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT shardids[1] AS shardgroup_9_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 9 \gset
SELECT citus_shard_property_set(:shardgroup_9_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                      result
---------------------------------------------------------------------
 [{"9": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"9": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"9": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
(3 rows)

SELECT nodeid AS worker_1_node FROM pg_dist_node WHERE nodeport=:worker_1_port \gset
SELECT nodeid AS worker_2_node FROM pg_dist_node WHERE nodeport=:worker_2_port \gset
SELECT pg_catalog.citus_split_shard_by_split_points(
    :shardgroup_9_shardid,
    ARRAY[((shardminvalue::bigint + shardmaxvalue::bigint) / 2)::text],
    ARRAY[:worker_1_node, :worker_2_node],
    'block_writes')
FROM pg_dist_shard
WHERE shardid = :shardgroup_9_shardid;
 citus_split_shard_by_split_points
---------------------------------------------------------------------

(1 row)

-- We shouldn't see shard group 9 because shard-split operation doesn't
-- preserve needsisolatednode flag when splitting the shard.
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
 result
---------------------------------------------------------------------
 {}
 {}
 {}
(3 rows)

SELECT shardids[1] AS shardgroup_12_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 12 \gset
SELECT citus_shard_property_set(:shardgroup_12_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                       result
---------------------------------------------------------------------
 [{"12": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"12": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"12": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
(3 rows)

SELECT shardids[1] AS shardgroup_10_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 10 \gset
SELECT pg_catalog.citus_split_shard_by_split_points(
    :shardgroup_10_shardid,
    ARRAY[((shardminvalue::bigint + shardmaxvalue::bigint) / 2)::text],
    ARRAY[:worker_1_node, :worker_2_node],
    'block_writes')
FROM pg_dist_shard
WHERE shardid = :shardgroup_10_shardid;
 citus_split_shard_by_split_points
---------------------------------------------------------------------

(1 row)

-- We should see old shard group 12 (now as 13 due to split
-- of a prior shard) because it's not the one we splitted.
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                       result
---------------------------------------------------------------------
 [{"13": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"13": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
 [{"13": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}]}]
(3 rows)

CREATE TABLE dist_3(a int);
SELECT create_distributed_table('dist_3', 'a', colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT shardids[1] AS shardgroup_17_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_3')
WHERE shardgroupindex = 17 \gset
SELECT citus_shard_property_set(:shardgroup_17_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_3')
$$)
ORDER BY result;
                     result
---------------------------------------------------------------------
 [{"17": [{"isolate_placement.dist_3": true}]}]
 [{"17": [{"isolate_placement.dist_3": true}]}]
 [{"17": [{"isolate_placement.dist_3": true}]}]
(3 rows)

-- verify that shard key value 100 is stored on shard group 17
select get_shard_id_for_distribution_column('dist_3', 100) = :shardgroup_17_shardid;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

SELECT 1 FROM isolate_tenant_to_new_shard('dist_3', 100, shard_transfer_mode => 'block_writes');
 ?column?
---------------------------------------------------------------------
        1
(1 row)

-- We shouldn't see shard group 17 because isolate_tenant_to_new_shard doesn't
-- preserve needsisolatednode flag when splitting the shard.
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_3')
$$)
ORDER BY result;
 result
---------------------------------------------------------------------
 {}
 {}
 {}
(3 rows)

SELECT shardids[1] AS shardgroup_18_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_3')
WHERE shardgroupindex = 18 \gset
SELECT citus_shard_property_set(:shardgroup_18_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_3')
$$)
ORDER BY result;
                     result
---------------------------------------------------------------------
 [{"18": [{"isolate_placement.dist_3": true}]}]
 [{"18": [{"isolate_placement.dist_3": true}]}]
 [{"18": [{"isolate_placement.dist_3": true}]}]
(3 rows)

-- verify that shard key value 1000 is _not_ stored on shard group 18
SELECT get_shard_id_for_distribution_column('dist_3', 1000) != :shardgroup_18_shardid;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

SELECT 1 FROM isolate_tenant_to_new_shard('dist_3', 1000, shard_transfer_mode => 'block_writes');
 ?column?
---------------------------------------------------------------------
        1
(1 row)

-- We should see shard group 18 (now as 20 due to split of a prior shard)
-- because it's not the one we splitted.
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_3')
$$)
ORDER BY result;
                     result
---------------------------------------------------------------------
 [{"20": [{"isolate_placement.dist_3": true}]}]
 [{"20": [{"isolate_placement.dist_3": true}]}]
 [{"20": [{"isolate_placement.dist_3": true}]}]
(3 rows)

CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT shardids[1] AS shardgroup_1_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.single_shard_1')
WHERE shardgroupindex = 1 \gset
SELECT citus_shard_property_set(:shardgroup_1_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

-- noop
SELECT citus_shard_property_set(:shardgroup_1_shardid, NULL);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(:shardgroup_1_shardid);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'single_shard_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.single_shard_1')
$$)
ORDER BY result;
                                              result
---------------------------------------------------------------------
 [{"1": [{"isolate_placement.single_shard_1": true}, {"isolate_placement.single_shard_2": true}]}]
 [{"1": [{"isolate_placement.single_shard_1": true}, {"isolate_placement.single_shard_2": true}]}]
 [{"1": [{"isolate_placement.single_shard_1": true}, {"isolate_placement.single_shard_2": true}]}]
(3 rows)

-- test invalid input
SELECT citus_shard_property_set(NULL, anti_affinity=>true);
ERROR:  shard_id cannot be NULL
SELECT citus_shard_property_set(0, anti_affinity=>true);
ERROR:  shard xxxxx does not exist
SELECT citus_shard_property_set(NULL, anti_affinity=>false);
ERROR:  shard_id cannot be NULL
SELECT citus_shard_property_set(0, anti_affinity=>false);
ERROR:  shard xxxxx does not exist
-- we verify whether shard exists even if anti_affinity is not provided
SELECT citus_shard_property_set(0, anti_affinity=>NULL);
ERROR:  shard xxxxx does not exist
CREATE TABLE append_table (a int, b int);
SELECT create_distributed_table('append_table', 'a', 'append');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT 1 FROM master_create_empty_shard('append_table');
 ?column?
---------------------------------------------------------------------
        1
(1 row)

CREATE TYPE composite_key_type AS (f1 int, f2 text);
CREATE TABLE range_table(key composite_key_type, value int);
SELECT create_distributed_table('range_table', 'key', 'range');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CALL public.create_range_partitioned_shards('range_table', '{"(0,a)","(25,a)"}','{"(24,z)","(49,z)"}');
CREATE TABLE ref_table(a int);
SELECT create_reference_table('ref_table');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE local_table(a int);
SELECT citus_add_local_table_to_metadata('local_table');
 citus_add_local_table_to_metadata
---------------------------------------------------------------------

(1 row)

-- all should fail
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'append_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'range_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'ref_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid = 'local_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>false) FROM pg_dist_shard WHERE logicalrelid = 'append_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>false) FROM pg_dist_shard WHERE logicalrelid = 'range_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>false) FROM pg_dist_shard WHERE logicalrelid = 'ref_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
SELECT citus_shard_property_set(shardid, anti_affinity=>false) FROM pg_dist_shard WHERE logicalrelid = 'local_table'::regclass LIMIT 1;
ERROR:  shard isolation is only supported for hash distributed tables
DROP TABLE range_table;
DROP TYPE composite_key_type;
SET client_min_messages TO WARNING;
DROP SCHEMA isolate_placement CASCADE;
CREATE SCHEMA isolate_placement;
SET search_path TO isolate_placement;
SET client_min_messages TO NOTICE;
CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
SELECT create_distributed_table('dist_1', 'a', shard_count=>4);
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_non_colocated(a int);
SELECT create_distributed_table('dist_non_colocated', 'a', shard_count=>4, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE append_table (a int, b int);
SELECT create_distributed_table('append_table', 'a', 'append');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT 1 FROM master_create_empty_shard('append_table');
 ?column?
---------------------------------------------------------------------
        1
(1 row)

CREATE TABLE range_table(a int, b int);
SELECT create_distributed_table('range_table', 'a', 'range');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CALL public.create_range_partitioned_shards('range_table', '{"0","25"}','{"26","50"}');
CREATE TABLE reference_table_1(a int);
SELECT create_reference_table('reference_table_1');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE local_table_1(a int);
SELECT citus_add_local_table_to_metadata('local_table_1');
 citus_add_local_table_to_metadata
---------------------------------------------------------------------

(1 row)

SELECT shardids[1] AS shardgroup_1_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_1')
WHERE shardgroupindex = 1 \gset
SELECT citus_shard_property_set(:shardgroup_1_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
CREATE TABLE reference_table_2(a int);
SELECT create_reference_table('reference_table_2');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE local_table_2(a int);
SELECT citus_add_local_table_to_metadata('local_table_2');
 citus_add_local_table_to_metadata
---------------------------------------------------------------------

(1 row)

-- make sure that we still have placements for both reference tables on all nodes
SELECT COUNT(DISTINCT(groupid))=3 FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'reference_table_1'::regclass;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

SELECT COUNT(DISTINCT(groupid))=3 FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'reference_table_2'::regclass;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

-- sanity check for local tables
SELECT groupid = 0 FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'local_table_1'::regclass;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

SELECT groupid = 0 FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'local_table_2'::regclass;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

CREATE TABLE dist_post_non_colocated(a int);
SELECT create_distributed_table('dist_post_non_colocated', 'a', shard_count=>4, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_post_concurrently_non_colocated(a int);
SELECT create_distributed_table_concurrently('dist_post_concurrently_non_colocated', 'a', shard_count=>4, colocate_with=>'none');
NOTICE:  relation dist_post_concurrently_non_colocated does not have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the relation will error out during create_distributed_table_concurrently unless there is a REPLICA IDENTITY or PRIMARY KEY. INSERT commands will still work.
 create_distributed_table_concurrently
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_post_colocated(a int);
SELECT create_distributed_table('dist_post_colocated', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_post_concurrently_colocated(a int);
SELECT create_distributed_table_concurrently('dist_post_concurrently_colocated', 'a', colocate_with=>'dist_1');
NOTICE:  relation dist_post_concurrently_colocated does not have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the relation will error out during create_distributed_table_concurrently unless there is a REPLICA IDENTITY or PRIMARY KEY. INSERT commands will still work.
 create_distributed_table_concurrently
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_post(a int);
SELECT create_distributed_table('single_shard_post', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE append_table_post(a int, b int);
SELECT create_distributed_table('append_table_post', 'a', 'append');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT 1 FROM master_create_empty_shard('append_table_post');
 ?column?
---------------------------------------------------------------------
        1
(1 row)

CREATE TABLE range_table_post(a int, b int);
SELECT create_distributed_table('range_table_post', 'a', 'range');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CALL public.create_range_partitioned_shards('range_table_post', '{"0","25"}','{"26","50"}');
SELECT result FROM run_command_on_all_nodes($$
    SELECT * FROM public.get_colocated_shards_needisolatednode('isolate_placement.dist_1')
$$)
ORDER BY result;
                                                                                              result
---------------------------------------------------------------------
 [{"1": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_post_colocated": true}, {"isolate_placement.dist_post_concurrently_colocated": true}]}]
 [{"1": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_post_colocated": true}, {"isolate_placement.dist_post_concurrently_colocated": true}]}]
 [{"1": [{"isolate_placement.dist_1": true}, {"isolate_placement.dist_2": true}, {"isolate_placement.dist_post_colocated": true}, {"isolate_placement.dist_post_concurrently_colocated": true}]}]
(3 rows)

-- Make sure that the node that contains shard-group 1 of isolate_placement.dist_1
-- doesn't have any other placements.
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SET client_min_messages TO ERROR;
SELECT citus_drain_node('localhost', :worker_1_port, shard_transfer_mode=>'block_writes');
 citus_drain_node
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT citus_set_node_property('localhost', :worker_1_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

-- drain node should have failed and the node should still have the same set of placements
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SET client_min_messages TO ERROR;
SELECT citus_drain_node('localhost', :worker_2_port, shard_transfer_mode=>'block_writes');
 citus_drain_node
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT citus_set_node_property('localhost', :worker_2_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

-- drain node should have failed and the node should still have the same set of placements
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

CREATE TABLE dist_3(a int);
SELECT create_distributed_table('dist_3', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SET citus.shard_replication_factor TO 2;
CREATE TABLE dist_replicated(a int);
-- fails as we only have one node that's not used to isolate a shard placement group
SELECT create_distributed_table('dist_replicated', 'a', shard_count=>4, colocate_with=>'none');
ERROR:  replication_factor (2) exceeds number of available worker nodes (1)
HINT:  Add more worker nodes or try again with a lower replication factor.
SET citus.shard_replication_factor TO 1;
CREATE TABLE dist_to_be_replicated(a int);
SELECT create_distributed_table('dist_to_be_replicated', 'a', shard_count=>4, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid = 'isolate_placement.dist_to_be_replicated'::regclass;
SET client_min_messages TO WARNING;
-- fails as we only have one node that's not used to isolate a shard placement group
SELECT replicate_table_shards('isolate_placement.dist_to_be_replicated', shard_replication_factor=>2, shard_transfer_mode=>'block_writes');
ERROR:  could not find a target for shard xxxxx
SET client_min_messages TO NOTICE;
SELECT citus_set_node_property('localhost', :master_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
-- succeeds as now we have two nodes that are not used to isolate a shard placement group
SELECT replicate_table_shards('isolate_placement.dist_to_be_replicated', shard_replication_factor=>2, shard_transfer_mode=>'block_writes');
 replicate_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 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
---------------------------------------------------------------------

(1 row)

DROP TABLE dist_replicated;
SET client_min_messages TO WARNING;
DROP SCHEMA isolate_placement CASCADE;
CREATE SCHEMA isolate_placement;
SET search_path TO isolate_placement;
SET client_min_messages TO NOTICE;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass, 'single_shard_2'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------


(2 rows)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
-- fails
CREATE TABLE dist_1(a int);
SELECT create_distributed_table('dist_1', 'a', shard_count=>4);
ERROR:  replication_factor (1) exceeds number of available worker nodes (0)
HINT:  Add more worker nodes or try again with a lower replication factor.
CREATE TABLE single_shard_3(a int);
SELECT create_distributed_table('single_shard_3', null, colocate_with=>'none');
ERROR:  couldn't find any available worker nodes
HINT:  Add more worker nodes
CREATE TABLE append_table (a int, b int);
SELECT create_distributed_table('append_table', 'a', 'append');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT 1 FROM master_create_empty_shard('append_table');
ERROR:  no worker nodes are available for placing shards
HINT:  Add more worker nodes.
CREATE TABLE range_table(a int, b int);
SELECT create_distributed_table('range_table', 'a', 'range');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CALL public.create_range_partitioned_shards('range_table', '{"0","25"}','{"26","50"}');
ERROR:  no worker nodes are available for placing shards
HINT:  Add more worker nodes.
CONTEXT:  SQL statement "SELECT master_create_empty_shard(rel::text)"
PL/pgSQL function public.create_range_partitioned_shards(regclass,text[],text[]) line XX at SQL statement
-- succeeds
CREATE TABLE reference_table_1(a int);
SELECT create_reference_table('reference_table_1');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE local_table_1(a int);
SELECT citus_add_local_table_to_metadata('local_table_1');
 citus_add_local_table_to_metadata
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_4(a int);
SELECT create_distributed_table('single_shard_4', null, colocate_with=>'single_shard_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_2', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SET client_min_messages TO WARNING;
DROP SCHEMA isolate_placement CASCADE;
CREATE SCHEMA isolate_placement;
SET search_path TO isolate_placement;
SET client_min_messages TO NOTICE;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- Make sure that we don't assume that a node is used to isolate a shard placement
-- group just because it contains a single shard placement group.
CREATE TABLE single_shard_3(a int);
SELECT create_distributed_table('single_shard_3', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
DROP SCHEMA isolate_placement CASCADE;
CREATE SCHEMA isolate_placement;
SET search_path TO isolate_placement;
SET client_min_messages TO NOTICE;
SELECT citus_set_node_property('localhost', :master_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int); -- will replicate this
CREATE TABLE dist_3(a int);
SELECT create_distributed_table('dist_1', 'a', shard_count=>1);
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', shard_count=>1, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_3', 'a', shard_count=>1, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('dist_1'::regclass, 'dist_2'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------


(2 rows)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
-- so that replicate_table_shards works
UPDATE pg_dist_partition SET repmodel = 'c' WHERE logicalrelid = 'isolate_placement.dist_2'::regclass;
SET client_min_messages TO WARNING;
-- succeeds but breaks the isolation requirement for either of dist_1 or dist_2 ..
SELECT replicate_table_shards('isolate_placement.dist_2', shard_replication_factor=>2, shard_transfer_mode=>'block_writes');
 replicate_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
-- .. so check the xor of the isolation requirements for dist_1 and dist_2
SELECT (public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1) OR public.verify_placements_in_shard_group_isolated('isolate_placement.dist_2', 1)) = true AND
       (public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1) AND public.verify_placements_in_shard_group_isolated('isolate_placement.dist_2', 1)) = false;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE dist_1, dist_2, dist_3;
SELECT citus_drain_node('localhost', :master_port, shard_transfer_mode=>'block_writes');
 citus_drain_node
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE dist_1(a int);
SELECT create_distributed_table('dist_1', 'a', shard_count=>4);
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT groupid AS single_shard_1_group_id FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass \gset
SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

-- show that we try to isolate placements where they were staying at the time rebalancer is invoked
SELECT groupid = :single_shard_1_group_id FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) WHERE logicalrelid = 'isolate_placement.single_shard_1'::regclass;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE dist_1, single_shard_1, single_shard_2;
SET citus.shard_replication_factor TO 2;
CREATE TABLE dist_1(a int);
CREATE TABLE dist_2(a int);
SELECT citus_set_node_property('localhost', :master_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_1', 'a', shard_count=>1);
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT create_distributed_table('dist_2', 'a', colocate_with=>'dist_1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SET citus.shard_replication_factor TO 1;
SELECT shardids[1] AS shardgroup_1_shardid
FROM public.get_enumerated_shard_groups('isolate_placement.dist_2')
WHERE shardgroupindex = 1 \gset
SELECT citus_shard_property_set(:shardgroup_1_shardid, anti_affinity=>true);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.dist_1', 1) = true;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE dist_1, dist_2;
SELECT citus_drain_node('localhost', :master_port, shard_transfer_mode=>'block_writes');
 citus_drain_node
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_3(a int);
SELECT create_distributed_table('single_shard_3', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_4(a int);
SELECT create_distributed_table('single_shard_4', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_2'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>false) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO WARNING;
SELECT rebalance_table_shards(shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_2', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1) = false;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE single_shard_1, single_shard_2, single_shard_3, single_shard_4;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- this would be placed on the same node as single_shard_1
CREATE TABLE single_shard_3(a int);
SELECT create_distributed_table('single_shard_3', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

DROP TABLE single_shard_2;
SELECT shardid, nodeid INTO single_shard_3_shardid_nodeid
FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) JOIN pg_dist_node USING (groupid)
WHERE logicalrelid = 'isolate_placement.single_shard_3'::regclass AND noderole = 'primary';
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_3'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

-- tell rebalancer that single_shard_3 cannot be placed on the node where it is currently placed
CREATE OR REPLACE FUNCTION test_shard_allowed_on_node(p_shardid bigint, p_nodeid int)
    RETURNS boolean AS
$$
    SELECT
        CASE
            WHEN (p_shardid = shardid and p_nodeid = nodeid) THEN false
            ELSE true
        END
    FROM single_shard_3_shardid_nodeid;
$$ LANGUAGE sql;
INSERT INTO pg_catalog.pg_dist_rebalance_strategy(
    name,
    default_strategy,
    shard_cost_function,
    node_capacity_function,
    shard_allowed_on_node_function,
    default_threshold,
    minimum_threshold,
    improvement_threshold
)
VALUES (
    'test_isolate_placement',
    false,
    'citus_shard_cost_1',
    'citus_node_capacity_1',
    'isolate_placement.test_shard_allowed_on_node',
    0,
    0,
    0
);
SET client_min_messages TO ERROR;
SELECT rebalance_table_shards(rebalance_strategy := 'test_isolate_placement', shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
-- test_shard_allowed_on_node() didn't cause rebalance_table_shards() to fail.
--
-- Right now single_shard_1 & single_shard_3 are placed on the same node. And
-- due to order we follow when assigning nodes to placement groups that need an
-- isolated node, we will try placing single_shard_1 to the node where it is
-- currently placed, and then we will try placing single_shard_3 to some other
-- node (as its current node is already assigned to single_shard_1), not to the
-- one we disallowed in test_shard_allowed_on_node().
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_3', 1);
 verify_placements_in_shard_group_isolated
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE single_shard_3_shardid_nodeid;
DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_isolate_placement';
DROP TABLE single_shard_1, single_shard_3;
CREATE TABLE single_shard_1(a int);
SELECT create_distributed_table('single_shard_1', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE single_shard_2(a int);
SELECT create_distributed_table('single_shard_2', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- this would be placed on the same node as single_shard_1
CREATE TABLE single_shard_3(a int);
SELECT create_distributed_table('single_shard_3', null, colocate_with=>'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

DROP TABLE single_shard_2;
SELECT shardid, nodeid INTO single_shard_3_shardid_nodeid
FROM pg_dist_shard JOIN pg_dist_placement USING (shardid) JOIN pg_dist_node USING (groupid)
WHERE logicalrelid = 'isolate_placement.single_shard_3'::regclass AND noderole = 'primary';
SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_1'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

SELECT citus_shard_property_set(shardid, anti_affinity=>true) FROM pg_dist_shard WHERE logicalrelid IN ('single_shard_3'::regclass);
 citus_shard_property_set
---------------------------------------------------------------------

(1 row)

-- Same test above but this time we tell rebalancer that single_shard_3 cannot be placed
-- on any node except the one where it is currently placed.
CREATE OR REPLACE FUNCTION test_shard_allowed_on_node(p_shardid bigint, p_nodeid int)
    RETURNS boolean AS
$$
    SELECT
        CASE
            WHEN (p_shardid = shardid and p_nodeid != nodeid) THEN false
            ELSE true
        END
    FROM single_shard_3_shardid_nodeid;
$$ LANGUAGE sql;
INSERT INTO pg_catalog.pg_dist_rebalance_strategy(
    name,
    default_strategy,
    shard_cost_function,
    node_capacity_function,
    shard_allowed_on_node_function,
    default_threshold,
    minimum_threshold,
    improvement_threshold
)
VALUES (
    'test_isolate_placement',
    false,
    'citus_shard_cost_1',
    'citus_node_capacity_1',
    'isolate_placement.test_shard_allowed_on_node',
    0,
    0,
    0
);
SET client_min_messages TO ERROR;
SELECT rebalance_table_shards(rebalance_strategy := 'test_isolate_placement', shard_transfer_mode=>'block_writes');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO NOTICE;
-- This time, test_shard_allowed_on_node() caused rebalance_table_shards() to
-- fail.
--
-- Right now single_shard_1 & single_shard_3 are placed on the same node. And
-- due to order we follow when assigning nodes to placement groups that need an
-- isolated node, we will try placing single_shard_1 to the node where it is
-- currently placed, and then we will try placing single_shard_3 to some other
-- node (as its current node is already assigned to single_shard_1). However,
-- test_shard_allowed_on_node() doesn't allow that.
SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_1', 1) = false;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

SELECT public.verify_placements_in_shard_group_isolated('isolate_placement.single_shard_3', 1) = false;
 ?column?
---------------------------------------------------------------------
 t
(1 row)

DROP TABLE single_shard_3_shardid_nodeid;
DELETE FROM pg_catalog.pg_dist_rebalance_strategy WHERE name='test_isolate_placement';
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
---------------------------------------------------------------------

(1 row)