reflect on feedback

src/backend/distributed/metadata/metadata_sync.c

@@ -179,6 +179,7 @@ PG_FUNCTION_INFO_V1(citus_internal_delete_colocation_metadata);
 PG_FUNCTION_INFO_V1(citus_internal_add_tenant_schema);
 PG_FUNCTION_INFO_V1(citus_internal_delete_tenant_schema);
 PG_FUNCTION_INFO_V1(citus_internal_update_none_dist_table_metadata);
+PG_FUNCTION_INFO_V1(citus_internal_shard_group_set_needsisolatednode);
 
 
 static bool got_SIGTERM = false;
@@ -1202,7 +1203,7 @@ ShardListInsertCommand(List *shardIntervalList)
 	StringInfo insertPlacementCommand = makeStringInfo();
 	appendStringInfo(insertPlacementCommand,
 					 "WITH placement_data(shardid, "
-					 "shardlength, groupid, placementid)  AS (VALUES ");
+					 "shardlength, groupid, placementid, needsisolatednode)  AS (VALUES ");
 
 	ShardInterval *shardInterval = NULL;
 	bool firstPlacementProcessed = false;
@@ -1225,19 +1226,20 @@ ShardListInsertCommand(List *shardIntervalList)
 			firstPlacementProcessed = true;
 
 			appendStringInfo(insertPlacementCommand,
-							 "(%ld, %ld, %d, %ld)",
+							 "(%ld, %ld, %d, %ld, %s)",
 							 shardId,
 							 placement->shardLength,
 							 placement->groupId,
-							 placement->placementId);
+							 placement->placementId,
+							 placement->needsIsolatedNode ? "true" : "false");
 		}
 	}
 
 	appendStringInfo(insertPlacementCommand, ") ");
 
 	appendStringInfo(insertPlacementCommand,
 					 "SELECT citus_internal_add_placement_metadata("
-					 "shardid, shardlength, groupid, placementid) "
+					 "shardid, shardlength, groupid, placementid, needsisolatednode) "
 					 "FROM placement_data;");
 
 	/* now add shards to insertShardCommand */
@@ -3903,6 +3905,33 @@ citus_internal_update_none_dist_table_metadata(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * citus_internal_shard_group_set_needsisolatednode is an internal UDF to
+ * set needsisolatednode flag for all the placements within the shard group
+ * that given shard belongs to.
+ */
+Datum
+citus_internal_shard_group_set_needsisolatednode(PG_FUNCTION_ARGS)
+{
+	CheckCitusVersion(ERROR);
+
+	PG_ENSURE_ARGNOTNULL(0, "shard_id");
+	uint64 shardId = PG_GETARG_INT64(0);
+
+	PG_ENSURE_ARGNOTNULL(1, "enabled");
+	bool enabled = PG_GETARG_BOOL(1);
+
+	if (!ShouldSkipMetadataChecks())
+	{
+		EnsureCoordinatorInitiatedOperation();
+	}
+
+	ShardGroupSetNeedsIsolatedNode(shardId, enabled);
+
+	PG_RETURN_VOID();
+}
+
+
 /*
  * SyncNewColocationGroup synchronizes a new pg_dist_colocation entry to a worker.
  */
@@ -4102,6 +4131,22 @@ UpdateNoneDistTableMetadataCommand(Oid relationId, char replicationModel,
 }
 
 
+/*
+ * ShardGroupSetNeedsIsolatedNodeCommand returns a command to call
+ * citus_internal_shard_group_set_needsisolatednode().
+ */
+char *
+ShardGroupSetNeedsIsolatedNodeCommand(uint64 shardId, bool enabled)
+{
+	StringInfo command = makeStringInfo();
+	appendStringInfo(command,
+					 "SELECT pg_catalog.citus_internal_shard_group_set_needsisolatednode(%lu, %s)",
+					 shardId, enabled ? "true" : "false");
+
+	return command->data;
+}
+
+
 /*
  * AddPlacementMetadataCommand returns a command to call
  * citus_internal_add_placement_metadata().