Only allow moves of shards of distributed tables (#5072)

Moving shards of reference tables was possible in at least one case:
```sql
select citus_disable_node('localhost', 9702);
create table r(x int);
select create_reference_table('r');
set citus.replicate_reference_tables_on_activate = off;
select citus_activate_node('localhost', 9702);
select citus_move_shard_placement(102008, 'localhost', 9701, 'localhost', 9702);
```

This would then remove the reference table shard on the source, causing
all kinds of issues. This fixes that by disallowing all shard moves
except for shards of distributed tables.

Co-authored-by: Onur Tirtir <onurcantirtir@gmail.com>
(cherry picked from commit d1d386a904)

src/backend/distributed/operations/repair_shards.c

@@ -298,7 +298,7 @@ citus_move_shard_placement(PG_FUNCTION_ARGS)
 	ListCell *colocatedShardCell = NULL;
 
 	Oid relationId = RelationIdForShard(shardId);
-	ErrorIfMoveCitusLocalTable(relationId);
+	ErrorIfMoveUnsupportedTableType(relationId);
 	ErrorIfTargetNodeIsNotSafeToMove(targetNodeName, targetNodePort);
 
 	ShardInterval *shardInterval = LoadShardInterval(shardId);
@@ -478,23 +478,40 @@ master_move_shard_placement(PG_FUNCTION_ARGS)
 
 
 /*
- * ErrorIfMoveCitusLocalTable is a helper function for rebalance_table_shards
+ * ErrorIfMoveUnsupportedTableType is a helper function for rebalance_table_shards
  * and citus_move_shard_placement udf's to error out if relation with relationId
- * is a citus local table.
+ * is not a distributed table.
  */
 void
-ErrorIfMoveCitusLocalTable(Oid relationId)
+ErrorIfMoveUnsupportedTableType(Oid relationId)
 {
-	if (!IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+	if (IsCitusTableType(relationId, DISTRIBUTED_TABLE))
 	{
 		return;
 	}
 
 	char *qualifiedRelationName = generate_qualified_relation_name(relationId);
-	ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("table %s is a local table, moving shard of "
-						   "a local table added to metadata is currently "
-						   "not supported", qualifiedRelationName)));
+	if (!IsCitusTable(relationId))
+	{
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("table %s is a regular postgres table, you can "
+							   "only move shards of a citus table",
+							   qualifiedRelationName)));
+	}
+	else if (IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+	{
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("table %s is a local table, moving shard of "
+							   "a local table added to metadata is currently "
+							   "not supported", qualifiedRelationName)));
+	}
+	else if (IsCitusTableType(relationId, REFERENCE_TABLE))
+	{
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("table %s is a reference table, moving shard of "
+							   "a reference table is not supported",
+							   qualifiedRelationName)));
+	}
 }
 
 

src/backend/distributed/operations/shard_rebalancer.c

@@ -794,7 +794,7 @@ rebalance_table_shards(PG_FUNCTION_ARGS)
 	if (!PG_ARGISNULL(0))
 	{
 		Oid relationId = PG_GETARG_OID(0);
-		ErrorIfMoveCitusLocalTable(relationId);
+		ErrorIfMoveUnsupportedTableType(relationId);
 
 		relationIdList = list_make1_oid(relationId);
 	}
@@ -997,7 +997,7 @@ get_rebalance_table_shards_plan(PG_FUNCTION_ARGS)
 	if (!PG_ARGISNULL(0))
 	{
 		Oid relationId = PG_GETARG_OID(0);
-		ErrorIfMoveCitusLocalTable(relationId);
+		ErrorIfMoveUnsupportedTableType(relationId);
 
 		relationIdList = list_make1_oid(relationId);
 	}

src/include/distributed/coordinator_protocol.h

@@ -282,7 +282,6 @@ extern ShardPlacement * SearchShardPlacementInListOrError(List *shardPlacementLi
 														  uint32 nodePort);
 extern void ErrorIfTargetNodeIsNotSafeToMove(const char *targetNodeName, int
 											 targetNodePort);
