Cover DROP unique/primary CONSTRAINT commands dropping fkeys

src/backend/distributed/commands/foreign_constraint.c

@@ -740,6 +740,23 @@ TableReferencing(Oid relationId)
 }
 
 
+/*
+ * ConstraintWithNameIsOfType 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.

src/backend/distributed/commands/table.c

@@ -18,6 +18,7 @@
 #include "catalog/index.h"
 #include "catalog/pg_class.h"
 #include "catalog/pg_constraint.h"
+#include "catalog/pg_depend.h"
 #include "commands/tablecmds.h"
 #include "distributed/citus_ruleutils.h"
 #include "distributed/colocation_utils.h"
@@ -59,6 +60,7 @@ static bool AlterTableHasCommandByFunc(AlterTableStmt *alterTableStatement,
 static bool AlterTableCmdAddsOrDropsFkey(AlterTableCmd *command, Oid relationId);
 static bool AlterTableCmdAddsFKey(AlterTableCmd *command, Oid relationId);
 static bool AlterTableCmdDropsFkey(AlterTableCmd *command, Oid relationId);
+static bool AnyForeignKeyDependsOnIndex(Oid constraintId);
 static void ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement);
 static void ErrorIfCitusLocalTablePartitionCommand(AlterTableCmd *alterTableCmd,
 												   Oid parentRelationId);
@@ -1085,6 +1087,47 @@ AlterTableCmdDropsFkey(AlterTableCmd *command, Oid relationId)
 		{
 			return true;
 		}
+		else if (ConstraintIsAUniquenessConstraint(constraintName, relationId))
+		{
+			bool missingOk = false;
+			Oid uniquenessConstraintId =
+				get_relation_constraint_oid(relationId, constraintName, missingOk);
+			Oid indexId = get_constraint_index(uniquenessConstraintId);
+			if (AnyForeignKeyDependsOnIndex(indexId))
+			{
+				return true;
+			}
+		}
+	}
+
+
+/*
+ * AnyForeignKeyDependsOnIndex scans pg_depend and returns true if given index
+ * is valid and any foreign key depends on it.
+ */
+static 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;

src/include/distributed/commands.h

@@ -167,6 +167,7 @@ extern bool HasForeignKeyToCitusLocalTable(Oid relationId);
 extern bool HasForeignKeyToReferenceTable(Oid relationOid);
 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);