Always CASCADE while dropping a shard

src/backend/distributed/master/master_expire_table_cache.c

@@ -145,6 +145,7 @@ DropShardsFromWorker(WorkerNode *workerNode, Oid relationId, List *shardInterval
 	char *relationName = get_rel_name(relationId);
 	char relationKind = get_rel_relkind(relationId);
 	StringInfo workerCommand = makeStringInfo();
+	StringInfo shardNames = makeStringInfo();
 	ListCell *shardIntervalCell = NULL;
 
 	if (shardIntervalList == NIL)
@@ -152,20 +153,6 @@ DropShardsFromWorker(WorkerNode *workerNode, Oid relationId, List *shardInterval
 		return;
 	}
 
-	if (relationKind == RELKIND_RELATION)
-	{
-		appendStringInfo(workerCommand, DROP_REGULAR_TABLE_COMMAND, "");
-	}
-	else if (relationKind == RELKIND_FOREIGN_TABLE)
-	{
-		appendStringInfo(workerCommand, DROP_FOREIGN_TABLE_COMMAND, "");
-	}
-	else
-	{
-		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-						errmsg("expire target is not a regular or foreign table")));
-	}
-
 	foreach(shardIntervalCell, shardIntervalList)
 	{
 		ShardInterval *shardInterval = (ShardInterval *) lfirst(shardIntervalCell);
@@ -174,16 +161,28 @@ DropShardsFromWorker(WorkerNode *workerNode, Oid relationId, List *shardInterval
 
 		AppendShardIdToName(&shardName, shardInterval->shardId);
 		quotedShardName = quote_qualified_identifier(schemaName, shardName);
-		appendStringInfo(workerCommand, "%s", quotedShardName);
+		appendStringInfo(shardNames, "%s", quotedShardName);
 
 		/* append a comma after the shard name if there are more shards */
 		if (lnext(shardIntervalCell) != NULL)
 		{
-			appendStringInfo(workerCommand, ", ");
+			appendStringInfo(shardNames, ", ");
 		}
 	}
 
-	appendStringInfo(workerCommand, " CASCADE");
+	if (relationKind == RELKIND_RELATION)
+	{
+		appendStringInfo(workerCommand, DROP_REGULAR_TABLE_COMMAND, shardNames->data);
+	}
+	else if (relationKind == RELKIND_FOREIGN_TABLE)
+	{
+		appendStringInfo(workerCommand, DROP_FOREIGN_TABLE_COMMAND, shardNames->data);
+	}
+	else
+	{
+		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						errmsg("expire target is not a regular or foreign table")));
+	}
 
 	ExecuteRemoteCommand(workerNode->workerName, workerNode->workerPort, workerCommand);
 }

src/include/distributed/master_protocol.h

@@ -64,8 +64,8 @@
 #define SHARD_RANGE_QUERY "SELECT min(%s), max(%s) FROM %s"
 #define SHARD_TABLE_SIZE_QUERY "SELECT pg_table_size(%s)"
 #define SHARD_CSTORE_TABLE_SIZE_QUERY "SELECT cstore_table_size(%s)"
-#define DROP_REGULAR_TABLE_COMMAND "DROP TABLE IF EXISTS %s"
-#define DROP_FOREIGN_TABLE_COMMAND "DROP FOREIGN TABLE IF EXISTS %s"
+#define DROP_REGULAR_TABLE_COMMAND "DROP TABLE IF EXISTS %s CASCADE"
+#define DROP_FOREIGN_TABLE_COMMAND "DROP FOREIGN TABLE IF EXISTS %s CASCADE"
 #define CREATE_SCHEMA_COMMAND "CREATE SCHEMA IF NOT EXISTS %s"
 #define CREATE_EMPTY_SHARD_QUERY "SELECT master_create_empty_shard('%s')"
 #define FINALIZED_SHARD_PLACEMENTS_QUERY \

src/test/regress/expected/multi_create_table_constraints.out

@@ -484,3 +484,50 @@ DROP TABLE pk_on_part_col, uq_part_col, uq_two_columns CASCADE;
 DROP TABLE ex_on_part_col, ex_on_two_columns, ex_on_two_columns_prt, ex_multiple_excludes, ex_overlaps CASCADE;
 DROP TABLE ex_on_part_col_named, ex_on_two_columns_named, ex_overlaps_named CASCADE;
 DROP TABLE uq_range_tables, check_example CASCADE;
+-- test dropping table with foreign keys
+SET citus.shard_count = 4;
+SET citus.shard_replication_factor = 1;
+CREATE TABLE raw_table_1 (user_id int, UNIQUE(user_id));
+SELECT create_distributed_table('raw_table_1', 'user_id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+CREATE TABLE raw_table_2 (user_id int REFERENCES raw_table_1(user_id), UNIQUE(user_id));
+SELECT create_distributed_table('raw_table_2', 'user_id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+-- see that the constraint exists
+\d raw_table_2
+  Table "public.raw_table_2"
+ Column  |  Type   | Modifiers 
+---------+---------+-----------
+ user_id | integer | 
+Indexes:
+    "raw_table_2_user_id_key" UNIQUE CONSTRAINT, btree (user_id)
+Foreign-key constraints:
+    "raw_table_2_user_id_fkey" FOREIGN KEY (user_id) REFERENCES raw_table_1(user_id)
+
+-- should be prevented by the foreign key
+DROP TABLE raw_table_1;
+ERROR:  cannot drop table raw_table_1 because other objects depend on it
+DETAIL:  constraint raw_table_2_user_id_fkey on table raw_table_2 depends on table raw_table_1
+HINT:  Use DROP ... CASCADE to drop the dependent objects too.
+-- should cleanly drop the remote shards
+DROP TABLE raw_table_1 CASCADE;
+NOTICE:  drop cascades to constraint raw_table_2_user_id_fkey on table raw_table_2
+-- see that the constraint also dropped
+\d raw_table_2
+  Table "public.raw_table_2"
+ Column  |  Type   | Modifiers 
+---------+---------+-----------
+ user_id | integer | 
+Indexes:
+    "raw_table_2_user_id_key" UNIQUE CONSTRAINT, btree (user_id)
+
+-- drop the table as well
+DROP TABLE raw_table_2;

src/test/regress/sql/multi_create_table_constraints.sql

@@ -249,3 +249,28 @@ DROP TABLE pk_on_part_col, uq_part_col, uq_two_columns CASCADE;
 DROP TABLE ex_on_part_col, ex_on_two_columns, ex_on_two_columns_prt, ex_multiple_excludes, ex_overlaps CASCADE;
 DROP TABLE ex_on_part_col_named, ex_on_two_columns_named, ex_overlaps_named CASCADE;
 DROP TABLE uq_range_tables, check_example CASCADE;
+
+-- test dropping table with foreign keys
+SET citus.shard_count = 4;
+SET citus.shard_replication_factor = 1;
+
+CREATE TABLE raw_table_1 (user_id int, UNIQUE(user_id));
+SELECT create_distributed_table('raw_table_1', 'user_id');
+
+CREATE TABLE raw_table_2 (user_id int REFERENCES raw_table_1(user_id), UNIQUE(user_id));
+SELECT create_distributed_table('raw_table_2', 'user_id');
+
+-- see that the constraint exists
+\d raw_table_2
+
+-- should be prevented by the foreign key
+DROP TABLE raw_table_1;
+
+-- should cleanly drop the remote shards
+DROP TABLE raw_table_1 CASCADE;
+
+-- see that the constraint also dropped
+\d raw_table_2
+
+-- drop the table as well
+DROP TABLE raw_table_2;