Automatically undistribute citus local tables when no more fkeys with reference tables (#4538)

src/backend/distributed/commands/alter_table.c

@@ -41,6 +41,7 @@
 #include "distributed/deparser.h"
 #include "distributed/distribution_column.h"
 #include "distributed/listutils.h"
+#include "distributed/local_executor.h"
 #include "distributed/metadata/dependency.h"
 #include "distributed/metadata_cache.h"
 #include "distributed/metadata_sync.h"
@@ -487,6 +488,15 @@ AlterTableSetAccessMethod(TableConversionParameters *params)
 TableConversionReturn *
 ConvertTable(TableConversionState *con)
 {
+	/*
+	 * We undistribute citus local tables that are not chained with any reference
+	 * tables via foreign keys at the end of the utility hook.
+	 * Here we temporarily set the related GUC to off to disable the logic for
+	 * internally executed DDL's that might invoke this mechanism unnecessarily.
+	 */
+	bool oldEnableLocalReferenceForeignKeys = EnableLocalReferenceForeignKeys;
+	SetLocalEnableLocalReferenceForeignKeys(false);
+
 	if (con->conversionType == UNDISTRIBUTE_TABLE && con->cascadeViaForeignKeys &&
 		(TableReferencing(con->relationId) || TableReferenced(con->relationId)))
 	{
@@ -501,6 +511,7 @@ ConvertTable(TableConversionState *con)
 		 * Undistributed every foreign key connected relation in our foreign key
 		 * subgraph including itself, so return here.
 		 */
+		SetLocalEnableLocalReferenceForeignKeys(oldEnableLocalReferenceForeignKeys);
 		return NULL;
 	}
 	char *newAccessMethod = con->accessMethod ? con->accessMethod :
@@ -742,6 +753,8 @@ ConvertTable(TableConversionState *con)
 	/* increment command counter so that next command can see the new table */
 	CommandCounterIncrement();
 
+	SetLocalEnableLocalReferenceForeignKeys(oldEnableLocalReferenceForeignKeys);
+
 	return ret;
 }
 

src/backend/distributed/commands/cascade_table_operation_for_connected_relations.c

@@ -28,13 +28,13 @@
 #include "distributed/reference_table_utils.h"
 #include "distributed/relation_access_tracking.h"
 #include "distributed/worker_protocol.h"
+#include "miscadmin.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
 
 
 static void EnsureSequentialModeForCitusTableCascadeFunction(List *relationIdList);
-static bool RelationIdListHasReferenceTable(List *relationIdList);
 static void LockRelationsWithLockMode(List *relationIdList, LOCKMODE lockMode);
 static List * RemovePartitionRelationIds(List *relationIdList);
 static List * GetFKeyCreationCommandsForRelationIdList(List *relationIdList);
@@ -222,7 +222,7 @@ EnsureSequentialModeForCitusTableCascadeFunction(List *relationIdList)
  * RelationIdListHasReferenceTable returns true if relationIdList has a relation
  * id that belongs to a reference table.
  */
-static bool
+bool
 RelationIdListHasReferenceTable(List *relationIdList)
 {
 	Oid relationId = InvalidOid;
@@ -282,8 +282,34 @@ DropRelationIdListForeignKeys(List *relationIdList, int fKeyFlags)
 void
 DropRelationForeignKeys(Oid relationId, int fKeyFlags)
 {
+	/*
+	 * We undistribute citus local tables that are not chained with any reference
+	 * tables via foreign keys at the end of the utility hook.
+	 * Here we temporarily set the related GUC to off to disable the logic for
+	 * internally executed DDL's that might invoke this mechanism unnecessarily.
+	 */
+	bool oldEnableLocalReferenceForeignKeys = EnableLocalReferenceForeignKeys;
+	SetLocalEnableLocalReferenceForeignKeys(false);
+
 	List *dropFkeyCascadeCommandList = GetRelationDropFkeyCommands(relationId, fKeyFlags);
 	ExecuteAndLogDDLCommandList(dropFkeyCascadeCommandList);
+
+	SetLocalEnableLocalReferenceForeignKeys(oldEnableLocalReferenceForeignKeys);
+}
+
+
+/*
+ * SetLocalEnableLocalReferenceForeignKeys is simply a C interface for setting
+ * the following:
+ *      SET LOCAL citus.enable_local_reference_table_foreign_keys = 'on'|'off';
+ */
+void
+SetLocalEnableLocalReferenceForeignKeys(bool state)
+{
+	char *stateStr = state ? "on" : "off";
+	set_config_option("citus.enable_local_reference_table_foreign_keys", stateStr,
+					  (superuser() ? PGC_SUSET : PGC_USERSET), PGC_S_SESSION,
+					  GUC_ACTION_LOCAL, true, 0, false);
 }
 
 

src/backend/distributed/commands/create_citus_local_table.c

@@ -72,6 +72,7 @@ static void FinalizeCitusLocalTableCreation(Oid relationId);
 
 
 PG_FUNCTION_INFO_V1(create_citus_local_table);
+PG_FUNCTION_INFO_V1(remove_local_tables_from_metadata);
 
 
 /*
@@ -84,6 +85,23 @@ create_citus_local_table(PG_FUNCTION_ARGS)
 {
 	CheckCitusVersion(ERROR);
 
+	if (ShouldEnableLocalReferenceForeignKeys())
+	{
+		/*
+		 * When foreign keys between reference tables and postgres tables are
+		 * enabled, we automatically undistribute citus local tables that are
+		 * not chained with any reference tables back to postgres tables.
+		 * So give a warning to user for that.
+		 */
+		ereport(WARNING, (errmsg("citus local tables that are not chained with "
+								 "reference tables via foreign keys might be "
+								 "automatically converted back to postgres tables"),
+						  errhint("Consider setting "
+								  "citus.enable_local_reference_table_foreign_keys "
+								  "to 'off' to disable automatically undistributing "
+								  "citus local tables")));
+	}
+
 	Oid relationId = PG_GETARG_OID(0);
 	bool cascadeViaForeignKeys = PG_GETARG_BOOL(1);
 
@@ -93,6 +111,22 @@ create_citus_local_table(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * remove_local_tables_from_metadata undistributes citus local
+ * tables that are not chained with any reference tables via foreign keys.
+ */
+Datum
+remove_local_tables_from_metadata(PG_FUNCTION_ARGS)
+{
+	CheckCitusVersion(ERROR);
+	EnsureCoordinator();
+
+	UndistributeDisconnectedCitusLocalTables();
+
+	PG_RETURN_VOID();
+}
+
+
 /*
  * CreateCitusLocalTable is the internal method that creates a citus table
  * from the table with relationId. The created table would have the following

src/backend/distributed/commands/drop_distributed_table.c

@@ -12,6 +12,7 @@
 #include "miscadmin.h"
 
 #include "distributed/commands/utility_hook.h"
+#include "distributed/commands.h"
 #include "distributed/metadata_utility.h"
 #include "distributed/coordinator_protocol.h"
 #include "distributed/metadata_sync.h"
@@ -30,6 +31,7 @@ static void MasterRemoveDistributedTableMetadataFromWorkers(Oid relationId,
 PG_FUNCTION_INFO_V1(master_drop_distributed_table_metadata);
 PG_FUNCTION_INFO_V1(master_remove_partition_metadata);
 PG_FUNCTION_INFO_V1(master_remove_distributed_table_metadata_from_workers);
+PG_FUNCTION_INFO_V1(notify_constraint_dropped);
 
 
 /*
@@ -148,3 +150,30 @@ MasterRemoveDistributedTableMetadataFromWorkers(Oid relationId, char *schemaName
 	char *deleteDistributionCommand = DistributionDeleteCommand(schemaName, tableName);
 	SendCommandToWorkersWithMetadata(deleteDistributionCommand);
 }
+
+
+/*
+ * notify_constraint_dropped simply calls NotifyUtilityHookConstraintDropped
+ * to set ConstraintDropped to true.
+ * This udf is designed to be called from citus_drop_trigger to tell us we
+ * dropped a table constraint.
+ */
+Datum
+notify_constraint_dropped(PG_FUNCTION_ARGS)
+{
+	CheckCitusVersion(ERROR);
+
+	/*
+	 * We reset this only in utility hook, so we should not set this flag
+	 * otherwise if we are not in utility hook.
+	 * In some cases -where dropping foreign key not issued via utility
+	 * hook-, we would not be able to undistribute such citus local tables
+	 * but we are ok with that.
+	 */
+	if (UtilityHookLevel >= 1)
+	{
+		NotifyUtilityHookConstraintDropped();
+	}
+
+	PG_RETURN_VOID();
+}

src/backend/distributed/commands/table.c

@@ -48,7 +48,6 @@ bool EnableLocalReferenceForeignKeys = true;
 
 /* Local functions forward declarations for unsupported command checks */
 static void PostprocessCreateTableStmtForeignKeys(CreateStmt *createStatement);
-static bool ShouldEnableLocalReferenceForeignKeys(void);
 static void PostprocessCreateTableStmtPartitionOf(CreateStmt *createStatement,
 												  const char *queryString);
 static bool AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(
@@ -251,7 +250,7 @@ PostprocessCreateTableStmtForeignKeys(CreateStmt *createStatement)
  * the value set by the user
  *
  */
-static bool
+bool
 ShouldEnableLocalReferenceForeignKeys(void)
 {
 	if (!EnableLocalReferenceForeignKeys)

src/backend/distributed/commands/utility_hook.c

@@ -46,6 +46,7 @@
 #include "distributed/commands/utility_hook.h" /* IWYU pragma: keep */
 #include "distributed/deparser.h"
 #include "distributed/deparse_shard_query.h"
+#include "distributed/foreign_key_relationship.h"
 #include "distributed/listutils.h"
 #include "distributed/local_executor.h"
 #include "distributed/maintenanced.h"
@@ -55,6 +56,7 @@
 #include "distributed/multi_executor.h"
 #include "distributed/multi_explain.h"
 #include "distributed/multi_physical_planner.h"
+#include "distributed/reference_table_utils.h"
 #include "distributed/resource_lock.h"
 #include "distributed/transmit.h"
 #include "distributed/version_compat.h"
@@ -72,6 +74,10 @@ PropSetCmdBehavior PropagateSetCommands = PROPSETCMD_NONE; /* SET prop off */
 static bool shouldInvalidateForeignKeyGraph = false;
 static int activeAlterTables = 0;
 static int activeDropSchemaOrDBs = 0;
+static bool ConstraintDropped = false;
+
+
+int UtilityHookLevel = 0;
 
 
 /* Local functions forward declarations for helper functions */
@@ -88,6 +94,7 @@ static void IncrementUtilityHookCountersIfNecessary(Node *parsetree);
 static void PostStandardProcessUtility(Node *parsetree);
 static void DecrementUtilityHookCountersIfNecessary(Node *parsetree);
 static bool IsDropSchemaOrDB(Node *parsetree);
+static bool ShouldUndistributeCitusLocalTables(void);
 
 
 /*
@@ -237,8 +244,41 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 		return;
 	}
 
-	ProcessUtilityInternal(pstmt, queryString, context,
+	UtilityHookLevel++;
-						   params, queryEnv, dest, completionTag);
+
+	PG_TRY();
+	{
+		ProcessUtilityInternal(pstmt, queryString, context, params, queryEnv, dest,
+							   completionTag);
+
+		if (UtilityHookLevel == 1)
+		{
+			/*
+			 * When Citus local tables are disconnected from the foreign key graph, which
+			 * can happen due to various kinds of drop commands, we immediately
+			 * undistribute them at the end of the command.
+			 */
+			if (ShouldUndistributeCitusLocalTables())
+			{
+				UndistributeDisconnectedCitusLocalTables();
+			}
+			ResetConstraintDropped();
+		}
+
+		UtilityHookLevel--;
+	}
+	PG_CATCH();
+	{
+		if (UtilityHookLevel == 1)
+		{
+			ResetConstraintDropped();
+		}
+
+		UtilityHookLevel--;
+
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
 }
 
 
@@ -647,6 +687,138 @@ ProcessUtilityInternal(PlannedStmt *pstmt,
 }
 
 
+/*
+ * UndistributeDisconnectedCitusLocalTables undistributes citus local tables that
+ * are not connected to any reference tables via their individual foreign key
+ * subgraphs.
+ */
+void
+UndistributeDisconnectedCitusLocalTables(void)
+{
+	List *citusLocalTableIdList = CitusTableTypeIdList(CITUS_LOCAL_TABLE);
+	citusLocalTableIdList = SortList(citusLocalTableIdList, CompareOids);
+
+	Oid citusLocalTableId = InvalidOid;
+	foreach_oid(citusLocalTableId, citusLocalTableIdList)
+	{
+		/* acquire ShareRowExclusiveLock to prevent concurrent foreign key creation */
+		LOCKMODE lockMode = ShareRowExclusiveLock;
+		LockRelationOid(citusLocalTableId, lockMode);
+
+		HeapTuple heapTuple =
+			SearchSysCache1(RELOID, ObjectIdGetDatum(citusLocalTableId));
+		if (!HeapTupleIsValid(heapTuple))
+		{
+			/*
+			 * UndistributeTable drops relation, skip if already undistributed
+			 * via cascade.
+			 */
+			continue;
+		}
+		ReleaseSysCache(heapTuple);
+
+		if (ConnectedToReferenceTableViaFKey(citusLocalTableId))
+		{
+			/* still connected to a reference table, skip it */
+			UnlockRelationOid(citusLocalTableId, lockMode);
+			continue;
+		}
+
+		/*
+		 * Citus local table is not connected to any reference tables, then
+		 * undistribute it via cascade. Here, instead of first dropping foreing
+		 * keys then undistributing the table, we just set cascadeViaForeignKeys
+		 * to true for simplicity.
+		 */
+		TableConversionParameters params = {
+			.relationId = citusLocalTableId,
+			.cascadeViaForeignKeys = true
+		};
+		UndistributeTable(&params);
+	}
+}
+
+
+/*
+ * ShouldUndistributeCitusLocalTables returns true if we might need to check
+ * citus local tables for their connectivity to reference tables.
+ */
+static bool
+ShouldUndistributeCitusLocalTables(void)
+{
+	if (!ConstraintDropped)
+	{
+		/*
+		 * citus_drop_trigger executes notify_constraint_dropped to set
+		 * ConstraintDropped to true, which means that last command dropped
+		 * a table constraint.
+		 */
+		return false;
+	}
+
+	if (!CitusHasBeenLoaded())
+	{
+		/*
+		 * If we are dropping citus, we should not try to undistribute citus
+		 * local tables as they will also be dropped.
+		 */
+		return false;
+	}
+
+	if (!InCoordinatedTransaction())
+	{
+		/* not interacting with any Citus objects */
+		return false;
+	}
+
+	if (IsCitusInitiatedBackend())
+	{
+		/* connection from the coordinator operating on a shard */
+		return false;
+	}
+
+	if (!ShouldEnableLocalReferenceForeignKeys())
+	{
+		/*
+		 * If foreign keys between reference tables and local tables are
+		 * disabled, then user might be using create_citus_local_table for
+		 * their own purposes. In that case, we should not undistribute
+		 * citus local tables.
+		 */
+		return false;
+	}
+
+	if (!IsCoordinator())
+	{
+		/* we should not perform this operation in worker nodes */
+		return false;
+	}
+
+	return true;
+}
+
+
+/*
+ * NotifyUtilityHookConstraintDropped sets ConstraintDropped to true to tell us
+ * last command dropped a table constraint.
+ */
+void
+NotifyUtilityHookConstraintDropped(void)
+{
+	ConstraintDropped = true;
+}
+
+
+/*
+ * ResetConstraintDropped sets ConstraintDropped to false.
+ */
+void
+ResetConstraintDropped(void)
+{
+	ConstraintDropped = false;
+}
+
+
 /*
  * IsDropSchemaOrDB returns true if parsetree represents DROP SCHEMA ...or
  * a DROP DATABASE.

src/backend/distributed/connection/connection_management.c

@@ -1292,20 +1292,13 @@ AfterXactHostConnectionHandling(ConnectionHashEntry *entry, bool isCommit)
 static bool
 ShouldShutdownConnection(MultiConnection *connection, const int cachedConnectionCount)
 {
-	bool isCitusInitiatedBackend = false;
-
 	/*
 	 * When we are in a backend that was created to serve an internal connection
 	 * from the coordinator or another worker, we disable connection caching to avoid
 	 * escalating the number of cached connections. We can recognize such backends
 	 * from their application name.
 	 */
-	if (application_name != NULL && strcmp(application_name, CITUS_APPLICATION_NAME) == 0)
+	return IsCitusInitiatedBackend() ||
-	{
-		isCitusInitiatedBackend = true;
-	}
-
-	return isCitusInitiatedBackend ||
 		   connection->initilizationState != POOL_STATE_INITIALIZED ||
 		   cachedConnectionCount >= MaxCachedConnectionsPerWorker ||
 		   connection->forceCloseAtTransactionEnd ||
@@ -1314,6 +1307,17 @@ ShouldShutdownConnection(MultiConnection *connection, const int cachedConnection
 }
 
 
+/*
+ * IsCitusInitiatedBackend returns true if we are in a backend that citus
+ * initiated via remote connection.
+ */
+bool
+IsCitusInitiatedBackend(void)
+{
+	return application_name && strcmp(application_name, CITUS_APPLICATION_NAME) == 0;
+}
+
+
 /*
  * ResetConnection preserves the given connection for later usage by
  * resetting its states.

src/backend/distributed/sql/citus--9.5-1--10.0-1.sql

@@ -27,6 +27,7 @@ DROP FUNCTION IF EXISTS pg_catalog.citus_total_relation_size(regclass);
 #include "udfs/citus_move_shard_placement/10.0-1.sql"
 #include "udfs/citus_drop_trigger/10.0-1.sql"
 #include "udfs/worker_change_sequence_dependency/10.0-1.sql"
+#include "udfs/remove_local_tables_from_metadata/10.0-1.sql"
 
 #include "../../columnar/sql/columnar--9.5-1--10.0-1.sql"
 

src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql

@@ -89,6 +89,9 @@ CREATE FUNCTION pg_catalog.master_create_worker_shards(table_name text, shard_co
     AS 'MODULE_PATHNAME'
     LANGUAGE C STRICT;
 
+DROP FUNCTION pg_catalog.notify_constraint_dropped();
+DROP FUNCTION pg_catalog.remove_local_tables_from_metadata();
+
 #include "../udfs/citus_drop_trigger/9.5-1.sql"
 #include "../udfs/citus_total_relation_size/7.0-1.sql"
 #include "../udfs/upgrade_to_reference_table/8.0-1.sql"

src/backend/distributed/sql/udfs/citus_drop_trigger/10.0-1.sql

@@ -1,9 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.notify_constraint_dropped()
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$notify_constraint_dropped$$;
+
 CREATE OR REPLACE FUNCTION pg_catalog.citus_drop_trigger()
     RETURNS event_trigger
     LANGUAGE plpgsql
     SET search_path = pg_catalog
     AS $cdbdt$
 DECLARE
+    constraint_event_count INTEGER;
     v_obj record;
     sequence_names text[] := '{}';
     table_colocation_id integer;
@@ -25,6 +31,18 @@ BEGIN
     LOOP
         PERFORM master_unmark_object_distributed(v_obj.classid, v_obj.objid, v_obj.objsubid);
     END LOOP;
+
+    SELECT COUNT(*) INTO constraint_event_count
+    FROM pg_event_trigger_dropped_objects()
+    WHERE object_type IN ('table constraint');
+
+    IF constraint_event_count > 0
+    THEN
+        -- Tell utility hook that a table constraint is dropped so we might
+        -- need to undistribute some of the citus local tables that are not
+        -- connected to any reference tables.
+        PERFORM notify_constraint_dropped();
+    END IF;
 END;
 $cdbdt$;
 COMMENT ON FUNCTION pg_catalog.citus_drop_trigger()

src/backend/distributed/sql/udfs/citus_drop_trigger/latest.sql

@@ -1,9 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.notify_constraint_dropped()
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$notify_constraint_dropped$$;
+
 CREATE OR REPLACE FUNCTION pg_catalog.citus_drop_trigger()
     RETURNS event_trigger
     LANGUAGE plpgsql
     SET search_path = pg_catalog
     AS $cdbdt$
 DECLARE
+    constraint_event_count INTEGER;
     v_obj record;
     sequence_names text[] := '{}';
     table_colocation_id integer;
@@ -25,6 +31,18 @@ BEGIN
     LOOP
         PERFORM master_unmark_object_distributed(v_obj.classid, v_obj.objid, v_obj.objsubid);
     END LOOP;
+
+    SELECT COUNT(*) INTO constraint_event_count
+    FROM pg_event_trigger_dropped_objects()
+    WHERE object_type IN ('table constraint');
+
+    IF constraint_event_count > 0
+    THEN
+        -- Tell utility hook that a table constraint is dropped so we might
+        -- need to undistribute some of the citus local tables that are not
+        -- connected to any reference tables.
+        PERFORM notify_constraint_dropped();
+    END IF;
 END;
 $cdbdt$;
 COMMENT ON FUNCTION pg_catalog.citus_drop_trigger()

src/backend/distributed/sql/udfs/remove_local_tables_from_metadata/10.0-1.sql

@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION pg_catalog.remove_local_tables_from_metadata()
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'MODULE_PATHNAME', $$remove_local_tables_from_metadata$$;
+COMMENT ON FUNCTION pg_catalog.remove_local_tables_from_metadata()
+	IS 'undistribute citus local tables that are not chained with any reference tables via foreign keys';

src/backend/distributed/sql/udfs/remove_local_tables_from_metadata/latest.sql

@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION pg_catalog.remove_local_tables_from_metadata()
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'MODULE_PATHNAME', $$remove_local_tables_from_metadata$$;
+COMMENT ON FUNCTION pg_catalog.remove_local_tables_from_metadata()
+	IS 'undistribute citus local tables that are not chained with any reference tables via foreign keys';

src/backend/distributed/test/foreign_key_relationship_query.c

@@ -14,12 +14,15 @@
 #include "fmgr.h"
 #include "funcapi.h"
 
+#include "catalog/dependency.h"
+#include "catalog/pg_constraint.h"
 #include "distributed/foreign_key_relationship.h"
 #include "distributed/coordinator_protocol.h"
 #include "distributed/listutils.h"
 #include "distributed/metadata_cache.h"
 #include "distributed/tuplestore.h"
 #include "distributed/version_compat.h"
+#include "utils/builtins.h"
 
 
 #define GET_FKEY_CONNECTED_RELATIONS_COLUMNS 1
@@ -29,6 +32,42 @@
 PG_FUNCTION_INFO_V1(get_referencing_relation_id_list);
 PG_FUNCTION_INFO_V1(get_referenced_relation_id_list);
 PG_FUNCTION_INFO_V1(get_foreign_key_connected_relations);
+PG_FUNCTION_INFO_V1(drop_constraint_cascade_via_perform_deletion);
+
+
+/*
+ * drop_constraint_cascade_via_perform_deletion simply drops constraint on
+ * relation via performDeletion.
+ */
+Datum
+drop_constraint_cascade_via_perform_deletion(PG_FUNCTION_ARGS)
+{
+	Oid relationId = PG_GETARG_OID(0);
+
+	if (PG_ARGISNULL(1))
+	{
+		/* avoid unexpected crashes in regression tests */
+		ereport(ERROR, (errmsg("cannot perform operation without constraint "
+							   "name argument")));
+	}
+
+	text *constraintNameText = PG_GETARG_TEXT_P(1);
+	char *constraintName = text_to_cstring(constraintNameText);
+
+	/* error if constraint does not exist */
+	bool missingOk = false;
+	Oid constraintId = get_relation_constraint_oid(relationId, constraintName, missingOk);
+
+	ObjectAddress constraintObjectAddress;
+	constraintObjectAddress.classId = ConstraintRelationId;
+	constraintObjectAddress.objectId = constraintId;
+	constraintObjectAddress.objectSubId = 0;
+
+	performDeletion(&constraintObjectAddress, DROP_CASCADE, 0);
+
+	PG_RETURN_VOID();
+}
+
 
 /*
  * get_referencing_relation_id_list returns the list of table oids that is referencing

src/backend/distributed/utils/foreign_key_relationship.c

@@ -23,9 +23,11 @@
 #include "access/table.h"
 #endif
 #include "catalog/pg_constraint.h"
+#include "distributed/commands.h"
 #include "distributed/foreign_key_relationship.h"
 #include "distributed/hash_helpers.h"
 #include "distributed/listutils.h"
+#include "distributed/metadata_cache.h"
 #include "distributed/version_compat.h"
 #include "nodes/pg_list.h"
 #include "storage/lockdefs.h"
@@ -144,6 +146,24 @@ GetForeignKeyConnectedRelationIdList(Oid relationId)
 }
 
 
+/*
+ * ConnectedToReferenceTableViaFKey returns true if given relationId is
+ * connected to a reference table via its foreign key subgraph.
+ */
+bool
+ConnectedToReferenceTableViaFKey(Oid relationId)
+{
+	/*
+	 * As we will operate on foreign key connected relations, here we
+	 * invalidate foreign key graph so that we act on fresh graph.
+	 */
+	InvalidateForeignKeyGraph();
+
+	List *fkeyConnectedRelations = GetForeignKeyConnectedRelationIdList(relationId);
+	return RelationIdListHasReferenceTable(fkeyConnectedRelations);
+}
+
+
 /*
  * GetRelationshipNodesForFKeyConnectedRelations performs breadth-first search
  * starting from input ForeignConstraintRelationshipNode and returns a list

src/include/distributed/commands.h

@@ -360,6 +360,7 @@ extern List * PreprocessDropTableStmt(Node *stmt, const char *queryString,
 									  ProcessUtilityContext processUtilityContext);
 extern void PostprocessCreateTableStmt(CreateStmt *createStatement,
 									   const char *queryString);
+extern bool ShouldEnableLocalReferenceForeignKeys(void);
 extern List * PostprocessAlterTableStmtAttachPartition(
 	AlterTableStmt *alterTableStatement,
 	const char *queryString);
@@ -481,7 +482,9 @@ extern void CascadeOperationForConnectedRelations(Oid relationId, LOCKMODE relLo
 												  CascadeOperationType
 												  cascadeOperationType);
 extern void ErrorIfAnyPartitionRelationInvolvedInNonInheritedFKey(List *relationIdList);
+extern bool RelationIdListHasReferenceTable(List *relationIdList);
 extern void DropRelationForeignKeys(Oid relationId, int flags);
+extern void SetLocalEnableLocalReferenceForeignKeys(bool state);
 extern void ExecuteAndLogDDLCommandList(List *ddlCommandList);
 extern void ExecuteAndLogDDLCommand(const char *commandString);
 extern void ExecuteForeignKeyCreateCommandList(List *ddlCommandList,

src/include/distributed/commands/utility_hook.h

@@ -35,6 +35,8 @@ extern bool EnableDependencyCreation;
 extern bool EnableCreateTypePropagation;
 extern bool EnableAlterRolePropagation;
 extern bool EnableAlterRoleSetPropagation;
+extern int UtilityHookLevel;
+
 
 /*
  * A DDLJob encapsulates the remote tasks and commands needed to process all or
@@ -77,6 +79,9 @@ extern List * DDLTaskList(Oid relationId, const char *commandString);
 extern List * NodeDDLTaskList(TargetWorkerSet targets, List *commands);
 extern bool AlterTableInProgress(void);
 extern bool DropSchemaOrDBInProgress(void);
+extern void UndistributeDisconnectedCitusLocalTables(void);
+extern void NotifyUtilityHookConstraintDropped(void);
+extern void ResetConstraintDropped(void);
 extern void ExecuteDistributedDDLJob(DDLJob *ddlJob);
 
 /* forward declarations for sending custom commands to a distributed table */

src/include/distributed/connection_management.h

@@ -252,6 +252,7 @@ extern void FinishConnectionListEstablishment(List *multiConnectionList);
 extern void FinishConnectionEstablishment(MultiConnection *connection);
 extern void ClaimConnectionExclusively(MultiConnection *connection);
 extern void UnclaimConnection(MultiConnection *connection);
+extern bool IsCitusInitiatedBackend(void);
 
 /* time utilities */
 extern double MillisecondsPassedSince(instr_time moment);

src/include/distributed/foreign_key_relationship.h

@@ -16,6 +16,7 @@
 #include "nodes/primnodes.h"
 
 extern List * GetForeignKeyConnectedRelationIdList(Oid relationId);
+extern bool ConnectedToReferenceTableViaFKey(Oid relationId);
 extern List * ReferencedRelationIdList(Oid relationId);
 extern List * ReferencingRelationIdList(Oid relationId);
 extern void SetForeignConstraintRelationshipGraphInvalid(void);

src/test/regress/bin/normalize.sed

@@ -200,3 +200,7 @@ s/(NOTICE:  executing.*)\([0-9]+, 'citus_local_tables_test_schema', [0-9]+(.*)/\
 s/citus_local_table_4_idx_[0-9]+/citus_local_table_4_idx_xxxxxx/g
 s/citus_local_table_4_[0-9]+/citus_local_table_4_xxxxxx/g
 s/ERROR:  cannot append to shardId [0-9]+/ERROR:  cannot append to shardId xxxxxx/g
+
+# hide warning/hint message that we get when executing create_citus_local_table
+/citus local tables that are not chained with reference tables via foreign keys might be automatically converted back to postgres tables$/d
+/Consider setting citus.enable_local_reference_table_foreign_keys to 'off' to disable automatically undistributing citus local tables$/d

src/test/regress/expected/auto_undist_citus_local.out

@@ -0,0 +1,635 @@
+-- regression tests regarding foreign key
+-- drops cascading into undistributing Citus
+-- local tables to Postgres local tables
+CREATE SCHEMA drop_fkey_cascade;
+SET search_path TO drop_fkey_cascade;
+SET client_min_messages TO WARNING;
+SET citus.next_shard_id TO 1810000;
+SELECT 1 FROM master_add_node('localhost', :master_port, groupId => 0);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+-- show that DROP CONSTRAINT cascades to undistributing citus_local_table
+CREATE TABLE citus_local_table(l1 int);
+SELECT create_citus_local_table('citus_local_table');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table(r1 int primary key);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1) ON DELETE CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table | n          | c
+ reference_table   | n          | t
+(2 rows)
+
+CREATE OR REPLACE FUNCTION drop_constraint_cascade_via_perform_deletion(IN table_name regclass, IN constraint_name text)
+RETURNS VOID
+LANGUAGE C STRICT
+AS 'citus', $$drop_constraint_cascade_via_perform_deletion$$;
+BEGIN;
+  SELECT drop_constraint_cascade_via_perform_deletion('citus_local_table', 'fkey_local_to_ref');
+ drop_constraint_cascade_via_perform_deletion
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- we dropped constraint without going through utility hook,
+  -- so we should still see citus_local_table
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table | n          | c
+ reference_table   | n          | t
+(2 rows)
+
+ROLLBACK;
+ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+DROP TABLE citus_local_table, reference_table;
+-- show that DROP COLUMN cascades to undistributing citus_local_table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+ALTER TABLE reference_table DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+DROP TABLE citus_local_table, reference_table;
+-- show that DROP COLUMN that cascades into drop foreign key undistributes local table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+ALTER TABLE citus_local_table DROP COLUMN l1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+DROP TABLE citus_local_table, reference_table;
+-- show that PRIMARY KEY that cascades into drop foreign key undistributes local table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_pkey CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+-- show that DROP UNIQUE INDEX that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+CREATE TABLE reference_table(r1 int, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE UNIQUE INDEX ref_unique ON reference_table(r1);
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+DROP INDEX ref_unique CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+-- show that UNIQUE CONSTRAINT that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_r1_key  CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+-- show that DROP TABLE that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+DROP TABLE reference_table CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+-- show that UNIQUE CONSTRAINT that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table;
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table   | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_r1_key  CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+  logicalrelid   | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table | n          | t
+(1 row)
+
+-- show that DROP SCHEMA that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+CREATE SCHEMA ref_table_drop_schema;
+CREATE TABLE ref_table_drop_schema.reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('ref_table_drop_schema.reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES ref_table_drop_schema.reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'ref_table_drop_schema.reference_table'::regclass) ORDER BY logicalrelid;
+             logicalrelid              | partmethod | repmodel
+---------------------------------------------------------------------
+ ref_table_drop_schema.reference_table | n          | t
+ citus_local_table                     | n          | c
+(2 rows)
+
+DROP SCHEMA ref_table_drop_schema CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+-- drop column cascade that doesn't cascade into citus local table
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table_2(r1 int UNIQUE REFERENCES reference_table_1(r1), r2 int);
+SELECT create_reference_table('reference_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_2(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+ALTER TABLE reference_table_1 DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+-- local table has multiple foreign keys to two tables
+-- drop one at a time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+DROP TABLE reference_table_1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(2 rows)
+
+CREATE TABLE distributed_table (d1 int);
+SELECT create_distributed_table('distributed_table', 'd1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- drop an unrelated distributed table too
+DROP TABLE reference_table_2, distributed_table CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+-- local table has multiple foreign keys to two tables
+-- drop both at the same time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+DROP TABLE reference_table_1, reference_table_2 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+-- local table has multiple foreign keys to two tables
+-- drop one at a time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+BEGIN;
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l1_fkey;
+  SAVEPOINT sp1;
+  -- this should undistribute citus_local_table
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l2_fkey;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+(2 rows)
+
+  ROLLBACK TO SAVEPOINT sp1;
+  -- rollback'ed second drop constraint, so we should still see citus_local_table
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+ citus_local_table | n          | c
+(3 rows)
+
+  -- this should undistribute citus_local_table again
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l2_fkey;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+ reference_table_2 | n          | t
+(2 rows)
+
+COMMIT;
+-- a single drop column cascades into multiple undistributes
+DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, reference_table_1;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+CREATE TABLE citus_local_table_2(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE REFERENCES citus_local_table_1(l2));
+CREATE TABLE citus_local_table_3(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES citus_local_table_2(l2));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+    logicalrelid     | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1   | n          | t
+ citus_local_table_1 | n          | c
+ citus_local_table_2 | n          | c
+ citus_local_table_3 | n          | c
+(4 rows)
+
+ALTER TABLE reference_table_1 DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+(1 row)
+
+-- a single drop table cascades into multiple undistributes
+DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, citus_local_table_3, citus_local_table_2, reference_table_1;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+CREATE TABLE citus_local_table_2(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE REFERENCES citus_local_table_1(l2));
+CREATE TABLE citus_local_table_3(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES citus_local_table_2(l2));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+    logicalrelid     | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1   | n          | t
+ citus_local_table_1 | n          | c
+ citus_local_table_2 | n          | c
+ citus_local_table_3 | n          | c
+(4 rows)
+
+-- test DROP OWNED BY
+-- Citus does not support "ALTER TABLE OWNER TO" commands. Also, not to deal with tests output
+-- difference between community and enterprise, let's disable enable_ddl_propagation here.
+SET citus.enable_ddl_propagation to OFF;
+CREATE USER another_user;
+SELECT run_command_on_workers('CREATE USER another_user');
+      run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"CREATE ROLE")
+ (localhost,57638,t,"CREATE ROLE")
+(2 rows)
+
+ALTER TABLE reference_table_1 OWNER TO another_user;
+SELECT run_command_on_placements('reference_table_1', 'ALTER TABLE %s OWNER TO another_user');
+         run_command_on_placements
+---------------------------------------------------------------------
+ (localhost,57636,1810039,t,"ALTER TABLE")
+ (localhost,57637,1810039,t,"ALTER TABLE")
+ (localhost,57638,1810039,t,"ALTER TABLE")
+(3 rows)
+
+SET citus.enable_ddl_propagation to ON;
+BEGIN;
+  DROP OWNED BY another_user cascade;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ( 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+ROLLBACK;
+DROP TABLE reference_table_1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ( 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+ logicalrelid | partmethod | repmodel
+---------------------------------------------------------------------
+(0 rows)
+
+-- dropping constraints inside a plpgsql procedure should be fine
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+    logicalrelid     | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1   | n          | t
+ citus_local_table_1 | n          | c
+(2 rows)
+
+CREATE OR REPLACE FUNCTION drop_constraint_via_func()
+RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+ALTER TABLE citus_local_table_1 DROP CONSTRAINT citus_local_table_1_l1_fkey;
+END;$$;
+BEGIN;
+  SELECT drop_constraint_via_func();
+ drop_constraint_via_func
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+(1 row)
+
+ROLLBACK;
+create or replace procedure drop_constraint_via_proc()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+        ALTER TABLE citus_local_table_1 DROP CONSTRAINT citus_local_table_1_l1_fkey;
+    commit;
+end;$$;
+call drop_constraint_via_proc();
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+(1 row)
+
+-- even if the procedure is called from another procedure
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+    logicalrelid     | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1   | n          | t
+ citus_local_table_1 | n          | c
+(2 rows)
+
+create or replace procedure drop_constraint_via_proc_top_level()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+        CALL drop_constraint_via_proc();
+    commit;
+end;$$;
+CALL drop_constraint_via_proc_top_level();
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+(1 row)
+
+-- even if the procedure is called from an exception handler
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+    logicalrelid     | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1   | n          | t
+ citus_local_table_1 | n          | c
+(2 rows)
+
+create or replace procedure drop_constraint_via_proc_exception()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+		PERFORM 1/0;
+    	EXCEPTION
+             when others then
+        CALL drop_constraint_via_proc();
+    commit;
+end;$$;
+CALL drop_constraint_via_proc_exception();
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+   logicalrelid    | partmethod | repmodel
+---------------------------------------------------------------------
+ reference_table_1 | n          | t
+(1 row)
+
+DROP SCHEMA drop_fkey_cascade CASCADE;

src/test/regress/expected/citus_local_tables.out

@@ -259,6 +259,8 @@ NOTICE:  foreign-data wrapper "fake_fdw" does not have an extension defined
 
 (1 row)
 
+DROP FOREIGN TABLE foreign_table;
+NOTICE:  executing the command locally: DROP FOREIGN TABLE IF EXISTS citus_local_tables_test_schema.foreign_table_xxxxx CASCADE
 -- drop them for next tests
 DROP TABLE citus_local_table_1, citus_local_table_2, distributed_table;
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS citus_local_tables_test_schema.citus_local_table_2_xxxxx CASCADE
@@ -413,6 +415,67 @@ NOTICE:  executing the command locally: CREATE UNIQUE INDEX   uniqueindex2_15040
 ---- utility command execution ----
 ---------------------------------------------------------------------
 SET search_path TO citus_local_tables_test_schema;
+CREATE TABLE dummy_reference_table (a INT PRIMARY KEY);
+SELECT create_reference_table('dummy_reference_table');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+BEGIN;
+  SET client_min_messages TO ERROR;
+  SELECT remove_local_tables_from_metadata();
+ remove_local_tables_from_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- should not see any citus local tables
+  SELECT logicalrelid::regclass::text FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='citus_local_tables_test_schema' AND
+        partmethod = 'n' AND repmodel = 'c'
+  ORDER BY 1;
+ logicalrelid
+---------------------------------------------------------------------
+(0 rows)
+
+ROLLBACK;
+-- define foreign keys between dummy_reference_table and citus local tables
+-- not to undistribute them automatically
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'citus_local_tables_test_schema', 'ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);')
+ALTER TABLE citus_local_table_2 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'citus_local_tables_test_schema', 'ALTER TABLE citus_local_table_2 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);')
+ALTER TABLE unlogged_table ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'citus_local_tables_test_schema', 'ALTER TABLE unlogged_table ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);')
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'citus_local_tables_test_schema', 'ALTER TABLE local_table_3 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);')
+ALTER TABLE dummy_reference_table ADD CONSTRAINT fkey_from_dummy_ref FOREIGN KEY (a) REFERENCES "CiTUS!LocalTables"."LocalTabLE.1!?!"(id);
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'CiTUS!LocalTables', 'ALTER TABLE dummy_reference_table ADD CONSTRAINT fkey_from_dummy_ref FOREIGN KEY (a) REFERENCES "CiTUS!LocalTables"."LocalTabLE.1!?!"(id);')
+BEGIN;
+  SET client_min_messages TO ERROR;
+  SELECT remove_local_tables_from_metadata();
+ remove_local_tables_from_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- now we defined foreign keys with above citus local tables, we should still see them
+  SELECT logicalrelid::regclass::text FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='citus_local_tables_test_schema' AND
+        partmethod = 'n' AND repmodel = 'c'
+  ORDER BY 1;
+    logicalrelid
+---------------------------------------------------------------------
+ citus_local_table_1
+ citus_local_table_2
+ local_table_3
+ unlogged_table
+(4 rows)
+
+ROLLBACK;
 -- between citus local tables and distributed tables
 ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_c_to_dist FOREIGN KEY(a) references distributed_table(a);
 ERROR:  cannot create foreign key constraint since foreign keys from reference tables and citus local tables to distributed tables are not supported
@@ -675,21 +738,17 @@ NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_comm
 -- verify creating citus local table with extended statistics
 CREATE TABLE test_citus_local_table_with_stats(a int, b int);
 CREATE STATISTICS stx1 ON a, b FROM test_citus_local_table_with_stats;
-SELECT create_citus_local_table('test_citus_local_table_with_stats');
+ALTER TABLE test_citus_local_table_with_stats ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
- create_citus_local_table
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (xxxxx, 'citus_local_tables_test_schema', xxxxx, 'citus_local_tables_test_schema', 'ALTER TABLE test_citus_local_table_with_stats ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);')
----------------------------------------------------------------------
-
-(1 row)
-
 CREATE STATISTICS "CiTUS!LocalTables"."Bad\'StatName" ON a, b FROM test_citus_local_table_with_stats;
-NOTICE:  executing the command locally: SELECT worker_apply_shard_ddl_command (1504042, 'citus_local_tables_test_schema', E'CREATE STATISTICS "CiTUS!LocalTables"."Bad\\''StatName" ON a, b FROM citus_local_tables_test_schema.test_citus_local_table_with_stats')
+NOTICE:  executing the command locally: SELECT worker_apply_shard_ddl_command (1504043, 'citus_local_tables_test_schema', E'CREATE STATISTICS "CiTUS!LocalTables"."Bad\\''StatName" ON a, b FROM citus_local_tables_test_schema.test_citus_local_table_with_stats')
 SELECT stxname FROM pg_statistic_ext ORDER BY stxname;
         stxname
 ---------------------------------------------------------------------
  Bad\'StatName
- Bad\'StatName_1504042
+ Bad\'StatName_1504043
  stx1
- stx1_1504042
+ stx1_1504043
 (4 rows)
 
 -- observe the debug messages telling that we switch to sequential

src/test/regress/expected/fkeys_between_local_ref.out

@@ -114,6 +114,14 @@ BEGIN;
  t
 (1 row)
 
+  -- dropping that column would undistribute those 4 citus local tables
+  ALTER TABLE local_table_1 DROP COLUMN col_1 CASCADE;
+  SELECT COUNT(*)=0 FROM citus_local_tables_in_schema;
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
 ROLLBACK;
 -- this actually attempts to convert local tables to citus local tables but errors out
 -- as citus doesn't support defining foreign keys via add column commands
@@ -132,6 +140,25 @@ BEGIN;
  t
 (1 row)
 
+  -- dropping foreign key from local_table_2 would only undistribute local_table_2 & local_table_5
+  ALTER TABLE local_table_2 DROP CONSTRAINT fkey_1;
+  SELECT logicalrelid::regclass::text FROM citus_local_tables_in_schema ORDER BY logicalrelid;
+ logicalrelid
+---------------------------------------------------------------------
+ local_table_1
+ local_table_3
+ local_table_4
+(3 rows)
+
+    -- dropping local_table_1 would undistribute last two citus local tables as local_table_1
+    -- was the bridge to reference table
+  DROP TABLE local_table_1 CASCADE;
+  SELECT COUNT(*)=0 FROM citus_local_tables_in_schema;
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
 ROLLBACK;
 -- they fail as local_table_99 does not exist
 ALTER TABLE local_table_99 ADD CONSTRAINT fkey FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
@@ -331,6 +358,21 @@ BEGIN;
  reference_table_1                                    | n          | t
 (8 rows)
 
+  DROP TABLE local_table_3 CASCADE;
+  DROP SCHEMA another_schema_fkeys_between_local_ref CASCADE;
+  -- now we shouldn't see local_table_5 since now it is not connected to any reference tables
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref')
+  ORDER BY tablename;
+     tablename     | partmethod | repmodel
+---------------------------------------------------------------------
+ distributed_table | h          | c
+ local_table_1     | n          | c
+ local_table_2     | n          | c
+ local_table_4     | n          | c
+ reference_table_1 | n          | t
+(5 rows)
+
 ROLLBACK;
 BEGIN;
   CREATE TABLE local_table_6 (col_1 INT PRIMARY KEY);
@@ -391,6 +433,31 @@ BEGIN;
  reference_table_1 | n          | t
 (8 rows)
 
+  CREATE SCHEMA another_schema_fkeys_between_local_ref;
+  CREATE TABLE another_schema_fkeys_between_local_ref.reference_table_3 (col_1 INT UNIQUE);
+  SELECT create_reference_table('another_schema_fkeys_between_local_ref.reference_table_3');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  TRUNCATE local_table_4 CASCADE;
+  ALTER TABLE local_table_4 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES another_schema_fkeys_between_local_ref.reference_table_3(col_1);
+  DROP TABLE local_table_5 CASCADE;
+  ALTER TABLE local_table_2 DROP CONSTRAINT fkey_1;
+  DROP SCHEMA another_schema_fkeys_between_local_ref CASCADE;
+  -- now we shouldn't see any citus local tables
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref')
+  ORDER BY tablename;
+     tablename     | partmethod | repmodel
+---------------------------------------------------------------------
+ distributed_table | h          | c
+ local_table_2     | n          | t
+ local_table_6     | n          | t
+ reference_table_1 | n          | t
+(4 rows)
+
 ROLLBACK;
 BEGIN;
   -- disable foreign keys to reference tables

src/test/regress/expected/multi_extension.out

@@ -495,6 +495,8 @@ SELECT * FROM print_extension_changes();
                                                                                  | function citus_update_table_statistics(regclass)
                                                                                  | function columnar.columnar_handler(internal)
                                                                                  | function create_citus_local_table(regclass,boolean)
+                                                                                 | function notify_constraint_dropped()
+                                                                                 | function remove_local_tables_from_metadata()
                                                                                  | function time_partition_range(regclass)
                                                                                  | function undistribute_table(regclass,boolean)
                                                                                  | function worker_change_sequence_dependency(regclass,regclass,regclass)
@@ -506,7 +508,7 @@ SELECT * FROM print_extension_changes();
                                                                                  | view citus_shards
                                                                                  | view citus_tables
                                                                                  | view time_partitions
-(61 rows)
+(63 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/multi_extension_0.out

@@ -491,6 +491,8 @@ SELECT * FROM print_extension_changes();
                                                                                  | function citus_update_shard_statistics(bigint)
                                                                                  | function citus_update_table_statistics(regclass)
                                                                                  | function create_citus_local_table(regclass,boolean)
+                                                                                 | function notify_constraint_dropped()
+                                                                                 | function remove_local_tables_from_metadata()
                                                                                  | function time_partition_range(regclass)
                                                                                  | function undistribute_table(regclass,boolean)
                                                                                  | function worker_change_sequence_dependency(regclass,regclass,regclass)
@@ -502,7 +504,7 @@ SELECT * FROM print_extension_changes();
                                                                                  | view citus_shards
                                                                                  | view citus_tables
                                                                                  | view time_partitions
-(57 rows)
+(59 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/ref_citus_local_fkeys.out

@@ -52,19 +52,25 @@ NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.cit
 -- show that we support drop constraint
 ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;
 NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;')
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506000 citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 -- we support ON UPDATE CASCADE behaviour in "ALTER TABLE ADD fkey citus_local_table (to reference table)" commands
 ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1) ON UPDATE CASCADE;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1) ON UPDATE CASCADE;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506002, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1) ON UPDATE CASCADE;')
 -- show that on update cascade works
 INSERT INTO reference_table VALUES (12);
 NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506001 (r1) VALUES (12)
 INSERT INTO citus_local_table VALUES (12);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506000 (l1) VALUES (12)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506002 (l1) VALUES (12)
 UPDATE reference_table SET r1=13 WHERE r1=12;
 NOTICE:  executing the command locally: UPDATE ref_citus_local_fkeys.reference_table_1506001 reference_table SET r1 = 13 WHERE (r1 OPERATOR(pg_catalog.=) 12)
 -- should print a row with 13
 SELECT * FROM citus_local_table ORDER BY l1;
-NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506000 citus_local_table ORDER BY l1
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506002 citus_local_table ORDER BY l1
  l1
 ---------------------------------------------------------------------
  13
@@ -72,32 +78,43 @@ NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.cit
 
 -- drop constraint for next commands
 ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506002, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;')
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506002 citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 INSERT INTO citus_local_table VALUES (2);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506000 (l1) VALUES (2)
 -- show that we are checking for foreign key constraint while defining, below should fail
 ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506003, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
-ERROR:  insert or update on table "citus_local_table_1506000" violates foreign key constraint "fkey_local_to_ref_1506000"
+ERROR:  insert or update on table "citus_local_table_1506003" violates foreign key constraint "fkey_local_to_ref_1506003"
 INSERT INTO reference_table VALUES (2);
 NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506001 (r1) VALUES (2)
 -- this should work
 ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506004, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
 -- show that we are checking for foreign key constraint after defining, this should fail
 INSERT INTO citus_local_table VALUES (1);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506000 (l1) VALUES (1)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506004 (l1) VALUES (1)
-ERROR:  insert or update on table "citus_local_table_1506000" violates foreign key constraint "fkey_local_to_ref_1506000"
+ERROR:  insert or update on table "citus_local_table_1506004" violates foreign key constraint "fkey_local_to_ref_1506004"
 INSERT INTO reference_table VALUES (1);
 NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506001 (r1) VALUES (1)
 -- this should work
 INSERT INTO citus_local_table VALUES (1);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506000 (l1) VALUES (1)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506004 (l1) VALUES (1)
 -- drop and add constraint for next commands
 ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506004, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;')
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506004 citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506000, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506005, 'ref_citus_local_fkeys', 1506001, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1);')
 -- show that drop table without CASCADE errors out
 DROP TABLE reference_table;
 ERROR:  cannot drop table reference_table because other objects depend on it
