From db5bae8ffa27d363f23dc63f6d5afe365f961962 Mon Sep 17 00:00:00 2001
From: Onur Tirtir <onurcantirtir@gmail.com>
Date: Fri, 25 Dec 2020 10:35:09 +0300
Subject: [PATCH] Cover DROP COLUMN commands dropping columns in the either
 side of fkeys

---
 .../distributed/commands/foreign_constraint.c     | 15 +++++++++++++++
 src/backend/distributed/commands/table.c          | 11 +++++++++++
 src/include/distributed/commands.h                |  2 +-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/commands/foreign_constraint.c b/src/backend/distributed/commands/foreign_constraint.c
index 12ec279e6..a78de3be1 100644
--- a/src/backend/distributed/commands/foreign_constraint.c
+++ b/src/backend/distributed/commands/foreign_constraint.c
@@ -483,6 +483,21 @@ ForeignConstraintFindDistKeys(HeapTuple pgConstraintTuple,
 }
 
 
+/*
+ * ColumnAppearsInForeignKey returns true if there is a foreign key constraint
+ * from/to given column.
+ */
+bool
+ColumnAppearsInForeignKey(char *columnName, Oid relationId)
+{
+	int searchForeignKeyColumnFlags = SEARCH_REFERENCING_RELATION |
+									  SEARCH_REFERENCED_RELATION;
+	List *foreignKeyIdsColumnAppeared =
+		GetForeignKeyIdsForColumn(columnName, relationId, searchForeignKeyColumnFlags);
+	return list_length(foreignKeyIdsColumnAppeared) > 0;
+}
+
+
 /*
  * ColumnAppearsInForeignKeyToReferenceTable checks if there is a foreign key
  * constraint from/to any reference table on the given column.
diff --git a/src/backend/distributed/commands/table.c b/src/backend/distributed/commands/table.c
index e417a6fa0..3afc324a3 100644
--- a/src/backend/distributed/commands/table.c
+++ b/src/backend/distributed/commands/table.c
@@ -1099,6 +1099,17 @@ AlterTableCmdDropsFkey(AlterTableCmd *command, Oid relationId)
 			}
 		}
 	}
+	else if (alterTableType == AT_DropColumn)
+	{
+		char *columnName = command->name;
+		if (ColumnAppearsInForeignKey(columnName, relationId))
+		{
+			return true;
+		}
+	}
+
+	return false;
+}
 
 
 /*
diff --git a/src/include/distributed/commands.h b/src/include/distributed/commands.h
index 5765b0590..370f3b9bd 100644
--- a/src/include/distributed/commands.h
+++ b/src/include/distributed/commands.h
@@ -159,7 +159,7 @@ extern void ErrorIfUnsupportedForeignConstraintExists(Relation relation,
 													  Var *distributionColumn,
 													  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);