Simplify foreign key graph invalidation

Object access hook is much simpler and safer
than trying to find all cases manually.

src/backend/distributed/commands/foreign_constraint.c

@@ -498,21 +498,6 @@ ForeignConstraintFindDistKeys(HeapTuple pgConstraintTuple,
 }
 
 
-/*
- * ColumnAppearsInForeignKey returns true if there is a foreign key constraint
- * from/to given column. False otherwise.
- */
-bool
-ColumnAppearsInForeignKey(char *columnName, Oid relationId)
-{
-	int searchForeignKeyColumnFlags = SEARCH_REFERENCING_RELATION |
-									  SEARCH_REFERENCED_RELATION;
-	List *foreignKeysColumnAppeared =
-		GetForeignKeyIdsForColumn(columnName, relationId, searchForeignKeyColumnFlags);
-	return list_length(foreignKeysColumnAppeared) > 0;
-}
-
-
 /*
  * ColumnAppearsInForeignKeyToReferenceTable checks if there is a foreign key
  * constraint from/to any reference table on the given column.
@@ -820,23 +805,6 @@ TableReferencing(Oid relationId)
 }
 
 
-/*
- * ConstraintIsAUniquenessConstraint is a wrapper around ConstraintWithNameIsOfType
- * that returns true if given constraint name identifies a uniqueness constraint, i.e:
- *   - primary key constraint, or
- *   - unique constraint
- */
-bool
-ConstraintIsAUniquenessConstraint(char *inputConstaintName, Oid relationId)
-{
-	bool isUniqueConstraint = ConstraintWithNameIsOfType(inputConstaintName, relationId,
-														 CONSTRAINT_UNIQUE);
-	bool isPrimaryConstraint = ConstraintWithNameIsOfType(inputConstaintName, relationId,
-														  CONSTRAINT_PRIMARY);
-	return isUniqueConstraint || isPrimaryConstraint;
-}
-
-
 /*
  * ConstraintIsAForeignKey is a wrapper around ConstraintWithNameIsOfType that returns true
  * if given constraint name identifies a foreign key constraint.
@@ -888,6 +856,29 @@ ConstraintWithIdIsOfType(Oid constraintId, char targetConstraintType)
 }
 
 
+/*
+ * ConstraintOnRelationId returns the relationId that the constraint is
+ * defined on.
+ */
+Oid
+ConstraintOnRelationId(Oid constraintId)
+{
+	HeapTuple heapTuple = SearchSysCache1(CONSTROID, ObjectIdGetDatum(constraintId));
+	if (!HeapTupleIsValid(heapTuple))
+	{
+		/* no such constraint */
+		return InvalidOid;
+	}
+
+	Form_pg_constraint constraintForm = (Form_pg_constraint) GETSTRUCT(heapTuple);
+	Oid conrelid = constraintForm->conrelid;
+
+	ReleaseSysCache(heapTuple);
+
+	return conrelid;
+}
+
+
 /*
  * FindForeignKeyOidWithName searches the foreign key constraint with
  * inputConstraintName in the given list of foreign key constraint OIDs.

src/backend/distributed/commands/index.c

@@ -754,11 +754,6 @@ PreprocessDropIndexStmt(Node *node, const char *dropIndexCommand,
 
 		ErrorIfUnsupportedDropIndexStmt(dropIndexStatement);
 
-		if (AnyForeignKeyDependsOnIndex(distributedIndexId))
-		{
-			MarkInvalidateForeignKeyGraph();
-		}
-
 		ObjectAddressSet(ddlJob->targetObjectAddress, RelationRelationId,
 						 distributedRelationId);
 

src/backend/distributed/commands/schema.c

@@ -112,7 +112,6 @@ PreprocessDropSchemaStmt(Node *node, const char *queryString,
 	{
 		if (SchemaHasDistributedTableWithFKey(strVal(schemaVal)))
 		{
-			MarkInvalidateForeignKeyGraph();
 			break;
 		}
 	}

src/backend/distributed/commands/table.c

@@ -90,7 +90,6 @@ static List * GetRangeVarListFromFKeyConstraintList(List *fKeyConstraintList);
 static List * GetRelationIdListFromRangeVarList(List *rangeVarList, LOCKMODE lockmode,
 												bool missingOk);
 static bool AlterTableCommandTypeIsTrigger(AlterTableType alterTableType);
-static bool AlterTableDropsForeignKey(AlterTableStmt *alterTableStatement);
 static void ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement);
 static List * InterShardDDLTaskList(Oid leftRelationId, Oid rightRelationId,
 									const char *commandString);
@@ -163,12 +162,6 @@ PreprocessDropTableStmt(Node *node, const char *queryString,
 			LockColocationId(cacheEntry->colocationId, ShareLock);
 		}
 
-		/* invalidate foreign key cache if the table involved in any foreign key */
-		if ((TableReferenced(relationId) || TableReferencing(relationId)))
-		{
-			MarkInvalidateForeignKeyGraph();
-		}
-
 		/* we're only interested in partitioned and mx tables */
 		if (!ShouldSyncTableMetadata(relationId) || !PartitionedTable(relationId))
 		{
@@ -796,18 +789,6 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 		return PreprocessAlterTableStmt(node, alterTableCommand, processUtilityContext);
 	}
 