@@ -106,12 +123,18 @@ BEGIN;
   DROP TABLE reference_table CASCADE;
 NOTICE:  drop cascades to constraint fkey_local_to_ref on table citus_local_table
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.reference_table_xxxxx CASCADE
-NOTICE:  drop cascades to constraint fkey_local_to_ref_1506000 on table ref_citus_local_fkeys.citus_local_table_1506000
+NOTICE:  drop cascades to constraint fkey_local_to_ref_1506005 on table ref_citus_local_fkeys.citus_local_table_1506005
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506005 citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 ROLLBACK;
 -- drop tables finally
 DROP TABLE citus_local_table, reference_table;
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.reference_table_xxxxx CASCADE
-NOTICE:  drop cascades to constraint fkey_local_to_ref_1506000 on table ref_citus_local_fkeys.citus_local_table_1506000
+NOTICE:  drop cascades to constraint fkey_local_to_ref_1506005 on table ref_citus_local_fkeys.citus_local_table_1506005
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
 ---------------------------------------------------------------------
 -- foreign key from reference table to citus local table --
@@ -141,11 +164,11 @@ SELECT create_reference_table('reference_table');
 (1 row)
 
 INSERT INTO reference_table VALUES (3);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506003 (r1) VALUES (3)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506007 (r1) VALUES (3)
 -- show that we are checking for foreign key constraint while defining, this should fail
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1);
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506003, 'ref_citus_local_fkeys', 1506002, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1);')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506007, 'ref_citus_local_fkeys', 1506006, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1);')
-ERROR:  insert or update on table "reference_table_1506003" violates foreign key constraint "fkey_ref_to_local_1506003"
+ERROR:  insert or update on table "reference_table_1506007" violates foreign key constraint "fkey_ref_to_local_1506007"
 -- we do not support CASCADE / SET NULL / SET DEFAULT behavior in "ALTER TABLE ADD fkey reference_table (to citus_local_table)" commands
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE CASCADE;
 ERROR:  cannot define foreign key constraint, foreign keys from reference tables to citus local tables can only be defined with NO ACTION or RESTRICT behaviors
