CopyShardForeignConstraintCommandList API change for grouped constraints

src/backend/distributed/master/master_repair_shards.c

@@ -348,8 +348,31 @@ CopyShardCommandList(ShardInterval *shardInterval,
 List *
 CopyShardForeignConstraintCommandList(ShardInterval *shardInterval)
 {
-	List *copyShardForeignConstraintCommandList = NIL;
+	List *colocatedShardForeignConstraintCommandList = NIL;
+	List *referenceTableForeignConstraintList = NIL;
 
+	CopyShardForeignConstraintCommandListGrouped(shardInterval,
+												 &
+												 colocatedShardForeignConstraintCommandList,
+												 &referenceTableForeignConstraintList);
+
+	return list_concat(colocatedShardForeignConstraintCommandList,
+					   referenceTableForeignConstraintList);
+}
+
+
+/*
+ * CopyShardForeignConstraintCommandListGrouped generates command lists
+ * to create foreign constraints existing in source shard after copying it to other
+ * node in separate groups for foreign constraints in between hash distributed tables
+ * and from a hash distributed to reference tables.
+ */
+void
+CopyShardForeignConstraintCommandListGrouped(ShardInterval *shardInterval,
+											 List **
+											 colocatedShardForeignConstraintCommandList,
+											 List **referenceTableForeignConstraintList)
+{
 	Oid schemaId = get_rel_namespace(shardInterval->relationId);
 	char *schemaName = get_namespace_name(schemaId);
 	char *escapedSchemaName = quote_literal_cstr(schemaName);
@@ -364,6 +387,9 @@ CopyShardForeignConstraintCommandList(ShardInterval *shardInterval)
 		shardIndex = ShardIndex(shardInterval);
 	}
 
+	*colocatedShardForeignConstraintCommandList = NIL;
+	*referenceTableForeignConstraintList = NIL;
+
 	foreach(commandCell, commandList)
 	{
 		char *command = (char *) lfirst(commandCell);
@@ -374,6 +400,7 @@ CopyShardForeignConstraintCommandList(ShardInterval *shardInterval)
 		char *referencedSchemaName = NULL;
 		char *escapedReferencedSchemaName = NULL;
 		uint64 referencedShardId = INVALID_SHARD_ID;
+		bool colocatedForeignKey = false;
 
 		StringInfo applyForeignConstraintCommand = makeStringInfo();
 
@@ -389,21 +416,43 @@ CopyShardForeignConstraintCommandList(ShardInterval *shardInterval)
 		referencedSchemaId = get_rel_namespace(referencedRelationId);
 		referencedSchemaName = get_namespace_name(referencedSchemaId);
 		escapedReferencedSchemaName = quote_literal_cstr(referencedSchemaName);
-		referencedShardId = ColocatedShardIdInRelation(referencedRelationId, shardIndex);
+
+		if (PartitionMethod(referencedRelationId) == DISTRIBUTE_BY_NONE)
+		{
+			List *shardList = LoadShardList(referencedRelationId);
+			uint64 *shardIdPointer = (uint64 *) linitial(shardList);
+
+			referencedShardId = (*shardIdPointer);
+		}
+		else
+		{
+			referencedShardId = ColocatedShardIdInRelation(referencedRelationId,
+														   shardIndex);
+
+			colocatedForeignKey = true;
+		}
 
 		appendStringInfo(applyForeignConstraintCommand,
 						 WORKER_APPLY_INTER_SHARD_DDL_COMMAND, shardInterval->shardId,
 						 escapedSchemaName, referencedShardId,
 						 escapedReferencedSchemaName, escapedCommand);
 
-		copyShardForeignConstraintCommandList = lappend(
+		if (colocatedForeignKey)
-			copyShardForeignConstraintCommandList,
+		{
-			applyForeignConstraintCommand->data);
+			*colocatedShardForeignConstraintCommandList = lappend(
+				*colocatedShardForeignConstraintCommandList,
+				applyForeignConstraintCommand->data);
+		}
+		else
+		{
+			*referenceTableForeignConstraintList = lappend(
+				*referenceTableForeignConstraintList,
+				applyForeignConstraintCommand->data);
+		}
 	}
-
-	return copyShardForeignConstraintCommandList;
 }
 
+
 /*
  * ConstuctQualifiedShardName creates the fully qualified name string of the
  * given shard in <schema>.<table_name>_<shard_id> format.

src/include/distributed/master_protocol.h

@@ -164,6 +164,11 @@ extern Datum master_copy_shard_placement(PG_FUNCTION_ARGS);
 extern List * CopyShardCommandList(ShardInterval *shardInterval, char *sourceNodeName,
 								   int32 sourceNodePort);
 extern List * CopyShardForeignConstraintCommandList(ShardInterval *shardInterval);
+extern void CopyShardForeignConstraintCommandListGrouped(ShardInterval *shardInterval,
+														 List **
+														 colocatedShardForeignConstraintCommandList,
+														 List **
+														 referenceTableForeignConstraintList);
 extern ShardPlacement * SearchShardPlacementInList(List *shardPlacementList,
 												   char *nodeName, uint32 nodePort,
 												   bool missingOk);