-	if (AlterTableDropsForeignKey(alterTableStatement))
-	{
-		/*
-		 * The foreign key graph keeps track of the foreign keys including local tables.
-		 * So, even if a foreign key on a local table is dropped, we should invalidate
-		 * the graph so that the next commands can see the graph up-to-date.
-		 * We are aware that utility hook would still invalidate foreign key graph
-		 * even when command fails, but currently we are ok with that.
-		 */
-		MarkInvalidateForeignKeyGraph();
-	}
-
 	bool referencingIsLocalTable = !IsCitusTable(leftRelationId);
 	if (referencingIsLocalTable)
 	{
@@ -1663,99 +1644,6 @@ ConstrTypeUsesIndex(ConstrType constrType)
 }
 
 
-/*
- * AlterTableDropsForeignKey returns true if the given AlterTableStmt drops
- * a foreign key. False otherwise.
- */
-static bool
-AlterTableDropsForeignKey(AlterTableStmt *alterTableStatement)
-{
-	LOCKMODE lockmode = AlterTableGetLockLevel(alterTableStatement->cmds);
-	Oid relationId = AlterTableLookupRelation(alterTableStatement, lockmode);
-
-	AlterTableCmd *command = NULL;
-	foreach_ptr(command, alterTableStatement->cmds)
-	{
-		AlterTableType alterTableType = command->subtype;
-
-		if (alterTableType == AT_DropColumn)
-		{
-			char *columnName = command->name;
-			if (ColumnAppearsInForeignKey(columnName, relationId))
-			{
-				/* dropping a column in the either side of the fkey will drop the fkey */
-				return true;
-			}
-		}
-
-		/*
-		 * In order to drop the foreign key, other than DROP COLUMN, the command must be
-		 * DROP CONSTRAINT command.
-		 */
-		if (alterTableType != AT_DropConstraint)
-		{
-			continue;
-		}
-
-		char *constraintName = command->name;
-		if (ConstraintIsAForeignKey(constraintName, relationId))
-		{
-			return true;
-		}
-		else if (ConstraintIsAUniquenessConstraint(constraintName, relationId))
-		{
-			/*
-			 * If the uniqueness constraint of the column that the foreign key depends on
-			 * is getting dropped, then the foreign key will also be dropped.
-			 */
-			bool missingOk = false;
-			Oid uniquenessConstraintId =
-				get_relation_constraint_oid(relationId, constraintName, missingOk);
-			Oid indexId = get_constraint_index(uniquenessConstraintId);
-			if (AnyForeignKeyDependsOnIndex(indexId))
-			{
-				return true;
-			}
-		}
-	}
-
-	return false;
-}
-
-
-/*
- * AnyForeignKeyDependsOnIndex scans pg_depend and returns true if given index
- * is valid and any foreign key depends on it.
- */
-bool
-AnyForeignKeyDependsOnIndex(Oid indexId)
-{
-	Oid dependentObjectClassId = RelationRelationId;
-	Oid dependentObjectId = indexId;
-	List *dependencyTupleList =
-		GetPgDependTuplesForDependingObjects(dependentObjectClassId, dependentObjectId);
-
-	HeapTuple dependencyTuple = NULL;
-	foreach_ptr(dependencyTuple, dependencyTupleList)
-	{
-		Form_pg_depend dependencyForm = (Form_pg_depend) GETSTRUCT(dependencyTuple);
-		Oid dependingClassId = dependencyForm->classid;
-		if (dependingClassId != ConstraintRelationId)
-		{
-			continue;
-		}
-
-		Oid dependingObjectId = dependencyForm->objid;
-		if (ConstraintWithIdIsOfType(dependingObjectId, CONSTRAINT_FOREIGN))
-		{
-			return true;
-		}
-	}
-
-	return false;
-}
-
-
 /*
  * PreprocessAlterTableStmt issues a warning.
  * ALTER TABLE ALL IN TABLESPACE statements have their node type as
@@ -2974,7 +2862,7 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 			case AT_ForceRowSecurity:
 			case AT_NoForceRowSecurity:
 			case AT_ValidateConstraint:
-			case AT_DropConstraint: /* we do the check for invalidation in AlterTableDropsForeignKey */
+			case AT_DropConstraint:
 #if PG_VERSION_NUM >= PG_VERSION_14
 			case AT_SetCompression:
 #endif