@@ -160,25 +183,32 @@ ERROR:  cannot define foreign key constraint, foreign keys from reference tables
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON UPDATE SET DEFAULT;
 ERROR:  cannot define foreign key constraint, foreign keys from reference tables to citus local tables can only be defined with NO ACTION or RESTRICT behaviors
 INSERT INTO citus_local_table VALUES (3);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506002 (l1) VALUES (3)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.citus_local_table_1506006 (l1) VALUES (3)
 -- .. but we allow such foreign keys with RESTRICT behavior
 BEGIN;
   ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE RESTRICT;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506003, 'ref_citus_local_fkeys', 1506002, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE RESTRICT;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506007, 'ref_citus_local_fkeys', 1506006, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE RESTRICT;')
 ROLLBACK;
 -- .. and we allow such foreign keys with NO ACTION behavior
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506003, 'ref_citus_local_fkeys', 1506002, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506007, 'ref_citus_local_fkeys', 1506006, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;')
 -- show that adding/dropping foreign keys from reference to citus local
 -- tables works fine with remote execution too
 SET citus.enable_local_execution TO OFF;
 ALTER TABLE reference_table DROP CONSTRAINT fkey_ref_to_local;
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;
+ERROR:  cannot execute command because a local execution has accessed a placement in the transaction
 SET citus.enable_local_execution TO ON;
+ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506007, 'ref_citus_local_fkeys', 1506009, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;')
 -- show that we are checking for foreign key constraint after defining, this should fail
 INSERT INTO reference_table VALUES (4);
-NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506003 (r1) VALUES (4)
+NOTICE:  executing the command locally: INSERT INTO ref_citus_local_fkeys.reference_table_1506007 (r1) VALUES (4)
-ERROR:  insert or update on table "reference_table_1506003" violates foreign key constraint "fkey_ref_to_local_1506003"
+ERROR:  insert or update on table "reference_table_1506007" violates foreign key constraint "fkey_ref_to_local_1506007"
 -- enable the worker_2 to show that we don't try to set up the foreign keys
 -- between reference tables and citus local tables in worker_2 placements of
 -- the reference tables
@@ -192,7 +222,13 @@ NOTICE:  Replicating reference table "reference_table" to the node localhost:xxx
 -- show that we support drop constraint
 BEGIN;
   ALTER TABLE reference_table DROP CONSTRAINT fkey_ref_to_local;
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506003, 'ref_citus_local_fkeys', 1506002, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table DROP CONSTRAINT fkey_ref_to_local;')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506007, 'ref_citus_local_fkeys', 1506009, 'ref_citus_local_fkeys', 'ALTER TABLE reference_table DROP CONSTRAINT fkey_ref_to_local;')
+NOTICE:  creating a new table for ref_citus_local_fkeys.citus_local_table
+NOTICE:  Moving the data of ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: SELECT l1 FROM ref_citus_local_fkeys.citus_local_table_1506009 citus_local_table
+NOTICE:  Dropping the old ref_citus_local_fkeys.citus_local_table
+NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
+NOTICE:  Renaming the new table to ref_citus_local_fkeys.citus_local_table
 ROLLBACK;
 -- show that drop table errors as expected
 DROP TABLE citus_local_table;
