Refactor foreign key functions to use table types (#4424)

* Reuses extractReferencing/Referenced variables

* Refactors GetForeignKeyOids function to check table types

* Converts flags to inclusive

src/backend/distributed/commands/foreign_constraint.c

@@ -68,10 +68,7 @@ static List * GetForeignKeyIdsForColumn(char *columnName, Oid relationId,
 										int searchForeignKeyColumnFlags);
 static List * GetForeignConstraintCommandsInternal(Oid relationId, int flags);
 static Oid get_relation_constraint_oid_compat(HeapTuple heapTuple);
-static List * GetForeignKeyOidsToCitusLocalTables(Oid relationId);
+static bool IsTableTypeIncluded(Oid relationId, int flags);
-static List * GetForeignKeyOidsToReferenceTables(Oid relationId);
-static List * FilterFKeyOidListByReferencedTableType(List *foreignKeyOidList,
-													 CitusTableType citusTableType);
 
 /*
  * ConstraintIsAForeignKeyToReferenceTable checks if the given constraint is a
@@ -80,7 +77,8 @@ static List * FilterFKeyOidListByReferencedTableType(List *foreignKeyOidList,
 bool
 ConstraintIsAForeignKeyToReferenceTable(char *inputConstaintName, Oid relationId)
 {
-	List *foreignKeyOids = GetForeignKeyOidsToReferenceTables(relationId);
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_REFERENCE_TABLES;
+	List *foreignKeyOids = GetForeignKeyOids(relationId, flags);
 
 	Oid foreignKeyOid = FindForeignKeyOidWithName(foreignKeyOids, inputConstaintName);
 
@@ -122,7 +120,7 @@ ErrorIfUnsupportedForeignConstraintExists(Relation relation, char referencingDis
 {
 	Oid referencingTableId = relation->rd_id;
 
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
 	List *foreignKeyOids = GetForeignKeyOids(referencingTableId, flags);
 
 	Oid foreignKeyOid = InvalidOid;
@@ -595,7 +593,7 @@ GetForeignKeyIdsForColumn(char *columnName, Oid relationId,
 List *
 GetReferencingForeignConstaintCommands(Oid relationId)
 {
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
 	return GetForeignConstraintCommandsInternal(relationId, flags);
 }
 
@@ -661,23 +659,9 @@ get_relation_constraint_oid_compat(HeapTuple heapTuple)
 bool
 HasForeignKeyToCitusLocalTable(Oid relationId)
 {
-	List *foreignKeyOidList = GetForeignKeyOidsToCitusLocalTables(relationId);
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_CITUS_LOCAL_TABLES;
-	return list_length(foreignKeyOidList) > 0;
-}
-
-
-/*
- * GetForeignKeyOidsToCitusLocalTables returns list of OIDs for the foreign key
- * constraints on the given relationId that are referencing to citus local tables.
- */
-static List *
-GetForeignKeyOidsToCitusLocalTables(Oid relationId)
-{
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
 	List *foreignKeyOidList = GetForeignKeyOids(relationId, flags);
-	List *fKeyOidsToCitusLocalTables =
+	return list_length(foreignKeyOidList) > 0;
-		FilterFKeyOidListByReferencedTableType(foreignKeyOidList, CITUS_LOCAL_TABLE);
-	return fKeyOidsToCitusLocalTables;
 }
 
 
@@ -689,57 +673,13 @@ GetForeignKeyOidsToCitusLocalTables(Oid relationId)
 bool
 HasForeignKeyToReferenceTable(Oid relationId)
 {
-	List *foreignKeyOids = GetForeignKeyOidsToReferenceTables(relationId);
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_REFERENCE_TABLES;
+	List *foreignKeyOids = GetForeignKeyOids(relationId, flags);
 
 	return list_length(foreignKeyOids) > 0;
 }
 
 
-/*
- * GetForeignKeyOidsToReferenceTables returns list of OIDs for the foreign key
- * constraints on the given relationId that are referencing to reference tables.
- */
-static List *
-GetForeignKeyOidsToReferenceTables(Oid relationId)
-{
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
-	List *foreignKeyOidList = GetForeignKeyOids(relationId, flags);
-	List *fKeyOidsToReferenceTables =
-		FilterFKeyOidListByReferencedTableType(foreignKeyOidList, REFERENCE_TABLE);
-	return fKeyOidsToReferenceTables;
-}
-
-
-/*
- * FilterFKeyOidListByReferencedTableType takes a list of foreign key OIDs and
- * CitusTableType to filter the foreign key OIDs that CitusTableType matches
- * referenced relation's type.
- */
-static List *
-FilterFKeyOidListByReferencedTableType(List *foreignKeyOidList,
-									   CitusTableType citusTableType)
-{
-	List *filteredFKeyOidList = NIL;
-
-	Oid foreignKeyOid = InvalidOid;
-	foreach_oid(foreignKeyOid, foreignKeyOidList)
-	{
-		HeapTuple heapTuple = SearchSysCache1(CONSTROID, ObjectIdGetDatum(foreignKeyOid));
-		Form_pg_constraint constraintForm = (Form_pg_constraint) GETSTRUCT(heapTuple);
-
-		Oid referencedTableOid = constraintForm->confrelid;
-		if (IsCitusTableType(referencedTableOid, citusTableType))
-		{
-			filteredFKeyOidList = lappend_oid(filteredFKeyOidList, foreignKeyOid);
-		}
-
-		ReleaseSysCache(heapTuple);
-	}
-
-	return filteredFKeyOidList;
-}
-
-
 /*
  * TableReferenced function checks whether given table is referenced by another table
  * via foreign constraints. If it is referenced, this function returns true.
@@ -747,7 +687,7 @@ FilterFKeyOidListByReferencedTableType(List *foreignKeyOidList,
 bool
 TableReferenced(Oid relationId)
 {
-	int flags = INCLUDE_REFERENCED_CONSTRAINTS;
+	int flags = INCLUDE_REFERENCED_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
 	List *foreignKeyOids = GetForeignKeyOids(relationId, flags);
 
 	return list_length(foreignKeyOids) > 0;
@@ -793,7 +733,7 @@ HeapTupleOfForeignConstraintIncludesColumn(HeapTuple heapTuple, Oid relationId,
 bool
 TableReferencing(Oid relationId)
 {
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
 	List *foreignKeyOids = GetForeignKeyOids(relationId, flags);
 
 	return list_length(foreignKeyOids) > 0;
@@ -821,7 +761,7 @@ ConstraintIsAForeignKey(char *inputConstaintName, Oid relationId)
 Oid
 GetForeignKeyOidByName(char *inputConstaintName, Oid relationId)
 {
-	int flags = INCLUDE_REFERENCING_CONSTRAINTS;
+	int flags = INCLUDE_REFERENCING_CONSTRAINTS | INCLUDE_ALL_TABLE_TYPES;
 	List *foreignKeyOids = GetForeignKeyOids(relationId, flags);
 
 	Oid foreignKeyId = FindForeignKeyOidWithName(foreignKeyOids, inputConstaintName);
@@ -862,10 +802,12 @@ FindForeignKeyOidWithName(List *foreignKeyOids, const char *inputConstraintName)
 void
 ErrorIfTableHasExternalForeignKeys(Oid relationId)
 {
-	int flags = (INCLUDE_REFERENCING_CONSTRAINTS | EXCLUDE_SELF_REFERENCES);
+	int flags = (INCLUDE_REFERENCING_CONSTRAINTS | EXCLUDE_SELF_REFERENCES |
+				 INCLUDE_ALL_TABLE_TYPES);
 	List *foreignKeyIdsTableReferencing = GetForeignKeyOids(relationId, flags);
 
-	flags = (INCLUDE_REFERENCED_CONSTRAINTS | EXCLUDE_SELF_REFERENCES);
+	flags = (INCLUDE_REFERENCED_CONSTRAINTS | EXCLUDE_SELF_REFERENCES |
+			 INCLUDE_ALL_TABLE_TYPES);
 	List *foreignKeyIdsTableReferenced = GetForeignKeyOids(relationId, flags);
 
 	List *foreignKeysWithOtherTables = list_concat(foreignKeyIdsTableReferencing,
@@ -897,10 +839,8 @@ GetForeignKeyOids(Oid relationId, int flags)
 {
 	AttrNumber pgConstraintTargetAttrNumber = InvalidAttrNumber;
 
-	bool extractReferencing PG_USED_FOR_ASSERTS_ONLY =
+	bool extractReferencing = (flags & INCLUDE_REFERENCING_CONSTRAINTS);
-		(flags & INCLUDE_REFERENCING_CONSTRAINTS);
+	bool extractReferenced = (flags & INCLUDE_REFERENCED_CONSTRAINTS);
-	bool extractReferenced PG_USED_FOR_ASSERTS_ONLY =
-		(flags & INCLUDE_REFERENCED_CONSTRAINTS);
 
 	/*
 	 * Only one of them should be passed at a time since the way we scan
@@ -913,14 +853,14 @@ GetForeignKeyOids(Oid relationId, int flags)
 	bool useIndex = false;
 	Oid indexOid = InvalidOid;
 
-	if (flags & INCLUDE_REFERENCING_CONSTRAINTS)
+	if (extractReferencing)
 	{
 		pgConstraintTargetAttrNumber = Anum_pg_constraint_conrelid;
 
 		useIndex = true;
 		indexOid = ConstraintRelidTypidNameIndexId;
 	}
-	else if (flags & INCLUDE_REFERENCED_CONSTRAINTS)
+	else if (extractReferenced)
 	{
 		pgConstraintTargetAttrNumber = Anum_pg_constraint_confrelid;
 	}
@@ -970,6 +910,22 @@ GetForeignKeyOids(Oid relationId, int flags)
 			continue;
 		}
 
+		Oid otherTableId = InvalidOid;
+		if (extractReferencing)
+		{
+			otherTableId = constraintForm->confrelid;
+		}
+		else if (extractReferenced)
+		{
+			otherTableId = constraintForm->conrelid;
+		}
+
+		if (!IsTableTypeIncluded(otherTableId, flags))
+		{
+			heapTuple = systable_getnext(scanDescriptor);
+			continue;
+		}
+
 		foreignKeyOids = lappend_oid(foreignKeyOids, constraintId);
 
 		heapTuple = systable_getnext(scanDescriptor);
@@ -1010,3 +966,30 @@ GetReferencedTableId(Oid foreignKeyId)
 
 	return referencedTableId;
 }
+
+
+/*
+ * IsTableTypeIncluded returns true if type of the table with relationId (distributed,
+ * reference, Citus local or Postgres local) is included in the flags, false if not
+ */
+static bool
+IsTableTypeIncluded(Oid relationId, int flags)
+{
+	if (!IsCitusTable(relationId))
+	{
+		return (flags & INCLUDE_LOCAL_TABLES) != 0;
+	}
+	else if (IsCitusTableType(relationId, DISTRIBUTED_TABLE))
+	{
+		return (flags & INCLUDE_DISTRIBUTED_TABLES) != 0;
+	}
+	else if (IsCitusTableType(relationId, REFERENCE_TABLE))
+	{
+		return (flags & INCLUDE_REFERENCE_TABLES) != 0;
+	}
+	else if (IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+	{
+		return (flags & INCLUDE_CITUS_LOCAL_TABLES) != 0;
+	}
+	return false;
+}

src/include/distributed/commands.h

@@ -61,7 +61,24 @@ typedef enum ExtractForeignKeyConstraintsMode
 	INCLUDE_REFERENCED_CONSTRAINTS = 1 << 1,
 
 	/* exclude the self-referencing foreign keys */
-	EXCLUDE_SELF_REFERENCES = 1 << 2
+	EXCLUDE_SELF_REFERENCES = 1 << 2,
+
+	/* any combination of the 4 flags below is supported */
+	/* include foreign keys when the other table is a distributed table*/
+	INCLUDE_DISTRIBUTED_TABLES = 1 << 3,
+
+	/* include foreign keys when the other table is a reference table*/
+	INCLUDE_REFERENCE_TABLES = 1 << 4,
+
+	/* include foreign keys when the other table is a citus local table*/
+	INCLUDE_CITUS_LOCAL_TABLES = 1 << 5,
+
+	/* include foreign keys when the other table is a Postgres local table*/
+	INCLUDE_LOCAL_TABLES = 1 << 6,
+
+	/* include foreign keys regardless of the other table's type */
+	INCLUDE_ALL_TABLE_TYPES = INCLUDE_DISTRIBUTED_TABLES | INCLUDE_REFERENCE_TABLES |
+							  INCLUDE_CITUS_LOCAL_TABLES | INCLUDE_LOCAL_TABLES
 } ExtractForeignKeyConstraintMode;