--
-- MULTI_COLOCATED_SHARD_REBALANCE
--
ALTER SEQUENCE pg_catalog.pg_dist_shardid_seq RESTART 13000000;
SET citus.shard_count TO 6;
SET citus.shard_replication_factor TO 1;
-- create distributed tables
CREATE TABLE table1_group1 ( id int PRIMARY KEY);
SELECT create_distributed_table('table1_group1', 'id', 'hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE table2_group1 ( id int );
SELECT create_distributed_table('table2_group1', 'id', 'hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SET citus.shard_count TO 8;
CREATE TABLE table5_groupX ( id int );
SELECT create_distributed_table('table5_groupX', 'id', 'hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE table6_append ( id int );
SELECT create_distributed_table('table6_append', 'id', 'append');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

SELECT master_create_empty_shard('table6_append');
 master_create_empty_shard
---------------------------------------------------------------------
                  13000020
(1 row)

SELECT master_create_empty_shard('table6_append');
 master_create_empty_shard
---------------------------------------------------------------------
                  13000021
(1 row)

-- Mark tables as non-mx tables, in order to be able to test citus_copy_shard_placement
UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid IN
    ('table1_group1'::regclass, 'table2_group1'::regclass, 'table5_groupX'::regclass);
-- test copy
-- test copying colocated shards
-- status before shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000000 | table1_group1 |    57637
 13000001 | table1_group1 |    57638
 13000002 | table1_group1 |    57637
 13000003 | table1_group1 |    57638
 13000004 | table1_group1 |    57637
 13000005 | table1_group1 |    57638
 13000006 | table2_group1 |    57637
 13000007 | table2_group1 |    57638
 13000008 | table2_group1 |    57637
 13000009 | table2_group1 |    57638
 13000010 | table2_group1 |    57637
 13000011 | table2_group1 |    57638
(12 rows)

-- try to copy colocated shards without a replica identity
SELECT citus_copy_shard_placement(13000000, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
ERROR:  cannot use logical replication to transfer shards of the relation table2_group1 since it doesn't have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the shard will error out during logical replication unless there is a REPLICA IDENTITY or PRIMARY KEY.
HINT:  If you wish to continue without a replica identity set the shard_transfer_mode to 'force_logical' or 'block_writes'.
-- copy colocated shards
SELECT citus_copy_shard_placement(13000000, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

-- error out if trying to move a shard that already has a placement on the target
SELECT citus_move_shard_placement(13000000, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
ERROR:  shard xxxxx already exists in the target node
-- status after shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000000 | table1_group1 |    57637
 13000000 | table1_group1 |    57638
 13000001 | table1_group1 |    57638
 13000002 | table1_group1 |    57637
 13000003 | table1_group1 |    57638
 13000004 | table1_group1 |    57637
 13000005 | table1_group1 |    57638
 13000006 | table2_group1 |    57637
 13000006 | table2_group1 |    57638
 13000007 | table2_group1 |    57638
 13000008 | table2_group1 |    57637
 13000009 | table2_group1 |    57638
 13000010 | table2_group1 |    57637
 13000011 | table2_group1 |    57638
(14 rows)

-- also connect worker to verify we successfully copied given shard (and other colocated shards)
\c - - - :worker_2_port
SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table1_group1_13000000'::regclass;
 Column |  Type   | Modifiers
---------------------------------------------------------------------
 id     | integer | not null
(1 row)

SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table2_group1_13000006'::regclass;
 Column |  Type   | Modifiers
---------------------------------------------------------------------
 id     | integer |
(1 row)

\c - - - :master_port
-- copy colocated shards again to see warning
SELECT citus_copy_shard_placement(13000000, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
WARNING:  shard is already present on node localhost:xxxxx
DETAIL:  Copy may have already completed.
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

-- test copying NOT colocated shard
-- status before shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table5_groupX'::regclass
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000012 | table5_groupx |    57637
 13000013 | table5_groupx |    57638
 13000014 | table5_groupx |    57637
 13000015 | table5_groupx |    57638
 13000016 | table5_groupx |    57637
 13000017 | table5_groupx |    57638
 13000018 | table5_groupx |    57637
 13000019 | table5_groupx |    57638
(8 rows)

-- copy NOT colocated shard
SELECT citus_copy_shard_placement(13000012, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table5_groupX'::regclass
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000012 | table5_groupx |    57637
 13000012 | table5_groupx |    57638
 13000013 | table5_groupx |    57638
 13000014 | table5_groupx |    57637
 13000015 | table5_groupx |    57638
 13000016 | table5_groupx |    57637
 13000017 | table5_groupx |    57638
 13000018 | table5_groupx |    57637
 13000019 | table5_groupx |    57638
(9 rows)

-- test copying shard in append distributed table
-- status before shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table6_append'::regclass
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000020 | table6_append |    57638
 13000021 | table6_append |    57637
(2 rows)

-- copy shard in append distributed table
SELECT citus_copy_shard_placement(13000020, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical');
 citus_copy_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard copy
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table6_append'::regclass
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000020 | table6_append |    57637
 13000020 | table6_append |    57638
 13000021 | table6_append |    57637
(3 rows)

-- test move
-- test moving colocated shards
-- status before shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000000 | table1_group1 |    57637
 13000000 | table1_group1 |    57638
 13000001 | table1_group1 |    57638
 13000002 | table1_group1 |    57637
 13000003 | table1_group1 |    57638
 13000004 | table1_group1 |    57637
 13000005 | table1_group1 |    57638
 13000006 | table2_group1 |    57637
 13000006 | table2_group1 |    57638
 13000007 | table2_group1 |    57638
 13000008 | table2_group1 |    57637
 13000009 | table2_group1 |    57638
 13000010 | table2_group1 |    57637
 13000011 | table2_group1 |    57638
(14 rows)

-- move colocated shards
SELECT master_move_shard_placement(13000001, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
    AND sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000000 | table1_group1 |    57637
 13000000 | table1_group1 |    57638
 13000001 | table1_group1 |    57637
 13000002 | table1_group1 |    57637
 13000003 | table1_group1 |    57638
 13000004 | table1_group1 |    57637
 13000005 | table1_group1 |    57638
 13000006 | table2_group1 |    57637
 13000006 | table2_group1 |    57638
 13000007 | table2_group1 |    57637
 13000008 | table2_group1 |    57637
 13000009 | table2_group1 |    57638
 13000010 | table2_group1 |    57637
 13000011 | table2_group1 |    57638
(14 rows)

-- moving the shard again is idempotent
SELECT citus_move_shard_placement(13000001, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical');
WARNING:  shard is already present on node localhost:xxxxx
DETAIL:  Move may have already completed.
 citus_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- also connect worker to verify we successfully moved given shard (and other colocated shards)
\c - - - :worker_1_port
SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table1_group1_13000001'::regclass;
 Column |  Type   | Modifiers
---------------------------------------------------------------------
 id     | integer | not null
(1 row)

SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table2_group1_13000007'::regclass;
 Column |  Type   | Modifiers
---------------------------------------------------------------------
 id     | integer |
(1 row)

\c - - - :master_port
-- test moving NOT colocated shard
-- status before shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table5_groupX'::regclass
    AND sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000012 | table5_groupx |    57637
 13000012 | table5_groupx |    57638
 13000013 | table5_groupx |    57638
 13000014 | table5_groupx |    57637
 13000015 | table5_groupx |    57638
 13000016 | table5_groupx |    57637
 13000017 | table5_groupx |    57638
 13000018 | table5_groupx |    57637
 13000019 | table5_groupx |    57638
(9 rows)

-- move NOT colocated shard
SELECT master_move_shard_placement(13000013, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table5_groupX'::regclass AND
    sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000012 | table5_groupx |    57637
 13000012 | table5_groupx |    57638
 13000013 | table5_groupx |    57637
 13000014 | table5_groupx |    57637
 13000015 | table5_groupx |    57638
 13000016 | table5_groupx |    57637
 13000017 | table5_groupx |    57638
 13000018 | table5_groupx |    57637
 13000019 | table5_groupx |    57638
(9 rows)

-- test moving shard in append distributed table
-- status before shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table6_append'::regclass
    AND sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000020 | table6_append |    57637
 13000020 | table6_append |    57638
 13000021 | table6_append |    57637
(3 rows)

-- move shard in append distributed table
SELECT master_move_shard_placement(13000021, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard move
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
    p.logicalrelid = 'table6_append'::regclass AND
    sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000020 | table6_append |    57637
 13000020 | table6_append |    57638
 13000021 | table6_append |    57638
(3 rows)

-- try to move shard from wrong node
SELECT master_move_shard_placement(13000021, 'localhost', :master_port, 'localhost', :worker_1_port, 'force_logical');
ERROR:  could not find placement matching "localhost:xxxxx"
HINT:  Confirm the placement still exists and try again.
-- test shard move with foreign constraints
DROP TABLE IF EXISTS table1_group1, table2_group1;
SET citus.shard_count TO 6;
SET citus.shard_replication_factor TO 1;
-- create distributed tables
CREATE TABLE table1_group1 ( id int PRIMARY KEY);
SELECT create_distributed_table('table1_group1', 'id', 'hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE table2_group1 ( id int, table1_id int, FOREIGN KEY(table1_id) REFERENCES table1_group1(id));
SELECT create_distributed_table('table2_group1', 'table1_id', 'hash');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- Mark the tables as non-mx tables
UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid IN
    ('table1_group1'::regclass, 'table2_group1'::regclass);
-- status before shard rebalance
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
	colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
    AND sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000022 | table1_group1 |    57637
 13000023 | table1_group1 |    57638
 13000024 | table1_group1 |    57637
 13000025 | table1_group1 |    57638
 13000026 | table1_group1 |    57637
 13000027 | table1_group1 |    57638
 13000028 | table2_group1 |    57637
 13000029 | table2_group1 |    57638
 13000030 | table2_group1 |    57637
 13000031 | table2_group1 |    57638
 13000032 | table2_group1 |    57637
 13000033 | table2_group1 |    57638
(12 rows)

SELECT master_move_shard_placement(13000022, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'block_writes');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

-- status after shard rebalance
SELECT s.shardid, s.logicalrelid::regclass, sp.nodeport
FROM
    pg_dist_partition p, pg_dist_shard s, pg_dist_shard_placement sp
WHERE
    p.logicalrelid = s.logicalrelid AND
    s.shardid = sp.shardid AND
	colocationid = (SELECT colocationid FROM pg_dist_partition WHERE logicalrelid = 'table1_group1'::regclass)
    AND sp.shardstate != 4
ORDER BY s.shardid, sp.nodeport;
 shardid  | logicalrelid  | nodeport
---------------------------------------------------------------------
 13000022 | table1_group1 |    57638
 13000023 | table1_group1 |    57638
 13000024 | table1_group1 |    57637
 13000025 | table1_group1 |    57638
 13000026 | table1_group1 |    57637
 13000027 | table1_group1 |    57638
 13000028 | table2_group1 |    57638
 13000029 | table2_group1 |    57638
 13000030 | table2_group1 |    57637
 13000031 | table2_group1 |    57638
 13000032 | table2_group1 |    57637
 13000033 | table2_group1 |    57638
(12 rows)

-- also connect worker to verify we successfully moved given shard (and other colocated shards)
\c - - - :worker_2_port
SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table1_group1_13000022'::regclass;
 Column |  Type   | Modifiers
---------------------------------------------------------------------
 id     | integer | not null
(1 row)

SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table2_group1_13000028'::regclass;
  Column   |  Type   | Modifiers
---------------------------------------------------------------------
 id        | integer |
 table1_id | integer |
(2 rows)

-- make sure that we've created the foreign keys
SELECT  "Constraint", "Definition" FROM table_fkeys
  WHERE "Constraint" LIKE 'table2_group%' OR "Constraint" LIKE 'table1_group%';
              Constraint               |                          Definition
---------------------------------------------------------------------
 table2_group1_table1_id_fkey          | FOREIGN KEY (table1_id) REFERENCES table1_group1(id)
 table2_group1_table1_id_fkey_13000028 | FOREIGN KEY (table1_id) REFERENCES table1_group1_13000022(id)
 table2_group1_table1_id_fkey_13000029 | FOREIGN KEY (table1_id) REFERENCES table1_group1_13000023(id)
 table2_group1_table1_id_fkey_13000031 | FOREIGN KEY (table1_id) REFERENCES table1_group1_13000025(id)
 table2_group1_table1_id_fkey_13000033 | FOREIGN KEY (table1_id) REFERENCES table1_group1_13000027(id)
(5 rows)

\c - - - :master_port
-- test shard copy with foreign constraints
-- we expect it to error out because we do not support foreign constraints with replication factor > 1
SELECT citus_copy_shard_placement(13000022, 'localhost', :worker_2_port, 'localhost', :worker_1_port);
ERROR:  cannot replicate shards with foreign keys
-- lets also test that master_move_shard_placement doesn't break serials
CREATE TABLE serial_move_test (key int, other_val serial);
SET citus.shard_replication_factor TO 1;
SELECT create_distributed_table('serial_move_test', 'key');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- key 15 goes to shard xxxxx
INSERT INTO serial_move_test (key) VALUES (15) RETURNING *;
 key | other_val
---------------------------------------------------------------------
  15 |         1
(1 row)

INSERT INTO serial_move_test (key) VALUES (15) RETURNING *;
 key | other_val
---------------------------------------------------------------------
  15 |         2
(1 row)

-- confirm the shard id
SELECT * FROM run_command_on_placements('serial_move_test', 'SELECT DISTINCT key FROM %s WHERE key = 15') WHERE result = '15' AND shardid = 13000034;
 nodename  | nodeport | shardid  | success | result
---------------------------------------------------------------------
 localhost |    57637 | 13000034 | t       | 15
(1 row)

SELECT master_move_shard_placement(13000034, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- confirm the successfull move
SELECT * FROM run_command_on_placements('serial_move_test', 'SELECT DISTINCT key FROM %s WHERE key = 15') WHERE result = '15' AND shardid = 13000034;
 nodename  | nodeport | shardid  | success | result
---------------------------------------------------------------------
 localhost |    57638 | 13000034 | t       | 15
(1 row)

-- finally show that serials work fine afterwards
INSERT INTO serial_move_test (key) VALUES (15) RETURNING *;
 key | other_val
---------------------------------------------------------------------
  15 |         3
(1 row)

INSERT INTO serial_move_test (key) VALUES (15) RETURNING *;
 key | other_val
---------------------------------------------------------------------
  15 |         4
(1 row)

-- lets do some failure testing
CREATE TABLE logical_failure_test (key int);
SET citus.shard_replication_factor TO 1;
SET citus.shard_count TO 4;
SELECT create_distributed_table('logical_failure_test', 'key');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- ensure that the shard is created for this user
\c - - - :worker_2_port
\dt logical_failure_test_13000038
                     List of relations
 Schema |             Name              | Type  |  Owner
---------------------------------------------------------------------
 public | logical_failure_test_13000038 | table | postgres
(1 row)

DROP TABLE logical_failure_test_13000038;
-- should fail since the command wouldn't be able to connect to the worker_1
\c - - - :master_port
SELECT master_move_shard_placement(13000038, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical');
WARNING:  relation "public.logical_failure_test_13000038" does not exist
CONTEXT:  while executing command on localhost:xxxxx
ERROR:  cannot get the size because of a connection error
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

DROP TABLE logical_failure_test;
-- lets test the logical replication modes
CREATE TABLE test_with_pkey (key int PRIMARY KEY, value int NOT NULL);
SET citus.shard_replication_factor TO 1;
SET citus.shard_count TO 4;
SELECT create_distributed_table('test_with_pkey', 'key', colocate_with => 'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- should succeed since there is a replica identity defined
SELECT master_move_shard_placement(13000042, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- should succeed since we still have a replica identity
ALTER TABLE test_with_pkey REPLICA IDENTITY FULL;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'auto');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- make sure we have the replica identity after the move
SELECT result FROM run_command_on_placements( 'test_with_pkey', 'SELECT relreplident FROM pg_class WHERE relname = ''%s''') WHERE shardid = 13000042;
 result
---------------------------------------------------------------------
 f
(1 row)

-- this time should fail since we don't have replica identity any more
ALTER TABLE test_with_pkey REPLICA IDENTITY NOTHING;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'auto');
ERROR:  cannot use logical replication to transfer shards of the relation test_with_pkey since it doesn't have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the shard will error out during logical replication unless there is a REPLICA IDENTITY or PRIMARY KEY.
HINT:  If you wish to continue without a replica identity set the shard_transfer_mode to 'force_logical' or 'block_writes'.
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- make sure we have the replica identity after the move
SELECT result FROM run_command_on_placements( 'test_with_pkey', 'SELECT relreplident FROM pg_class WHERE relname = ''%s''') WHERE shardid = 13000042;
 result
---------------------------------------------------------------------
 n
(1 row)

-- should succeed since we still have a replica identity
ALTER TABLE test_with_pkey REPLICA IDENTITY USING INDEX test_with_pkey_pkey;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- make sure we have the replica identity after the move
SELECT result FROM run_command_on_placements( 'test_with_pkey', 'SELECT relreplident FROM pg_class WHERE relname = ''%s''') WHERE shardid = 13000042;
 result
---------------------------------------------------------------------
 i
(1 row)

-- one final test with shard_transfer_mode auto
CREATE UNIQUE INDEX req_rep_idx ON test_with_pkey(key, value);
ALTER TABLE test_with_pkey REPLICA IDENTITY USING INDEX req_rep_idx;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'auto');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- make sure we have the replica identity after the move
SELECT result FROM run_command_on_placements( 'test_with_pkey', 'SELECT relreplident FROM pg_class WHERE relname = ''%s''') WHERE shardid = 13000042;
 result
---------------------------------------------------------------------
 i
(1 row)

ALTER TABLE test_with_pkey REPLICA IDENTITY NOTHING;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_1_port, 'localhost', :worker_2_port, 'force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- make sure we have the replica identity after the move
SELECT result FROM run_command_on_placements( 'test_with_pkey', 'SELECT relreplident FROM pg_class WHERE relname = ''%s''') WHERE shardid = 13000042;
 result
---------------------------------------------------------------------
 n
(1 row)

-- should succeed but not use logical replication
ALTER TABLE test_with_pkey REPLICA IDENTITY NOTHING;
SET client_min_messages TO DEBUG1;
SELECT master_move_shard_placement(13000042, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'block_writes');
DEBUG:  table "test_with_pkey_13000042" does not exist, skipping
DETAIL:  from localhost:xxxxx
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO DEFAULT;
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- we don't support multiple shard moves in a single transaction
SELECT
    master_move_shard_placement(shardid, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='force_logical')
FROM
    pg_dist_shard_placement where nodeport = :worker_1_port AND
    shardid IN (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_with_pkey'::regclass);
ERROR:  moving multiple shard placements via logical replication in the same transaction is currently not supported
HINT:  If you wish to move multiple shard placements in a single transaction set the shard_transfer_mode to 'block_writes'.
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- similar test with explicit transaction block
BEGIN;
    SELECT master_move_shard_placement(13000042, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

    SELECT master_move_shard_placement(13000044, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='force_logical');
ERROR:  moving multiple shard placements via logical replication in the same transaction is currently not supported
HINT:  If you wish to move multiple shard placements in a single transaction set the shard_transfer_mode to 'block_writes'.
COMMIT;
    SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- we do support the same with block writes
SELECT
    master_move_shard_placement(shardid, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='block_writes')
FROM
    pg_dist_shard_placement where nodeport = :worker_1_port AND
    shardid IN (SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'test_with_pkey'::regclass);
 master_move_shard_placement
---------------------------------------------------------------------


(2 rows)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- we should be able to move shard placements after COMMIT/ABORT
BEGIN;
    SELECT master_move_shard_placement(13000043, 'localhost', :worker_2_port, 'localhost', :worker_1_port, shard_transfer_mode:='force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

COMMIT;
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

SELECT master_move_shard_placement(13000045, 'localhost', :worker_2_port, 'localhost', :worker_1_port, shard_transfer_mode:='force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

BEGIN;
    SELECT master_move_shard_placement(13000043, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

ABORT;
SELECT master_move_shard_placement(13000045, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode:='force_logical');
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- we should be able to move shard placements of partitioend tables
CREATE SCHEMA move_partitions;
CREATE TABLE move_partitions.events (
    id serial,
    t timestamptz default now(),
    payload text
)
PARTITION BY RANGE(t);
SET citus.shard_count TO 6;
SELECT create_distributed_table('move_partitions.events', 'id', colocate_with := 'none');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE move_partitions.events_1 PARTITION OF move_partitions.events
FOR VALUES FROM ('2015-01-01') TO ('2016-01-01');
INSERT INTO move_partitions.events (t, payload)
SELECT '2015-01-01'::date + (interval '1 day' * s), s FROM generate_series(1, 100) s;
SELECT count(*) FROM move_partitions.events;
 count
---------------------------------------------------------------------
   100
(1 row)

-- try to move automatically
SELECT master_move_shard_placement(shardid, 'localhost', :worker_2_port, 'localhost', :worker_1_port)
FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
WHERE logicalrelid = 'move_partitions.events'::regclass AND nodeport = :worker_2_port
AND shardstate != 4
ORDER BY shardid LIMIT 1;
ERROR:  cannot use logical replication to transfer shards of the relation events_1 since it doesn't have a REPLICA IDENTITY or PRIMARY KEY
DETAIL:  UPDATE and DELETE commands on the shard will error out during logical replication unless there is a REPLICA IDENTITY or PRIMARY KEY.
HINT:  If you wish to continue without a replica identity set the shard_transfer_mode to 'force_logical' or 'block_writes'.
SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

-- force logical replication
SELECT master_move_shard_placement(shardid, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'force_logical')
FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
WHERE logicalrelid = 'move_partitions.events'::regclass AND nodeport = :worker_2_port
ORDER BY shardid LIMIT 1;
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM move_partitions.events;
 count
---------------------------------------------------------------------
   100
(1 row)

-- add a primary key to the partition
ALTER TABLE move_partitions.events_1 ADD CONSTRAINT e_1_pk PRIMARY KEY (id);
-- should be able to move automatically now
SELECT master_move_shard_placement(shardid, 'localhost', :worker_2_port, 'localhost', :worker_1_port)
FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
WHERE logicalrelid = 'move_partitions.events'::regclass AND nodeport = :worker_2_port AND shardstate != 4
ORDER BY shardid LIMIT 1;
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM move_partitions.events;
 count
---------------------------------------------------------------------
   100
(1 row)

-- should also be able to move with block writes
SELECT master_move_shard_placement(shardid, 'localhost', :worker_2_port, 'localhost', :worker_1_port, 'block_writes')
FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
WHERE logicalrelid = 'move_partitions.events'::regclass AND nodeport = :worker_2_port AND shardstate != 4
ORDER BY shardid LIMIT 1;
 master_move_shard_placement
---------------------------------------------------------------------

(1 row)

SELECT public.wait_for_resource_cleanup();
 wait_for_resource_cleanup
---------------------------------------------------------------------

(1 row)

SELECT count(*) FROM move_partitions.events;
 count
---------------------------------------------------------------------
   100
(1 row)

-- should have moved all shards to node 1 (2*6 = 12)
SELECT count(*)
FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
WHERE logicalrelid::text LIKE 'move_partitions.events%' AND nodeport = :worker_1_port;
 count
---------------------------------------------------------------------
    12
(1 row)

DROP TABLE move_partitions.events;
-- set back to the defaults and drop the tables
SET client_min_messages TO DEFAULT;
DROP TABLE test_with_pkey;
DROP TABLE table2_group1;
DROP TABLE table1_group1;
DROP TABLE table5_groupX;
DROP TABLE table6_append;
DROP TABLE serial_move_test;
DROP SCHEMA move_partitions CASCADE;