-extern void ErrorIfMoveCitusLocalTable(Oid relationId);
 extern char LookupShardTransferMode(Oid shardReplicationModeOid);
 extern void BlockWritesToShardList(List *shardList);
 extern List * WorkerApplyShardDDLCommandList(List *ddlCommandList, int64 shardId);

src/include/distributed/repair_shards.h

@@ -13,3 +13,4 @@
 
 extern uint64 ShardListSizeInBytes(List *colocatedShardList,
 								   char *workerNodeName, uint32 workerNodePort);
+extern void ErrorIfMoveUnsupportedTableType(Oid relationId);

src/test/regress/expected/shard_rebalancer.out

@@ -1,13 +1,7 @@
 --
 -- MUTLI_SHARD_REBALANCER
 --
-CREATE TABLE dist_table_test(a int primary key);
-SELECT create_distributed_table('dist_table_test', 'a');
- create_distributed_table
----------------------------------------------------------------------
-
-(1 row)
-
+SET citus.next_shard_id TO 433000;
 CREATE TABLE ref_table_test(a int primary key);
 SELECT create_reference_table('ref_table_test');
  create_reference_table
@@ -15,6 +9,14 @@ SELECT create_reference_table('ref_table_test');
 
 (1 row)
 
+CREATE TABLE dist_table_test(a int primary key);
+SELECT create_distributed_table('dist_table_test', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE postgres_table_test(a int primary key);
 -- make sure that all rebalance operations works fine when
 -- reference tables are replicated to the coordinator
 SELECT 1 FROM master_add_node('localhost', :master_port, groupId=>0);
@@ -41,6 +43,7 @@ SELECT rebalance_table_shards();
 CALL citus_cleanup_orphaned_shards();
 -- test that calling rebalance_table_shards without specifying relation
 -- wouldn't move shard of the citus local table.
+SET citus.next_shard_id TO 433100;
 CREATE TABLE citus_local_table(a int, b int);
 SELECT citus_add_local_table_to_metadata('citus_local_table');
  citus_add_local_table_to_metadata
@@ -56,11 +59,31 @@ SELECT rebalance_table_shards();
 (1 row)
 
 CALL citus_cleanup_orphaned_shards();
+-- Check that rebalance_table_shards and get_rebalance_table_shards_plan fail
+-- for any type of table, but distributed tables.
+SELECT rebalance_table_shards('ref_table_test');
+ERROR:  table public.ref_table_test is a reference table, moving shard of a reference table is not supported
+SELECT rebalance_table_shards('postgres_table_test');
+ERROR:  table public.postgres_table_test is a regular postgres table, you can only move shards of a citus table
+SELECT rebalance_table_shards('citus_local_table');
+ERROR:  table public.citus_local_table is a local table, moving shard of a local table added to metadata is currently not supported
+SELECT get_rebalance_table_shards_plan('ref_table_test');
+ERROR:  table public.ref_table_test is a reference table, moving shard of a reference table is not supported
+SELECT get_rebalance_table_shards_plan('postgres_table_test');
+ERROR:  table public.postgres_table_test is a regular postgres table, you can only move shards of a citus table
+SELECT get_rebalance_table_shards_plan('citus_local_table');
+ERROR:  table public.citus_local_table is a local table, moving shard of a local table added to metadata is currently not supported
+-- Check that citus_move_shard_placement fails for shards belonging reference
+-- tables or citus local tables
+SELECT citus_move_shard_placement(433000, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
+ERROR:  table public.ref_table_test is a reference table, moving shard of a reference table is not supported
+SELECT citus_move_shard_placement(433100, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
+ERROR:  table public.citus_local_table is a local table, moving shard of a local table added to metadata is currently not supported
 -- show that citus local table shard is still on the coordinator
 SELECT tablename FROM pg_catalog.pg_tables where tablename like 'citus_local_table_%';
         tablename
 ---------------------------------------------------------------------
- citus_local_table_102047
+ citus_local_table_433100
 (1 row)
 
 -- also check that we still can access shard relation, not the shell table
@@ -111,7 +134,7 @@ CALL citus_cleanup_orphaned_shards();
 SELECT tablename FROM pg_catalog.pg_tables where tablename like 'citus_local_table_%';
         tablename
 ---------------------------------------------------------------------
- citus_local_table_102047
+ citus_local_table_433100
 (1 row)
 
 -- also check that we still can access shard relation, not the shell table
@@ -182,7 +205,7 @@ NOTICE:  Copying shard xxxxx from localhost:xxxxx to localhost:xxxxx ...
 
 (1 row)
 
-DROP TABLE dist_table_test, dist_table_test_2, ref_table_test;
+DROP TABLE dist_table_test, dist_table_test_2, ref_table_test, postgres_table_test;
 RESET citus.shard_count;
 RESET citus.shard_replication_factor;
 -- Create a user to test multiuser usage of rebalancer functions

src/test/regress/sql/shard_rebalancer.sql

@@ -2,10 +2,12 @@
 -- MUTLI_SHARD_REBALANCER
 --
 
-CREATE TABLE dist_table_test(a int primary key);
-SELECT create_distributed_table('dist_table_test', 'a');
+SET citus.next_shard_id TO 433000;
 CREATE TABLE ref_table_test(a int primary key);
 SELECT create_reference_table('ref_table_test');
+CREATE TABLE dist_table_test(a int primary key);
+SELECT create_distributed_table('dist_table_test', 'a');
+CREATE TABLE postgres_table_test(a int primary key);
 
 -- make sure that all rebalance operations works fine when
 -- reference tables are replicated to the coordinator
@@ -20,6 +22,7 @@ CALL citus_cleanup_orphaned_shards();
 
 -- test that calling rebalance_table_shards without specifying relation
 -- wouldn't move shard of the citus local table.
+SET citus.next_shard_id TO 433100;
 CREATE TABLE citus_local_table(a int, b int);
 SELECT citus_add_local_table_to_metadata('citus_local_table');
 INSERT INTO citus_local_table VALUES (1, 2);
@@ -27,6 +30,20 @@ INSERT INTO citus_local_table VALUES (1, 2);
 SELECT rebalance_table_shards();
 CALL citus_cleanup_orphaned_shards();
 
+-- Check that rebalance_table_shards and get_rebalance_table_shards_plan fail
+-- for any type of table, but distributed tables.
+SELECT rebalance_table_shards('ref_table_test');
+SELECT rebalance_table_shards('postgres_table_test');
+SELECT rebalance_table_shards('citus_local_table');
+SELECT get_rebalance_table_shards_plan('ref_table_test');
+SELECT get_rebalance_table_shards_plan('postgres_table_test');
+SELECT get_rebalance_table_shards_plan('citus_local_table');
+
+-- Check that citus_move_shard_placement fails for shards belonging reference
+-- tables or citus local tables
+SELECT citus_move_shard_placement(433000, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
+SELECT citus_move_shard_placement(433100, 'localhost', :worker_1_port, 'localhost', :worker_2_port);
+
 -- show that citus local table shard is still on the coordinator
 SELECT tablename FROM pg_catalog.pg_tables where tablename like 'citus_local_table_%';
 -- also check that we still can access shard relation, not the shell table
@@ -83,7 +100,7 @@ SELECT pg_sleep(.1); -- wait to make sure the config has changed before running
 SET citus.shard_replication_factor TO 2;
 SELECT replicate_table_shards('dist_table_test_2',  max_shard_copies := 4,  shard_transfer_mode:='block_writes');
 
-DROP TABLE dist_table_test, dist_table_test_2, ref_table_test;
+DROP TABLE dist_table_test, dist_table_test_2, ref_table_test, postgres_table_test;
 RESET citus.shard_count;
 RESET citus.shard_replication_factor;