Fix dropping fkey when distributing table

src/backend/distributed/commands/foreign_constraint.c

@@ -502,6 +502,21 @@ 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.
@@ -833,6 +848,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/index.c

@@ -708,6 +708,11 @@ PreprocessDropIndexStmt(Node *node, const char *dropIndexCommand,
 
 		ErrorIfUnsupportedDropIndexStmt(dropIndexStatement);
 
+		if (AnyForeignKeyDependsOnIndex(distributedIndexId))
+		{
+			MarkInvalidateForeignKeyGraph();
+		}
+
 		ddlJob->targetRelationId = distributedRelationId;
 		ddlJob->concurrentIndexCmd = dropIndexStatement->concurrent;
 

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"
@@ -28,6 +29,7 @@
 #include "distributed/listutils.h"
 #include "distributed/coordinator_protocol.h"
 #include "distributed/metadata_sync.h"
+#include "distributed/metadata/dependency.h"
 #include "distributed/multi_executor.h"
 #include "distributed/multi_partitioning_utils.h"
 #include "distributed/reference_table_utils.h"
@@ -68,6 +70,7 @@ 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 void ErrorIfCitusLocalTablePartitionCommand(AlterTableCmd *alterTableCmd,
 												   Oid parentRelationId);
@@ -507,6 +510,18 @@ 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)
 	{
@@ -1070,6 +1085,99 @@ AlterTableCommandTypeIsTrigger(AlterTableType alterTableType)
 }
 
 
