CREATE SCHEMA citus_drain_node;
SET search_path TO citus_drain_node;
SET citus.shard_count TO 4;
SET citus.shard_replication_factor TO 1;
SET citus.next_shard_id TO 974653;
SET client_min_messages TO ERROR;
SELECT * FROM citus_set_coordinator_host('localhost', :master_port);
 citus_set_coordinator_host
---------------------------------------------------------------------

(1 row)

SELECT * FROM master_set_node_property('localhost', :master_port, 'shouldhaveshards', true);
 master_set_node_property
---------------------------------------------------------------------

(1 row)

CREATE TABLE test (x INT, y INT);
SELECT create_distributed_table('test','x');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CALL citus_cleanup_orphaned_resources();
SELECT nodename, nodeport, COUNT(*)
  FROM pg_dist_placement AS placement,
       pg_dist_node AS node
 WHERE placement.groupid = node.groupid
   AND node.noderole = 'primary' GROUP BY nodename, nodeport ORDER BY 1,2;
 nodename  | nodeport | count
---------------------------------------------------------------------
 localhost |    57636 |     2
 localhost |    57637 |     1
 localhost |    57638 |     1
(3 rows)

SELECT * FROM citus_set_node_property('localhost', :worker_2_port, 'shouldhaveshards', false);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SELECT * from citus_drain_node('localhost', :worker_1_port, shard_transfer_mode :='force_logical');
 citus_drain_node
---------------------------------------------------------------------

(1 row)

CALL citus_cleanup_orphaned_resources();
SELECT nodename, nodeport, COUNT(*)
  FROM pg_dist_placement AS placement,
       pg_dist_node AS node
 WHERE placement.groupid = node.groupid
   AND node.noderole = 'primary' GROUP BY nodename, nodeport ORDER BY 1,2;
 nodename  | nodeport | count
---------------------------------------------------------------------
 localhost |    57636 |     3
 localhost |    57638 |     1
(2 rows)

SELECT * FROM citus_set_node_property('localhost', :worker_1_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SELECT * FROM citus_set_node_property('localhost', :worker_2_port, 'shouldhaveshards', true);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SELECT * FROM rebalance_table_shards(shard_transfer_mode :='force_logical');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

CALL citus_cleanup_orphaned_resources();
SELECT nodename, nodeport, COUNT(*)
  FROM pg_dist_placement AS placement,
       pg_dist_node AS node
 WHERE placement.groupid = node.groupid
   AND node.noderole = 'primary' GROUP BY nodename, nodeport ORDER BY 1,2;
 nodename  | nodeport | count
---------------------------------------------------------------------
 localhost |    57636 |     2
 localhost |    57637 |     1
 localhost |    57638 |     1
(3 rows)

SELECT * FROM citus_set_node_property('localhost', :master_port, 'shouldhaveshards', false);
 citus_set_node_property
---------------------------------------------------------------------

(1 row)

SELECT * FROM rebalance_table_shards(shard_transfer_mode :='force_logical');
 rebalance_table_shards
---------------------------------------------------------------------

(1 row)

CALL citus_cleanup_orphaned_resources();
SELECT nodename, nodeport, COUNT(*)
  FROM pg_dist_placement AS placement,
       pg_dist_node AS node
 WHERE placement.groupid = node.groupid
   AND node.noderole = 'primary' GROUP BY nodename, nodeport ORDER BY 1,2;
 nodename  | nodeport | count
---------------------------------------------------------------------
 localhost |    57637 |     2
 localhost |    57638 |     2
(2 rows)

RESET search_path;
SET client_min_messages TO WARNING;
DROP SCHEMA citus_drain_node CASCADE;