Introduce citus_copy_shard_placement UDF with nodeid

src/backend/distributed/operations/shard_rebalancer.c

@@ -1953,12 +1953,10 @@ RebalanceTableShardsBackground(RebalanceOptions *options, Oid shardReplicationMo
 		resetStringInfo(&buf);
 
 		appendStringInfo(&buf,
-						 "SELECT pg_catalog.citus_move_shard_placement(%ld,%s,%u,%s,%u,%s)",
+						 "SELECT pg_catalog.citus_move_shard_placement(%ld,%u,%u,%s)",
 						 move->shardId,
-						 quote_literal_cstr(move->sourceNode->workerName),
+						 move->sourceNode->nodeId,
-						 move->sourceNode->workerPort,
+						 move->targetNode->nodeId,
-						 quote_literal_cstr(move->targetNode->workerName),
-						 move->targetNode->workerPort,
 						 quote_literal_cstr(shardTranferModeLabel));
 
 		BackgroundTask *task = ScheduleBackgroundTask(jobId, GetUserId(), buf.data,

src/backend/distributed/operations/shard_transfer.c

@@ -131,6 +131,7 @@ static char * CreateShardCopyCommand(ShardInterval *shard, WorkerNode *targetNod
 
 /* declarations for dynamic loading */
 PG_FUNCTION_INFO_V1(citus_copy_shard_placement);
+PG_FUNCTION_INFO_V1(citus_copy_shard_placement_with_nodeid);
 PG_FUNCTION_INFO_V1(master_copy_shard_placement);
 PG_FUNCTION_INFO_V1(citus_move_shard_placement);
 PG_FUNCTION_INFO_V1(master_move_shard_placement);
@@ -169,6 +170,36 @@ citus_copy_shard_placement(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * citus_copy_shard_placement_with_nodeid implements a user-facing UDF to copy a placement
+ * from a source node to a target node, including all co-located placements.
+ */
+Datum
+citus_copy_shard_placement_with_nodeid(PG_FUNCTION_ARGS)
+{
+	CheckCitusVersion(ERROR);
+	EnsureCoordinator();
+
+	int64 shardId = PG_GETARG_INT64(0);
+	uint32 sourceNodeId = PG_GETARG_INT32(1);
+	uint32 targetNodeId = PG_GETARG_INT32(2);
+	Oid shardReplicationModeOid = PG_GETARG_OID(3);
+
+	bool missingOk = false;
+	WorkerNode *sourceNode = FindNodeWithNodeId(sourceNodeId, missingOk);
+	WorkerNode *targetNode = FindNodeWithNodeId(targetNodeId, missingOk);
+
+	char shardReplicationMode = LookupShardTransferMode(shardReplicationModeOid);
+
+	ReplicateColocatedShardPlacement(shardId,
+									 sourceNode->workerName, sourceNode->workerPort,
+									 targetNode->workerName, targetNode->workerPort,
+									 shardReplicationMode);
+
+	PG_RETURN_VOID();
+}
+
+
 /*
  * master_copy_shard_placement is a wrapper function for old UDF name.
  */

src/backend/distributed/sql/citus--11.1-1--11.2-1.sql

@@ -13,6 +13,8 @@ DROP FUNCTION pg_catalog.worker_append_table_to_shard(text, text, text, integer)
 #include "udfs/citus_task_wait/11.2-1.sql"
 #include "udfs/citus_prepare_pg_upgrade/11.2-1.sql"
 #include "udfs/citus_finish_pg_upgrade/11.2-1.sql"
+#include "udfs/citus_copy_shard_placement/11.2-1.sql"
+#include "udfs/citus_move_shard_placement_with_nodeid/11.2-1.sql"
 
 -- drop orphaned shards after inserting records for them into pg_dist_cleanup
 INSERT INTO pg_dist_cleanup

src/backend/distributed/sql/downgrades/citus--11.2-1--11.1-1.sql

@@ -25,3 +25,6 @@ COMMENT ON FUNCTION pg_catalog.worker_append_table_to_shard(text, text, text, in
 DROP FUNCTION pg_catalog.citus_task_wait(bigint, pg_catalog.citus_task_status);
 #include "../udfs/citus_prepare_pg_upgrade/11.1-1.sql"
 #include "../udfs/citus_finish_pg_upgrade/11.1-1.sql"
+
+DROP FUNCTION pg_catalog.citus_copy_shard_placement(bigint, integer, integer, citus.shard_transfer_mode);
+DROP FUNCTION pg_catalog.citus_move_shard_placement_with_nodeid;

src/backend/distributed/sql/udfs/citus_copy_shard_placement/11.2-1.sql

@@ -0,0 +1,35 @@
+-- citus_copy_shard_placement, but with nodeid
+CREATE FUNCTION pg_catalog.citus_copy_shard_placement(
+	shard_id bigint,
+	source_node_id integer,
+	target_node_id integer,
+	transfer_mode citus.shard_transfer_mode default 'auto')
+RETURNS void
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_copy_shard_placement_with_nodeid$$;
+
+COMMENT ON FUNCTION pg_catalog.citus_copy_shard_placement(
+	shard_id bigint,
+	source_node_id integer,
+	target_node_id integer,
+	transfer_mode citus.shard_transfer_mode)
+IS 'copy a shard from the source node to the destination node';
+
+CREATE OR REPLACE FUNCTION pg_catalog.citus_copy_shard_placement(
+	shard_id bigint,
+	source_node_name text,
+	source_node_port integer,
+	target_node_name text,
+	target_node_port integer,
+	transfer_mode citus.shard_transfer_mode default 'auto')
+RETURNS void
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_copy_shard_placement$$;
+
+COMMENT ON FUNCTION pg_catalog.citus_copy_shard_placement(shard_id bigint,
+	source_node_name text,
+	source_node_port integer,
+	target_node_name text,
+	target_node_port integer,
+	shard_transfer_mode citus.shard_transfer_mode)
+IS 'copy a shard from the source node to the destination node';

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

@@ -1,5 +1,21 @@
-DROP FUNCTION pg_catalog.citus_copy_shard_placement;
+-- citus_copy_shard_placement, but with nodeid
 CREATE FUNCTION pg_catalog.citus_copy_shard_placement(
+	shard_id bigint,
+	source_node_id integer,
+	target_node_id integer,
+	transfer_mode citus.shard_transfer_mode default 'auto')
+RETURNS void
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_copy_shard_placement_with_nodeid$$;
+
+COMMENT ON FUNCTION pg_catalog.citus_copy_shard_placement(
+	shard_id bigint,
+	source_node_id integer,
+	target_node_id integer,
+	transfer_mode citus.shard_transfer_mode)
+IS 'copy a shard from the source node to the destination node';
+
+CREATE OR REPLACE FUNCTION pg_catalog.citus_copy_shard_placement(
 	shard_id bigint,
 	source_node_name text,
 	source_node_port integer,

src/backend/distributed/utils/reference_table_utils.c

@@ -421,13 +421,11 @@ CopyShardPlacementToWorkerNodeQuery(ShardPlacement *sourceShardPlacement,
 
 	appendStringInfo(queryString,
 					 "SELECT citus_copy_shard_placement("
-					 UINT64_FORMAT ", %s, %d, %s, %d, "
+					 UINT64_FORMAT ", %d, %d, "
 								   "transfer_mode := %s)",
 					 sourceShardPlacement->shardId,
-					 quote_literal_cstr(sourceShardPlacement->nodeName),
+					 sourceShardPlacement->nodeId,
-					 sourceShardPlacement->nodePort,
+					 workerNode->nodeId,
-					 quote_literal_cstr(workerNode->workerName),
-					 workerNode->workerPort,
 					 quote_literal_cstr(transferModeString));
 
 	return queryString;

src/test/regress/expected/citus_copy_shard_placement.out

@@ -42,26 +42,28 @@ INSERT INTO data VALUES ('key-1', 'value-1');
 INSERT INTO data VALUES ('key-2', 'value-2');
 INSERT INTO history VALUES ('key-1', '2020-02-01', 'old');
 INSERT INTO history VALUES ('key-1', '2019-10-01', 'older');
+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
 -- verify we error out if no healthy placement exists at source
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_1_port,
+           :worker_1_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 ERROR:  could not find placement matching "localhost:xxxxx"
 HINT:  Confirm the placement still exists and try again.
 -- verify we error out if source and destination are the same
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 ERROR:  cannot copy shard to the same node
 -- verify we warn if target already contains a healthy placement
 SELECT citus_copy_shard_placement(
            (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='ref_table'::regclass::oid),
-           'localhost', :worker_1_port,
+           :worker_1_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 WARNING:  shard is already present on node localhost:xxxxx
 DETAIL:  Copy may have already completed.
@@ -75,15 +77,15 @@ INSERT INTO ref_table SELECT 1, value FROM data;
 ALTER TABLE data ADD CONSTRAINT distfk FOREIGN KEY (value) REFERENCES ref_table (b) MATCH FULL;
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_1_port);
+           :worker_1_node);
 ERROR:  cannot replicate shards with foreign keys
 ALTER TABLE data DROP CONSTRAINT distfk;
 -- replicate shard that contains key-1
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_1_port,
+           :worker_1_node,
            transfer_mode := 'block_writes');
  citus_copy_shard_placement
 ---------------------------------------------------------------------

src/test/regress/expected/multi_extension.out

@@ -1266,6 +1266,7 @@ SELECT * FROM multi_extension.print_extension_changes();
  function get_rebalance_progress() TABLE(sessionid integer, table_name regclass, shardid bigint, shard_size bigint, sourcename text, sourceport integer, targetname text, targetport integer, progress bigint, source_shard_size bigint, target_shard_size bigint, operation_type text) |
  function worker_append_table_to_shard(text,text,text,integer) void                                                                                                                                                                                                                     |
  function worker_split_shard_replication_setup(split_shard_info[]) SETOF replication_slot_info                                                                                                                                                                                          |
+                                                                                                                                                                                                                                                                                        | function citus_copy_shard_placement(bigint,integer,integer,citus.shard_transfer_mode) void
                                                                                                                                                                                                                                                                                         | function citus_get_node_clock() cluster_clock
                                                                                                                                                                                                                                                                                         | function citus_get_transaction_clock() cluster_clock
                                                                                                                                                                                                                                                                                         | function citus_internal_adjust_local_clock_to_remote(cluster_clock) void
@@ -1295,7 +1296,7 @@ SELECT * FROM multi_extension.print_extension_changes();
                                                                                                                                                                                                                                                                                         | operator family cluster_clock_ops for access method btree
                                                                                                                                                                                                                                                                                         | sequence pg_dist_clock_logical_seq
                                                                                                                                                                                                                                                                                         | type cluster_clock
-(32 rows)
+(33 rows)
 
 DROP TABLE multi_extension.prev_objects, multi_extension.extension_diff;
 -- show running version

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -35,6 +35,7 @@ ORDER BY 1;
  function citus_cleanup_orphaned_shards()
  function citus_conninfo_cache_invalidate()
  function citus_coordinator_nodeid()
+ function citus_copy_shard_placement(bigint,integer,integer,citus.shard_transfer_mode)
  function citus_copy_shard_placement(bigint,text,integer,text,integer,citus.shard_transfer_mode)
  function citus_create_restore_point(text)
  function citus_disable_node(text,integer,boolean)
@@ -312,5 +313,5 @@ ORDER BY 1;
  view citus_stat_statements
  view pg_dist_shard_placement
  view time_partitions
-(304 rows)
+(305 rows)
 

src/test/regress/sql/citus_copy_shard_placement.sql

@@ -34,25 +34,28 @@ INSERT INTO data VALUES ('key-2', 'value-2');
 INSERT INTO history VALUES ('key-1', '2020-02-01', 'old');
 INSERT INTO history VALUES ('key-1', '2019-10-01', 'older');
 
+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
+
 -- verify we error out if no healthy placement exists at source
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_1_port,
+           :worker_1_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 
 -- verify we error out if source and destination are the same
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 
 -- verify we warn if target already contains a healthy placement
 SELECT citus_copy_shard_placement(
            (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='ref_table'::regclass::oid),
-           'localhost', :worker_1_port,
+           :worker_1_node,
-           'localhost', :worker_2_port,
+           :worker_2_node,
            transfer_mode := 'block_writes');
 
 -- verify we error out if table has foreign key constraints
@@ -61,16 +64,16 @@ INSERT INTO ref_table SELECT 1, value FROM data;
 ALTER TABLE data ADD CONSTRAINT distfk FOREIGN KEY (value) REFERENCES ref_table (b) MATCH FULL;
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_1_port);
+           :worker_1_node);
 
 ALTER TABLE data DROP CONSTRAINT distfk;
 
 -- replicate shard that contains key-1
 SELECT citus_copy_shard_placement(
            get_shard_id_for_distribution_column('data', 'key-1'),
-           'localhost', :worker_2_port,
+           :worker_2_node,
-           'localhost', :worker_1_port,
+           :worker_1_node,
            transfer_mode := 'block_writes');
 
 -- forcefully mark the old replica as inactive