@@ -201,7 +237,7 @@ ERROR:  cannot drop table citus_local_table because other objects depend on it
 DROP TABLE citus_local_table CASCADE;
 NOTICE:  drop cascades to constraint fkey_ref_to_local on table reference_table
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS ref_citus_local_fkeys.citus_local_table_xxxxx CASCADE
-NOTICE:  drop cascades to constraint fkey_ref_to_local_1506003 on table ref_citus_local_fkeys.reference_table_1506003
+NOTICE:  drop cascades to constraint fkey_ref_to_local_1506007 on table ref_citus_local_fkeys.reference_table_1506007
 BEGIN;
   CREATE TABLE citus_local_table_1(a int, b int, unique (a,b));
   CREATE TABLE citus_local_table_2(a int, b int, unique (a,b));
@@ -219,7 +255,7 @@ BEGIN;
 
   -- show that we properly handle multi column foreign keys
   ALTER TABLE citus_local_table_1 ADD CONSTRAINT multi_fkey FOREIGN KEY (a, b) REFERENCES citus_local_table_2(a, b);
-NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506004, 'ref_citus_local_fkeys', 1506005, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table_1 ADD CONSTRAINT multi_fkey FOREIGN KEY (a, b) REFERENCES citus_local_table_2(a, b);')
+NOTICE:  executing the command locally: SELECT worker_apply_inter_shard_ddl_command (1506010, 'ref_citus_local_fkeys', 1506011, 'ref_citus_local_fkeys', 'ALTER TABLE citus_local_table_1 ADD CONSTRAINT multi_fkey FOREIGN KEY (a, b) REFERENCES citus_local_table_2(a, b);')
 COMMIT;
 -- when local execution is disabled, citus local table cannot be created
 BEGIN;

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -151,12 +151,14 @@ ORDER BY 1;
  function master_update_node(integer,text,integer,boolean,integer)
  function master_update_shard_statistics(bigint)
  function master_update_table_statistics(regclass)
