Cascade Citus Local Tables via fkeys

src/backend/distributed/commands/cascade_table_operation_for_connected_relations.c

@@ -74,32 +74,16 @@ CascadeOperationForConnectedRelations(Oid relationId, LOCKMODE lockMode,
 
 	LockRelationsWithLockMode(fKeyConnectedRelationIdList, lockMode);
 
-	/*
-	 * Before removing any partition relations, we should error out here if any
-	 * of connected relations is a partition table involved in a foreign key that
-	 * is not inherited from its parent table.
-	 * We should handle this case here as we remove partition relations in this
-	 * function	before ExecuteCascadeOperationForRelationIdList.
-	 */
-	ErrorIfAnyPartitionRelationInvolvedInNonInheritedFKey(fKeyConnectedRelationIdList);
-
-	/*
-	 * We shouldn't cascade through foreign keys on partition tables as citus
-	 * table functions already have their own logics to handle partition relations.
-	 */
-	List *nonPartitionRelationIdList =
-		RemovePartitionRelationIds(fKeyConnectedRelationIdList);
-
 	/*
 	 * Our foreign key subgraph can have distributed tables which might already
 	 * be modified in current transaction. So switch to sequential execution
 	 * before executing any ddl's to prevent erroring out later in this function.
 	 */
-	EnsureSequentialModeForCitusTableCascadeFunction(nonPartitionRelationIdList);
+	EnsureSequentialModeForCitusTableCascadeFunction(fKeyConnectedRelationIdList);
 
 	/* store foreign key creation commands before dropping them */
 	List *fKeyCreationCommands =
-		GetFKeyCreationCommandsForRelationIdList(nonPartitionRelationIdList);
+		GetFKeyCreationCommandsForRelationIdList(fKeyConnectedRelationIdList);
 
 	/*
 	 * Note that here we only drop referencing foreign keys for each relation.
@@ -107,8 +91,8 @@ CascadeOperationForConnectedRelations(Oid relationId, LOCKMODE lockMode,
 	 * relations' referencing foreign keys.
 	 */
 	int fKeyFlags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
-	DropRelationIdListForeignKeys(nonPartitionRelationIdList, fKeyFlags);
+	DropRelationIdListForeignKeys(fKeyConnectedRelationIdList, fKeyFlags);
-	ExecuteCascadeOperationForRelationIdList(nonPartitionRelationIdList,
+	ExecuteCascadeOperationForRelationIdList(fKeyConnectedRelationIdList,
 											 cascadeOperationType);
 
 	/* now recreate foreign keys on tables */

src/backend/distributed/commands/citus_add_local_table_to_metadata.c

@@ -48,7 +48,8 @@
 static void citus_add_local_table_to_metadata_internal(Oid relationId,
 													   bool cascadeViaForeignKeys);
 static void CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId,
-														  Oid shardRelationId);
+														  Oid shardRelationId,
+														  bool cascadeViaForeignKeys);
 static void ErrorIfUnsupportedCreateCitusLocalTable(Relation relation);
 static void ErrorIfUnsupportedCitusLocalTableKind(Oid relationId);
 static void ErrorIfUnsupportedCitusLocalColumnDefinition(Relation relation);
@@ -266,24 +267,6 @@ CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
 		 */
 		return;
 	}
-	else if (tableHasExternalForeignKeys)
-	{
-		/*
-		 * We do not allow creating citus local table if the table is involved in a
-		 * foreign key relationship with "any other table". Note that we allow self
-		 * references.
-		 */
-		char *qualifiedRelationName = generate_qualified_relation_name(relationId);
-		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						errmsg("relation %s is involved in a foreign key "
-							   "relationship with another table", qualifiedRelationName),
-						errhint("Use cascade_via_foreign_keys option to add "
-								"all the relations involved in a foreign key "
-								"relationship with %s to citus metadata by "
-								"executing SELECT citus_add_local_table_to_metadata($$%s$$, "
-								"cascade_via_foreign_keys=>true)",
-								qualifiedRelationName, qualifiedRelationName)));
-	}
 
 	ObjectAddress tableAddress = { 0 };
 	ObjectAddressSet(tableAddress, RelationRelationId, relationId);
@@ -353,7 +336,8 @@ CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
 	 * If it's a partitioned table, we need to create Citus Local Tables from its
 	 * partitions too.
 	 */
-	CreateChildLocalTablesIfRelationIsPartitioned(shellRelationId, shardRelationId);
+	CreateChildLocalTablesIfRelationIsPartitioned(shellRelationId, shardRelationId,
+												  cascadeViaForeignKeys);
 }
 
 
@@ -370,7 +354,8 @@ CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
  * ATTACH to establish the correct relationship.
  */
 static void
-CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId, Oid shardRelationId)
+CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId, Oid shardRelationId,
+											  bool cascadeViaForeignKeys)
 {
 	/* if this table is partitioned table, add its partitions to metadata too */
 	if (PartitionedTable(shardRelationId))
@@ -392,7 +377,7 @@ CreateChildLocalTablesIfRelationIsPartitioned(Oid shellRelationId, Oid shardRela
 																 qualifiedShellRelationName);
 
 			/* here we call CreateCitusLocalTable for all partitions */
-			CreateCitusLocalTable(partitionRelationId, false);
+			CreateCitusLocalTable(partitionRelationId, cascadeViaForeignKeys);
 
 			/*
 			 * CreateCitusLocalTable has created a Citus Local Table for the partition.

src/backend/distributed/commands/table.c

@@ -504,7 +504,7 @@ PreprocessAlterTableStmtAttachPartition(AlterTableStmt *alterTableStatement,
 			}
 
 			/* if parent of this table is distributed, distribute this table too */
-			if (IsCitusTable(relationId) &&
+			if (IsCitusTableType(relationId, DISTRIBUTED_TABLE) &&
 				!IsCitusTable(partitionRelationId))
 			{
 				Var *distributionColumn = DistPartitionKeyOrError(relationId);