+/*
+ * 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
@@ -1696,21 +1804,6 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 				break;
 			}
 
-			case AT_DropConstraint:
-			{
-				if (!OidIsValid(relationId))
-				{
-					return;
-				}
-
-				if (ConstraintIsAForeignKey(command->name, relationId))
-				{
-					MarkInvalidateForeignKeyGraph();
-				}
-
-				break;
-			}
-
 			case AT_EnableTrig:
 			case AT_EnableAlwaysTrig:
 			case AT_EnableReplicaTrig:
@@ -1740,6 +1833,7 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 			case AT_SetNotNull:
 			case AT_ReplicaIdentity:
 			case AT_ValidateConstraint:
+			case AT_DropConstraint: /* we do the check for invalidation in AlterTableDropsForeignKey */
 			{
 				/*
 				 * We will not perform any special check for:

src/include/distributed/commands.h

@@ -180,6 +180,7 @@ 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);
@@ -187,11 +188,13 @@ extern List * GetForeignConstraintToReferenceTablesCommands(Oid relationId);
 extern List * GetForeignConstraintToDistributedTablesCommands(Oid relationId);
 extern List * GetForeignConstraintFromDistributedTablesCommands(Oid relationId);
 extern List * GetForeignConstraintCommandsInternal(Oid relationId, int flags);
+extern bool AnyForeignKeyDependsOnIndex(Oid indexId);
 extern bool HasForeignKeyWithLocalTable(Oid relationId);
 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);

src/test/regress/expected/multi_foreign_key.out

@@ -938,5 +938,239 @@ SELECT create_reference_table('self_referencing_reference_table');
 (1 row)
 
 ALTER TABLE self_referencing_reference_table ADD CONSTRAINT fk FOREIGN KEY(id, other_column_ref) REFERENCES self_referencing_reference_table(id, other_column);
+-- make sure that if fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest2 DROP CONSTRAINT fkey1;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that if a column that is in a fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest2 DROP COLUMN y;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if a column that is in a multi-column index is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int, y int);
+CREATE UNIQUE INDEX indd ON dropfkeytest1(x, y);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x, y) REFERENCES dropfkeytest1(x, y);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest2 DROP COLUMN y CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if a column that is in a multi-column fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int, y int);
+CREATE UNIQUE INDEX indd ON dropfkeytest1(x, y);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x, y) REFERENCES dropfkeytest1(x, y);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest1 DROP COLUMN y CASCADE;
+NOTICE:  drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if an index which fkey relies on is dropped
+-- Citus can see up-to date fkey graph
+-- also irrelevant index drops doesn't affect this
+CREATE TABLE dropfkeytest1 (x int);
+CREATE UNIQUE INDEX i1 ON dropfkeytest1(x);
+CREATE UNIQUE INDEX unrelated_idx ON dropfkeytest1(x);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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 INDEX unrelated_idx CASCADE;
+-- should still error out since we didn't drop the index that foreign key depends on
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+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 INDEX i1 CASCADE;
+NOTICE:  drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if a uniqueness constraint which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest1 DROP CONSTRAINT dropfkeytest1_x_key CASCADE;
+NOTICE:  drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if a primary key which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int primary key);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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
+ALTER TABLE dropfkeytest1 DROP CONSTRAINT dropfkeytest1_pkey CASCADE;
+NOTICE:  drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+-- make sure that even if a schema which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE SCHEMA fkeytestsc;
+CREATE TABLE fkeytestsc.dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES fkeytestsc.dropfkeytest1(x);
+SELECT create_distributed_table ('fkeytestsc.dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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 fkeytestsc CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table fkeytestsc.dropfkeytest1
+drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TABLE dropfkeytest2 CASCADE;
+-- make sure that even if a table which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+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 TABLE dropfkeytest1 CASCADE;
+NOTICE:  drop cascades to constraint fkey1 on table dropfkeytest2
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
 -- we no longer need those tables
-DROP TABLE referenced_by_reference_table, references_to_reference_table, reference_table, reference_table_second, referenced_local_table, self_referencing_reference_table;
+DROP TABLE referenced_by_reference_table, references_to_reference_table, reference_table, reference_table_second, referenced_local_table, self_referencing_reference_table, dropfkeytest2;

src/test/regress/sql/multi_foreign_key.sql

@@ -563,5 +563,138 @@ CREATE TABLE self_referencing_reference_table(
 SELECT create_reference_table('self_referencing_reference_table');
 ALTER TABLE self_referencing_reference_table ADD CONSTRAINT fk FOREIGN KEY(id, other_column_ref) REFERENCES self_referencing_reference_table(id, other_column);
 
+-- make sure that if fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+ALTER TABLE dropfkeytest2 DROP CONSTRAINT fkey1;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that if a column that is in a fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+ALTER TABLE dropfkeytest2 DROP COLUMN y;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if a column that is in a multi-column index is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int, y int);
+CREATE UNIQUE INDEX indd ON dropfkeytest1(x, y);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x, y) REFERENCES dropfkeytest1(x, y);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+ALTER TABLE dropfkeytest2 DROP COLUMN y CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if a column that is in a multi-column fkey is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int, y int);
+CREATE UNIQUE INDEX indd ON dropfkeytest1(x, y);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (x, y) REFERENCES dropfkeytest1(x, y);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+ALTER TABLE dropfkeytest1 DROP COLUMN y CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if an index which fkey relies on is dropped
+-- Citus can see up-to date fkey graph
+-- also irrelevant index drops doesn't affect this
+CREATE TABLE dropfkeytest1 (x int);
+CREATE UNIQUE INDEX i1 ON dropfkeytest1(x);
+CREATE UNIQUE INDEX unrelated_idx ON dropfkeytest1(x);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+DROP INDEX unrelated_idx CASCADE;
+-- should still error out since we didn't drop the index that foreign key depends on
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+DROP INDEX i1 CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if a uniqueness constraint which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+ALTER TABLE dropfkeytest1 DROP CONSTRAINT dropfkeytest1_x_key CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if a primary key which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int primary key);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+ALTER TABLE dropfkeytest1 DROP CONSTRAINT dropfkeytest1_pkey CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest1, dropfkeytest2 CASCADE;
+
+-- make sure that even if a schema which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE SCHEMA fkeytestsc;
+CREATE TABLE fkeytestsc.dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES fkeytestsc.dropfkeytest1(x);
+SELECT create_distributed_table ('fkeytestsc.dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+DROP SCHEMA fkeytestsc CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
+DROP TABLE dropfkeytest2 CASCADE;
+
+-- make sure that even if a table which fkey depends on is dropped
+-- Citus can see up-to date fkey graph
+CREATE TABLE dropfkeytest1 (x int unique);
+CREATE TABLE dropfkeytest2 (x int8, y int8);
+ALTER TABLE dropfkeytest2 ADD CONSTRAINT fkey1 FOREIGN KEY (y) REFERENCES dropfkeytest1(x);
+SELECT create_distributed_table ('dropfkeytest1', 'x');
+-- this should error out
+SELECT create_distributed_table ('dropfkeytest2', 'y', colocate_with:='none');
+DROP TABLE dropfkeytest1 CASCADE;
+-- this should work
+SELECT create_distributed_table ('dropfkeytest2', 'x', colocate_with:='none');
+
 -- we no longer need those tables
-DROP TABLE referenced_by_reference_table, references_to_reference_table, reference_table, reference_table_second, referenced_local_table, self_referencing_reference_table;
+DROP TABLE referenced_by_reference_table, references_to_reference_table, reference_table, reference_table_second, referenced_local_table, self_referencing_reference_table, dropfkeytest2;