+ function notify_constraint_dropped()
  function poolinfo_valid(text)
  function read_intermediate_result(text,citus_copy_format)
  function read_intermediate_results(text[],citus_copy_format)
  function rebalance_table_shards(regclass,real,integer,bigint[],citus.shard_transfer_mode,boolean,name)
  function recover_prepared_transactions()
  function relation_is_a_known_shard(regclass)
+ function remove_local_tables_from_metadata()
  function replicate_reference_tables()
  function replicate_table_shards(regclass,integer,integer,bigint[],citus.shard_transfer_mode)
  function role_exists(name)
@@ -236,5 +238,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(220 rows)
+(222 rows)
 

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -147,12 +147,14 @@ ORDER BY 1;
  function master_update_node(integer,text,integer,boolean,integer)
  function master_update_shard_statistics(bigint)
  function master_update_table_statistics(regclass)
+ function notify_constraint_dropped()
  function poolinfo_valid(text)
  function read_intermediate_result(text,citus_copy_format)
  function read_intermediate_results(text[],citus_copy_format)
  function rebalance_table_shards(regclass,real,integer,bigint[],citus.shard_transfer_mode,boolean,name)
  function recover_prepared_transactions()
  function relation_is_a_known_shard(regclass)
+ function remove_local_tables_from_metadata()
  function replicate_reference_tables()
  function replicate_table_shards(regclass,integer,integer,bigint[],citus.shard_transfer_mode)
  function role_exists(name)