src/backend/distributed/commands/utility_hook.c

@@ -1514,16 +1514,6 @@ static void
 PostStandardProcessUtility(Node *parsetree)
 {
 	DecrementUtilityHookCountersIfNecessary(parsetree);
-
-	/*
-	 * Re-forming the foreign key graph relies on the command being executed
-	 * on the local table first. However, in order to decide whether the
-	 * command leads to an invalidation, we need to check before the command
-	 * is being executed since we read pg_constraint table. Thus, we maintain a
-	 * local flag and do the invalidation after multi_ProcessUtility,
-	 * before ExecuteDistributedDDLJob().
-	 */
-	InvalidateForeignKeyGraphForDDL();
 }
 
 
@@ -1558,22 +1548,6 @@ MarkInvalidateForeignKeyGraph()
 }
 
 
-/*
- * InvalidateForeignKeyGraphForDDL simply keeps track of whether
- * the foreign key graph should be invalidated due to a DDL.
- */
-void
-InvalidateForeignKeyGraphForDDL(void)
-{
-	if (shouldInvalidateForeignKeyGraph)
-	{
-		InvalidateForeignKeyGraph();
-
-		shouldInvalidateForeignKeyGraph = false;
-	}
-}
-
-
 /*
  * DDLTaskList builds a list of tasks to execute a DDL command on a
  * given list of shards.

src/backend/distributed/shared_library_init.c

@@ -2699,6 +2699,8 @@ IsSuperuser(char *roleName)
 }
 
 
+#include "catalog/pg_constraint.h"
+
 /*
  * CitusObjectAccessHook is called when an object is created.
  *
@@ -2722,4 +2724,17 @@ CitusObjectAccessHook(ObjectAccessType access, Oid classId, Oid objectId, int su
 		 * regardless if it's citus being created */
 		SetCreateCitusTransactionLevel(GetCurrentTransactionNestLevel());
 	}
+	else if (access == OAT_DROP && classId == ConstraintRelationId &&
+			 ConstraintWithIdIsOfType(objectId, CONSTRAINT_FOREIGN))
+	{
+		Oid relationId = ConstraintOnRelationId(objectId);
+		if (OidIsValid(relationId) && IsCitusTable(relationId))
+		{
+			/*
+			 * User drops foreign constraint on Citus table,
+			 * invalidate fkey graph.
+			 */
+			InvalidateForeignKeyGraph();
+		}
+	}
 }

src/include/distributed/commands.h

@@ -254,7 +254,6 @@ extern void ErrorIfUnsupportedForeignConstraintExists(Relation relation,
 													  uint32 colocationId);
 extern void ErrorOutForFKeyBetweenPostgresAndCitusLocalTable(Oid localTableId);
 extern bool ColumnReferencedByAnyForeignKey(char *columnName, Oid relationId);
-extern bool ColumnAppearsInForeignKey(char *columnName, Oid relationId);
 extern bool ColumnAppearsInForeignKeyToReferenceTable(char *columnName, Oid
 													  relationId);
 extern List * GetReferencingForeignConstaintCommands(Oid relationOid);
@@ -268,11 +267,11 @@ extern bool HasForeignKeyToReferenceTable(Oid relationOid);
 extern List * GetForeignKeysFromLocalTables(Oid relationId);
 extern bool TableReferenced(Oid relationOid);
 extern bool TableReferencing(Oid relationOid);
-extern bool ConstraintIsAUniquenessConstraint(char *inputConstaintName, Oid relationId);
 extern bool ConstraintIsAForeignKey(char *inputConstaintName, Oid relationOid);
 extern bool ConstraintWithNameIsOfType(char *inputConstaintName, Oid relationId,
 									   char targetConstraintType);
 extern bool ConstraintWithIdIsOfType(Oid constraintId, char targetConstraintType);
+extern Oid ConstraintOnRelationId(Oid constraintId);
 extern bool TableHasExternalForeignKeys(Oid relationId);
 extern List * GetForeignKeyOids(Oid relationId, int flags);
 extern Oid GetReferencedTableId(Oid foreignKeyId);

src/include/distributed/commands/utility_hook.h

@@ -93,7 +93,6 @@ extern void ProcessUtilityParseTree(Node *node, const char *queryString,
 									QueryCompletion *completionTag
 									);
 extern void MarkInvalidateForeignKeyGraph(void);
-extern void InvalidateForeignKeyGraphForDDL(void);
 extern List * DDLTaskList(Oid relationId, const char *commandString);
 extern List * NodeDDLTaskList(TargetWorkerSet targets, List *commands);
 extern bool AlterTableInProgress(void);