Fix foreign key non colocated (#4840)

DESCRIPTION: Fix a bug where one could create a foreign key between non-colocated tables
    
Due to the scoping and only setting the variable to true when matched, not false when not matched, it was possible for a foreign key to be created on non-colocated tables.

This patch changes the scope and uses an always assign the test result to the variable we prevent users from creating foreign keys between non-colocated tables.

src/backend/distributed/commands/foreign_constraint.c

@@ -134,10 +134,6 @@ ErrorIfUnsupportedForeignConstraint(Relation relation, char distributionMethod,
 	int referencedColumnCount = 0;
 	bool isNull = false;
 	int attrIdx = 0;
-	bool foreignConstraintOnPartitionColumn = false;
-	bool selfReferencingTable = false;
-	bool referencedTableIsAReferenceTable = false;
-	bool referencingColumnsIncludeDistKey = false;
 
 	pgConstraint = heap_open(ConstraintRelationId, AccessShareLock);
 	ScanKeyInit(&scanKey[0], Anum_pg_constraint_conrelid, BTEqualStrategyNumber, F_OIDEQ,
@@ -148,6 +144,10 @@ ErrorIfUnsupportedForeignConstraint(Relation relation, char distributionMethod,
 	heapTuple = systable_getnext(scanDescriptor);
 	while (HeapTupleIsValid(heapTuple))
 	{
+		bool foreignConstraintOnPartitionColumn = false;
+		bool selfReferencingTable = false;
+		bool referencedTableIsAReferenceTable = false;
+		bool referencingColumnsIncludeDistKey = false;
 		Form_pg_constraint constraintForm = (Form_pg_constraint) GETSTRUCT(heapTuple);
 		bool singleReplicatedTable = true;
 
@@ -196,10 +196,8 @@ ErrorIfUnsupportedForeignConstraint(Relation relation, char distributionMethod,
 			 * tables. This is why we make this check under !selfReferencingTable
 			 * and after !IsDistributedTable(referencedTableId).
 			 */
-			if (PartitionMethod(referencedTableId) == DISTRIBUTE_BY_NONE)
-			{
-				referencedTableIsAReferenceTable = true;
-			}
+			referencedTableIsAReferenceTable =
+				(PartitionMethod(referencedTableId) == DISTRIBUTE_BY_NONE);
 
 			/*
 			 * To enforce foreign constraints, tables must be co-located unless a

src/test/regress/expected/foreign_key_to_reference_table.out

@@ -1886,9 +1886,45 @@ ROLLBACK;
 DROP TABLE referenced_table CASCADE;
 NOTICE:  drop cascades to constraint fkey_to_ref on table referencing_table_4
 DROP TABLE referencing_table;
+-- tests specific to an edgecase in citus 8.x where it was possible to create foreign keys
+-- in between colocation groups due to a bug of foreign key to reference tables
+CREATE TABLE t1 (a int PRIMARY KEY, b text);
+CREATE TABLE t2 (a bigint PRIMARY KEY, b text);
+CREATE TABLE r1 (a int PRIMARY KEY, b text);
+SELECT create_distributed_table('t1', 'a');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+SELECT create_distributed_table('t2', 'a');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+SELECT create_reference_table('r1');
+ create_reference_table 
+------------------------
+ 
+(1 row)
+
+-- this always fails as it should be
+ALTER TABLE t1 ADD CONSTRAINT c1 FOREIGN KEY (a) REFERENCES t2(a);
+ERROR:  cannot create foreign key constraint since relations are not colocated or not referencing a reference table
+DETAIL:  A distributed table can only have foreign keys if it is referencing another colocated hash distributed table or a reference table
+-- after we create a foreign key to the reference table, that has a lower order by name,
+-- we would have been able to create a FK to non-colocated tables
+ALTER TABLE t1 ADD CONSTRAINT c2 FOREIGN KEY (a) REFERENCES r1(a);
+ALTER TABLE t1 ADD CONSTRAINT c3 FOREIGN KEY (a) REFERENCES t2(a);
+ERROR:  cannot create foreign key constraint since relations are not colocated or not referencing a reference table
+DETAIL:  A distributed table can only have foreign keys if it is referencing another colocated hash distributed table or a reference table
 DROP SCHEMA fkey_reference_table CASCADE;
-NOTICE:  drop cascades to 3 other objects
+NOTICE:  drop cascades to 6 other objects
 DETAIL:  drop cascades to type foreign_details
 drop cascades to view table_fkeys_in_workers
 drop cascades to type composite
+drop cascades to table t1
+drop cascades to table t2
+drop cascades to table r1
 SET search_path TO DEFAULT;

src/test/regress/sql/foreign_key_to_reference_table.sql

@@ -951,5 +951,24 @@ ROLLBACK;
 DROP TABLE referenced_table CASCADE;
 DROP TABLE referencing_table;
 
+-- tests specific to an edgecase in citus 8.x where it was possible to create foreign keys
+-- in between colocation groups due to a bug of foreign key to reference tables
+CREATE TABLE t1 (a int PRIMARY KEY, b text);
+CREATE TABLE t2 (a bigint PRIMARY KEY, b text);
+CREATE TABLE r1 (a int PRIMARY KEY, b text);
+
+SELECT create_distributed_table('t1', 'a');
+SELECT create_distributed_table('t2', 'a');
+SELECT create_reference_table('r1');
+
+-- this always fails as it should be
+ALTER TABLE t1 ADD CONSTRAINT c1 FOREIGN KEY (a) REFERENCES t2(a);
+
+-- after we create a foreign key to the reference table, that has a lower order by name,
+-- we would have been able to create a FK to non-colocated tables
+ALTER TABLE t1 ADD CONSTRAINT c2 FOREIGN KEY (a) REFERENCES r1(a);
+ALTER TABLE t1 ADD CONSTRAINT c3 FOREIGN KEY (a) REFERENCES t2(a);
+
+
 DROP SCHEMA fkey_reference_table CASCADE;
 SET search_path TO DEFAULT;