@@ -232,5 +234,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(216 rows)
+(218 rows)
 

src/test/regress/multi_schedule

@@ -319,7 +319,8 @@ test: multi_remove_node_reference_table
 # --------
 test: add_coordinator
 test: multi_reference_table citus_local_tables_queries
-test: foreign_key_to_reference_table citus_local_table_triggers
+test: foreign_key_to_reference_table
+test: citus_local_table_triggers
 test: replicate_reference_tables_to_coordinator
 test: coordinator_shouldhaveshards
 test: local_shard_utility_command_execution
@@ -328,6 +329,7 @@ test: multi_row_router_insert mixed_relkind_tests create_ref_dist_from_citus_loc
 test: undistribute_table_cascade
 test: create_citus_local_table_cascade
 test: fkeys_between_local_ref
+test: auto_undist_citus_local
 
 test: remove_coordinator
 

src/test/regress/sql/auto_undist_citus_local.sql

@@ -0,0 +1,320 @@
+-- regression tests regarding foreign key
+-- drops cascading into undistributing Citus
+-- local tables to Postgres local tables
+CREATE SCHEMA drop_fkey_cascade;
+SET search_path TO drop_fkey_cascade;
+SET client_min_messages TO WARNING;
+SET citus.next_shard_id TO 1810000;
+
+SELECT 1 FROM master_add_node('localhost', :master_port, groupId => 0);
+
+-- show that DROP CONSTRAINT cascades to undistributing citus_local_table
+CREATE TABLE citus_local_table(l1 int);
+SELECT create_citus_local_table('citus_local_table');
+CREATE TABLE reference_table(r1 int primary key);
+SELECT create_reference_table('reference_table');
+ALTER TABLE citus_local_table ADD CONSTRAINT fkey_local_to_ref FOREIGN KEY(l1) REFERENCES reference_table(r1) ON DELETE CASCADE;
+
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+CREATE OR REPLACE FUNCTION drop_constraint_cascade_via_perform_deletion(IN table_name regclass, IN constraint_name text)
+RETURNS VOID
+LANGUAGE C STRICT
+AS 'citus', $$drop_constraint_cascade_via_perform_deletion$$;
+
+BEGIN;
+  SELECT drop_constraint_cascade_via_perform_deletion('citus_local_table', 'fkey_local_to_ref');
+  -- we dropped constraint without going through utility hook,
+  -- so we should still see citus_local_table
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+ROLLBACK;
+
+ALTER TABLE citus_local_table DROP CONSTRAINT fkey_local_to_ref;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+DROP TABLE citus_local_table, reference_table;
+
+-- show that DROP COLUMN cascades to undistributing citus_local_table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+ALTER TABLE reference_table DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+DROP TABLE citus_local_table, reference_table;
+
+-- show that DROP COLUMN that cascades into drop foreign key undistributes local table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+ALTER TABLE citus_local_table DROP COLUMN l1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+
+DROP TABLE citus_local_table, reference_table;
+
+-- show that PRIMARY KEY that cascades into drop foreign key undistributes local table
+CREATE TABLE reference_table(r1 int primary key, r2 int);
+SELECT create_reference_table('reference_table');
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_pkey CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+-- show that DROP UNIQUE INDEX that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+
+CREATE TABLE reference_table(r1 int, r2 int);
+SELECT create_reference_table('reference_table');
+CREATE UNIQUE INDEX ref_unique ON reference_table(r1);
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+DROP INDEX ref_unique CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+-- show that UNIQUE CONSTRAINT that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_r1_key  CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+-- show that DROP TABLE that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+DROP TABLE reference_table CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+
+
+-- show that UNIQUE CONSTRAINT that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table;
+
+CREATE TABLE reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table');
+
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+ALTER TABLE reference_table DROP CONSTRAINT reference_table_r1_key  CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table'::regclass) ORDER BY logicalrelid;
+
+
+-- show that DROP SCHEMA that cascades into drop foreign key undistributes local table
+DROP TABLE citus_local_table, reference_table;
+
+CREATE SCHEMA ref_table_drop_schema;
+CREATE TABLE ref_table_drop_schema.reference_table(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('ref_table_drop_schema.reference_table');
+CREATE TABLE citus_local_table(l1 int REFERENCES ref_table_drop_schema.reference_table(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'ref_table_drop_schema.reference_table'::regclass) ORDER BY logicalrelid;
+DROP SCHEMA ref_table_drop_schema CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+
+
+-- drop column cascade that doesn't cascade into citus local table
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE reference_table_2(r1 int UNIQUE REFERENCES reference_table_1(r1), r2 int);
+SELECT create_reference_table('reference_table_2');
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_2(r1), l2 int);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+ALTER TABLE reference_table_1 DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+
+
+-- local table has multiple foreign keys to two tables
+-- drop one at a time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+
+DROP TABLE reference_table_1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+
+CREATE TABLE distributed_table (d1 int);
+SELECT create_distributed_table('distributed_table', 'd1');
+
+-- drop an unrelated distributed table too
+DROP TABLE reference_table_2, distributed_table CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+
+-- local table has multiple foreign keys to two tables
+-- drop both at the same time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+DROP TABLE reference_table_1, reference_table_2 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass) ORDER BY logicalrelid;
+
+-- local table has multiple foreign keys to two tables
+-- drop one at a time
+DROP TABLE IF EXISTS citus_local_table, reference_table_1, reference_table_2;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE reference_table_2(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_2');
+CREATE TABLE citus_local_table(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES reference_table_2(r1));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+
+BEGIN;
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l1_fkey;
+  SAVEPOINT sp1;
+
+  -- this should undistribute citus_local_table
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l2_fkey;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+  ROLLBACK TO SAVEPOINT sp1;
+
+  -- rollback'ed second drop constraint, so we should still see citus_local_table
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+
+  -- this should undistribute citus_local_table again
+  ALTER TABLE citus_local_table DROP CONSTRAINT citus_local_table_l2_fkey;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('citus_local_table'::regclass, 'reference_table_1'::regclass, 'reference_table_2'::regclass) ORDER BY logicalrelid;
+COMMIT;
+
+-- a single drop column cascades into multiple undistributes
+DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, reference_table_1;
+
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+CREATE TABLE citus_local_table_2(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE REFERENCES citus_local_table_1(l2));
+CREATE TABLE citus_local_table_3(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES citus_local_table_2(l2));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+ALTER TABLE reference_table_1 DROP COLUMN r1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+
+-- a single drop table cascades into multiple undistributes
+DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, citus_local_table_3, citus_local_table_2, reference_table_1;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+CREATE TABLE citus_local_table_2(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE REFERENCES citus_local_table_1(l2));
+CREATE TABLE citus_local_table_3(l1 int REFERENCES reference_table_1(r1), l2 int REFERENCES citus_local_table_2(l2));
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+
+-- test DROP OWNED BY
+
+-- Citus does not support "ALTER TABLE OWNER TO" commands. Also, not to deal with tests output
+-- difference between community and enterprise, let's disable enable_ddl_propagation here.
+SET citus.enable_ddl_propagation to OFF;
+
+CREATE USER another_user;
+SELECT run_command_on_workers('CREATE USER another_user');
+
+ALTER TABLE reference_table_1 OWNER TO another_user;
+SELECT run_command_on_placements('reference_table_1', 'ALTER TABLE %s OWNER TO another_user');
+
+SET citus.enable_ddl_propagation to ON;
+
+BEGIN;
+  DROP OWNED BY another_user cascade;
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ( 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+ROLLBACK;
+
+DROP TABLE reference_table_1 CASCADE;
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ( 'citus_local_table_1'::regclass, 'citus_local_table_2'::regclass, 'citus_local_table_3'::regclass) ORDER BY logicalrelid;
+
+
+-- dropping constraints inside a plpgsql procedure should be fine
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+
+CREATE OR REPLACE FUNCTION drop_constraint_via_func()
+RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+ALTER TABLE citus_local_table_1 DROP CONSTRAINT citus_local_table_1_l1_fkey;
+END;$$;
+
+BEGIN;
+  SELECT drop_constraint_via_func();
+  SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+ROLLBACK;
+
+create or replace procedure drop_constraint_via_proc()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+        ALTER TABLE citus_local_table_1 DROP CONSTRAINT citus_local_table_1_l1_fkey;
+    commit;
+end;$$;
+call drop_constraint_via_proc();
+
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+-- even if the procedure is called from another procedure
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+create or replace procedure drop_constraint_via_proc_top_level()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+        CALL drop_constraint_via_proc();
+    commit;
+end;$$;
+
+CALL drop_constraint_via_proc_top_level();
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+
+-- even if the procedure is called from an exception handler
+DROP TABLE IF EXISTS citus_local_table_1, reference_table_1 CASCADE;
+CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
+SELECT create_reference_table('reference_table_1');
+CREATE TABLE citus_local_table_1(l1 int REFERENCES reference_table_1(r1), l2 int UNIQUE);
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+create or replace procedure drop_constraint_via_proc_exception()
+language plpgsql
+as $$
+DECLARE
+    res INT := 0;
+begin
+		PERFORM 1/0;
+    	EXCEPTION
+             when others then
+        CALL drop_constraint_via_proc();
+    commit;
+end;$$;
+
+CALL drop_constraint_via_proc_exception();
+SELECT logicalrelid, partmethod, repmodel FROM pg_dist_partition WHERE logicalrelid IN ('reference_table_1'::regclass, 'citus_local_table_1'::regclass) ORDER BY logicalrelid;
+
+DROP SCHEMA drop_fkey_cascade CASCADE;

src/test/regress/sql/citus_local_tables.sql

@@ -201,6 +201,8 @@ CREATE FOREIGN TABLE foreign_table (
 -- & shell relation points to the same same server object
 SELECT create_citus_local_table('foreign_table');
 
+DROP FOREIGN TABLE foreign_table;
+
 -- drop them for next tests
 DROP TABLE citus_local_table_1, citus_local_table_2, distributed_table;
 
@@ -314,6 +316,41 @@ CREATE UNIQUE INDEX uniqueIndex2 ON "LocalTabLE.1!?!"(id);
 
 SET search_path TO citus_local_tables_test_schema;
 
+CREATE TABLE dummy_reference_table (a INT PRIMARY KEY);
+SELECT create_reference_table('dummy_reference_table');
+
+BEGIN;
+  SET client_min_messages TO ERROR;
+  SELECT remove_local_tables_from_metadata();
+
+  -- should not see any citus local tables
+  SELECT logicalrelid::regclass::text FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='citus_local_tables_test_schema' AND
+        partmethod = 'n' AND repmodel = 'c'
+  ORDER BY 1;
+ROLLBACK;
+
+-- define foreign keys between dummy_reference_table and citus local tables
+-- not to undistribute them automatically
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+ALTER TABLE citus_local_table_2 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+ALTER TABLE unlogged_table ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
+ALTER TABLE dummy_reference_table ADD CONSTRAINT fkey_from_dummy_ref FOREIGN KEY (a) REFERENCES "CiTUS!LocalTables"."LocalTabLE.1!?!"(id);
+
+BEGIN;
+  SET client_min_messages TO ERROR;
+  SELECT remove_local_tables_from_metadata();
+
+  -- now we defined foreign keys with above citus local tables, we should still see them
+  SELECT logicalrelid::regclass::text FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='citus_local_tables_test_schema' AND
+        partmethod = 'n' AND repmodel = 'c'
+  ORDER BY 1;
+ROLLBACK;
+
 -- between citus local tables and distributed tables
 ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_c_to_dist FOREIGN KEY(a) references distributed_table(a);
 ALTER TABLE distributed_table ADD CONSTRAINT fkey_dist_to_c FOREIGN KEY(a) references citus_local_table_1(a);
@@ -450,7 +487,7 @@ ALTER TABLE referencing_table ADD CONSTRAINT fkey_cl_to_cl FOREIGN KEY (a) REFER
 -- verify creating citus local table with extended statistics
 CREATE TABLE test_citus_local_table_with_stats(a int, b int);
 CREATE STATISTICS stx1 ON a, b FROM test_citus_local_table_with_stats;
-SELECT create_citus_local_table('test_citus_local_table_with_stats');
+ALTER TABLE test_citus_local_table_with_stats ADD CONSTRAINT fkey_to_dummy_ref FOREIGN KEY (a) REFERENCES dummy_reference_table(a);
 CREATE STATISTICS "CiTUS!LocalTables"."Bad\'StatName" ON a, b FROM test_citus_local_table_with_stats;
 SELECT stxname FROM pg_statistic_ext ORDER BY stxname;
 

src/test/regress/sql/fkeys_between_local_ref.sql

@@ -95,6 +95,10 @@ BEGIN;
 
   -- show that we converted all 4 local tables in this schema to citus local tables
   SELECT COUNT(*)=4 FROM citus_local_tables_in_schema;
+
+  -- dropping that column would undistribute those 4 citus local tables
+  ALTER TABLE local_table_1 DROP COLUMN col_1 CASCADE;
+  SELECT COUNT(*)=0 FROM citus_local_tables_in_schema;
 ROLLBACK;
 
 -- this actually attempts to convert local tables to citus local tables but errors out
@@ -111,6 +115,15 @@ BEGIN;
 
   -- now we have 5 citus local tables in this schema
   SELECT COUNT(*)=5 FROM citus_local_tables_in_schema;
+
+  -- dropping foreign key from local_table_2 would only undistribute local_table_2 & local_table_5
+  ALTER TABLE local_table_2 DROP CONSTRAINT fkey_1;
+  SELECT logicalrelid::regclass::text FROM citus_local_tables_in_schema ORDER BY logicalrelid;
+
+    -- dropping local_table_1 would undistribute last two citus local tables as local_table_1
+    -- was the bridge to reference table
+  DROP TABLE local_table_1 CASCADE;
+  SELECT COUNT(*)=0 FROM citus_local_tables_in_schema;
 ROLLBACK;
 
 -- they fail as local_table_99 does not exist
@@ -252,6 +265,14 @@ BEGIN;
   WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref' UNION
                                SELECT 'another_schema_fkeys_between_local_ref.local_table_6')
   ORDER BY tablename;
+
+  DROP TABLE local_table_3 CASCADE;
+  DROP SCHEMA another_schema_fkeys_between_local_ref CASCADE;
+
+  -- now we shouldn't see local_table_5 since now it is not connected to any reference tables
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref')
+  ORDER BY tablename;
 ROLLBACK;
 
 BEGIN;
@@ -283,6 +304,21 @@ BEGIN;
   SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
   WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref')
   ORDER BY tablename;
+
+  CREATE SCHEMA another_schema_fkeys_between_local_ref;
+  CREATE TABLE another_schema_fkeys_between_local_ref.reference_table_3 (col_1 INT UNIQUE);
+  SELECT create_reference_table('another_schema_fkeys_between_local_ref.reference_table_3');
+  TRUNCATE local_table_4 CASCADE;
+  ALTER TABLE local_table_4 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES another_schema_fkeys_between_local_ref.reference_table_3(col_1);
+
+  DROP TABLE local_table_5 CASCADE;
+  ALTER TABLE local_table_2 DROP CONSTRAINT fkey_1;
+  DROP SCHEMA another_schema_fkeys_between_local_ref CASCADE;
+
+  -- now we shouldn't see any citus local tables
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='fkeys_between_local_ref')
+  ORDER BY tablename;
 ROLLBACK;
 
 BEGIN;

src/test/regress/sql/ref_citus_local_fkeys.sql

@@ -120,6 +120,8 @@ ALTER TABLE reference_table DROP CONSTRAINT fkey_ref_to_local;
 ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;
 SET citus.enable_local_execution TO ON;
 
+ALTER TABLE reference_table ADD CONSTRAINT fkey_ref_to_local FOREIGN KEY(r1) REFERENCES citus_local_table(l1) ON DELETE NO ACTION;
+
 -- show that we are checking for foreign key constraint after defining, this should fail
 INSERT INTO reference_table VALUES (4);