Add cascade_via_foreign_keys option to create_citus_local_table (#4462)

src/backend/distributed/commands/cascade_table_operation_for_connected_relations.c

@@ -326,6 +326,11 @@ GetCascadeOperationFunction(CascadeOperationType cascadeOperationType)
 			return UndistributeTable;
 		}
 
+		case CREATE_CITUS_LOCAL_TABLE:
+		{
+			return CreateCitusLocalTable;
+		}
+
 		default:
 		{
 			/*

src/backend/distributed/commands/create_citus_local_table.c

@@ -38,7 +38,6 @@
 #include "utils/syscache.h"
 
 
-static void CreateCitusLocalTable(Oid relationId);
 static void ErrorIfUnsupportedCreateCitusLocalTable(Relation relation);
 static void ErrorIfUnsupportedCitusLocalTableKind(Oid relationId);
 static List * GetShellTableDDLEventsForCitusLocalTable(Oid relationId);
@@ -80,8 +79,9 @@ create_citus_local_table(PG_FUNCTION_ARGS)
 	CheckCitusVersion(ERROR);
 
 	Oid relationId = PG_GETARG_OID(0);
+	bool cascadeViaForeignKeys = PG_GETARG_BOOL(1);
 
-	CreateCitusLocalTable(relationId);
+	CreateCitusLocalTable(relationId, cascadeViaForeignKeys);
 
 	PG_RETURN_VOID();
 }
@@ -98,8 +98,8 @@ create_citus_local_table(PG_FUNCTION_ARGS)
  * Similar to reference tables, it has only 1 placement. In addition to that, that
  * single placement is only allowed to be on the coordinator.
  */
-static void
+void
-CreateCitusLocalTable(Oid relationId)
+CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
 {
 	/*
 	 * These checks should be done before acquiring any locks on relation.
@@ -118,7 +118,8 @@ CreateCitusLocalTable(Oid relationId)
 	 * we open the relation with try_relation_open instead of relation_open
 	 * to give a nice error in case the table is dropped by another backend.
 	 */
-	Relation relation = try_relation_open(relationId, AccessExclusiveLock);
+	LOCKMODE lockMode = AccessExclusiveLock;
+	Relation relation = try_relation_open(relationId, lockMode);
 
 	ErrorIfUnsupportedCreateCitusLocalTable(relation);
 
@@ -131,6 +132,36 @@ CreateCitusLocalTable(Oid relationId)
 	 */
 	relation_close(relation, NoLock);
 
+	bool tableHasExternalForeignKeys = TableHasExternalForeignKeys(relationId);
+	if (tableHasExternalForeignKeys && cascadeViaForeignKeys)
+	{
+		CascadeOperationForConnectedRelations(relationId, lockMode,
+											  CREATE_CITUS_LOCAL_TABLE);
+
+		/*
+		 * We converted every foreign key connected table in our subgraph
+		 * including itself to a citus local table, so return here.
+		 */
+		return;
+	}
+	else if (tableHasExternalForeignKeys)
+	{
+		/*
+		 * We do not allow creating citus local table if the table is involved in a
+		 * foreign key relationship with "any other table". Note that we allow self
+		 * references.
+		 */
+		char *qualifiedRelationName = generate_qualified_relation_name(relationId);
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("relation %s is involved in a foreign key "
+							   "relationship with another table", qualifiedRelationName),
+						errhint("Use cascade_via_foreign_keys option to convert "
+								"all the relations involved in a foreign key "
+								"relationship with %s to a citus local table by "
+								"executing SELECT create_citus_local_table($$%s$$, "
+								"cascade_via_foreign_keys=>true)",
+								qualifiedRelationName, qualifiedRelationName)));
+	}
 
 	ObjectAddress tableAddress = { 0 };
 	ObjectAddressSet(tableAddress, RelationRelationId, relationId);
@@ -217,13 +248,6 @@ ErrorIfUnsupportedCreateCitusLocalTable(Relation relation)
 	 */
 	ErrorIfRelationIsAKnownShard(relationId);
 
-	/*
-	 * We do not allow creating citus local table if the table is involved in a
-	 * foreign key relationship with "any other table". Note that we allow self
-	 * references.
-	 */
-	ErrorIfTableHasExternalForeignKeys(relationId);
-
 	/* we do not support policies in citus community */
 	ErrorIfUnsupportedPolicy(relation);
 }

src/backend/distributed/commands/foreign_constraint.c

@@ -825,11 +825,11 @@ FindForeignKeyOidWithName(List *foreignKeyOids, const char *inputConstraintName)
 
 
 /*
- * ErrorIfTableHasExternalForeignKeys errors out if the relation with relationId
+ * TableHasExternalForeignKeys returns true if the relation with relationId is
- * is involved in a foreign key relationship other than the self-referencing ones.
+ * involved in a foreign key relationship other than the self-referencing ones.
  */
-void
+bool
-ErrorIfTableHasExternalForeignKeys(Oid relationId)
+TableHasExternalForeignKeys(Oid relationId)
 {
 	int flags = (INCLUDE_REFERENCING_CONSTRAINTS | EXCLUDE_SELF_REFERENCES |
 				 INCLUDE_ALL_TABLE_TYPES);
@@ -844,16 +844,10 @@ ErrorIfTableHasExternalForeignKeys(Oid relationId)
 
 	if (list_length(foreignKeysWithOtherTables) == 0)
 	{
-		return;
+		return false;
 	}
 
-	const char *relationName = get_rel_name(relationId);
+	return true;
-	ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("relation \"%s\" is involved in a foreign key relationship "
-						   "with another table", relationName),
-					errhint("Drop foreign keys with other tables and re-define them "
-							"with ALTER TABLE commands after the current operation "
-							"is done.")));
 }
 
 

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

@@ -7,5 +7,6 @@ DROP FUNCTION IF EXISTS pg_catalog.citus_total_relation_size(regclass);
 #include "udfs/citus_tables/10.0-1.sql"
 #include "udfs/citus_finish_pg_upgrade/10.0-1.sql"
 #include "udfs/undistribute_table/10.0-1.sql"
+#include "udfs/create_citus_local_table/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

@@ -8,7 +8,9 @@
 DROP VIEW public.citus_tables;
 DROP FUNCTION pg_catalog.citus_total_relation_size(regclass,boolean);
 DROP FUNCTION pg_catalog.undistribute_table(regclass,boolean);
+DROP FUNCTION pg_catalog.create_citus_local_table(regclass,boolean);
 
 #include "../udfs/citus_total_relation_size/7.0-1.sql"
 #include "../udfs/upgrade_to_reference_table/8.0-1.sql"
 #include "../udfs/undistribute_table/9.5-1.sql"
+#include "../udfs/create_citus_local_table/9.5-1.sql"

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

@@ -0,0 +1,7 @@
+DROP FUNCTION pg_catalog.create_citus_local_table(regclass);
+CREATE OR REPLACE FUNCTION pg_catalog.create_citus_local_table(table_name regclass, cascade_via_foreign_keys boolean default false)
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'MODULE_PATHNAME', $$create_citus_local_table$$;
+COMMENT ON FUNCTION pg_catalog.create_citus_local_table(table_name regclass, cascade_via_foreign_keys boolean)
+	IS 'create a citus local table';

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

@@ -1,6 +1,7 @@
-CREATE OR REPLACE FUNCTION pg_catalog.create_citus_local_table(table_name regclass)
+DROP FUNCTION pg_catalog.create_citus_local_table(regclass);
+CREATE OR REPLACE FUNCTION pg_catalog.create_citus_local_table(table_name regclass, cascade_via_foreign_keys boolean default false)
 	RETURNS void
 	LANGUAGE C STRICT
 	AS 'MODULE_PATHNAME', $$create_citus_local_table$$;
-COMMENT ON FUNCTION pg_catalog.create_citus_local_table(table_name regclass)
+COMMENT ON FUNCTION pg_catalog.create_citus_local_table(table_name regclass, cascade_via_foreign_keys boolean)
 	IS 'create a citus local table';

src/include/distributed/commands.h

@@ -171,7 +171,7 @@ extern bool ConstraintIsAForeignKey(char *inputConstaintName, Oid relationOid);
 extern bool ConstraintWithNameIsOfType(char *inputConstaintName, Oid relationId,
 									   char targetConstraintType);
 extern bool ConstraintWithIdIsOfType(Oid constraintId, char targetConstraintType);
-extern void ErrorIfTableHasExternalForeignKeys(Oid relationId);
+extern bool TableHasExternalForeignKeys(Oid relationId);
 extern List * GetForeignKeyOids(Oid relationId, int flags);
 extern Oid GetReferencedTableId(Oid foreignKeyId);
 
@@ -387,7 +387,7 @@ extern List * CitusLocalTableTriggerCommandDDLJob(Oid relationId, char *triggerN
 												  const char *queryString);
 extern Oid GetTriggerFunctionId(Oid triggerId);
 
-/* cascade_citus_table_function.c */
+/* cascade_table_operation_for_connected_relations.c */
 
 /*
  * Flags that can be passed to CascadeOperationForConnectedRelations to specify
@@ -399,6 +399,9 @@ typedef enum CascadeOperationType
 
 	/* execute UndistributeTable on each relation */
 	UNDISTRIBUTE_TABLE = 1 << 1,
+
+	/* execute CreateCitusLocalTable on each relation */
+	CREATE_CITUS_LOCAL_TABLE = 1 << 2,
 } CascadeOperationType;
 
 extern void CascadeOperationForConnectedRelations(Oid relationId, LOCKMODE relLockMode,
@@ -407,6 +410,9 @@ extern void CascadeOperationForConnectedRelations(Oid relationId, LOCKMODE relLo
 extern void ExecuteAndLogDDLCommandList(List *ddlCommandList);
 extern void ExecuteAndLogDDLCommand(const char *commandString);
 
+/* create_citus_local_table.c */
+extern void CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys);
+
 extern bool ShouldPropagateSetCommand(VariableSetStmt *setStmt);
 extern void PostprocessVariableSetStmt(VariableSetStmt *setStmt, const char *setCommand);
 

src/test/regress/expected/citus_local_tables.out

@@ -349,9 +349,9 @@ CREATE TABLE local_table_3 (a int primary key, b int references local_table_3(a)
 -- below two should fail as we do not allow foreign keys between
 -- postgres local tables and citus local tables
 SELECT create_citus_local_table('local_table_1');
-ERROR:  relation "local_table_1" is involved in a foreign key relationship with another table
+ERROR:  relation citus_local_tables_test_schema.local_table_1 is involved in a foreign key relationship with another table
 SELECT create_citus_local_table('local_table_2');
-ERROR:  relation "local_table_2" is involved in a foreign key relationship with another table
+ERROR:  relation citus_local_tables_test_schema.local_table_2 is involved in a foreign key relationship with another table
 -- below should work as we allow initial self references in citus local tables
 SELECT create_citus_local_table('local_table_3');
  create_citus_local_table

src/test/regress/expected/create_citus_local_table_cascade.out

@@ -0,0 +1,224 @@
+\set VERBOSITY terse
+SET citus.next_shard_id TO 1516000;
+SET citus.shard_replication_factor TO 1;
+CREATE SCHEMA create_citus_local_table_cascade;
+SET search_path TO create_citus_local_table_cascade;
+SET client_min_messages to ERROR;
+-- ensure that coordinator is added to pg_dist_node
+SELECT 1 FROM master_add_node('localhost', :master_port, groupId => 0);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+CREATE TABLE local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE local_table_3 (col_1 INT UNIQUE);
+CREATE TABLE local_table_4 (col_1 INT UNIQUE);
+--                                        _
+--                                       | |
+--                                       | v
+-- local_table_2 -> local_table_1 -> local_table_4
+--                      ^ |               |
+--                      | v               |
+--                  local_table_3 <--------
+ALTER TABLE local_table_2 ADD CONSTRAINT fkey_1 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES local_table_3(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES local_table_4(col_1);
+ALTER TABLE local_table_4 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES local_table_3(col_1);
+ALTER TABLE local_table_4 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES local_table_4(col_1);
+-- show that all of below fails as we didn't provide cascade_via_foreign_keys=true
+SELECT create_citus_local_table('local_table_1');
+ERROR:  relation create_citus_local_table_cascade.local_table_1 is involved in a foreign key relationship with another table
+SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>false);
+ERROR:  relation create_citus_local_table_cascade.local_table_4 is involved in a foreign key relationship with another table
+-- In each of below two transaction blocks, show that we preserve foreign keys.
+-- Also show that we converted all local_table_xxx tables in current schema
+-- to citus local tables after create_citus_local_table (cascade).
+-- So in each transaction, both selects should return true.
+BEGIN;
+  SELECT conname, conrelid::regclass::text, confrelid::regclass::text
+  FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$'
+  ORDER BY 1,2,3;
+ conname |   conrelid    |   confrelid
+---------------------------------------------------------------------
+ fkey_1  | local_table_2 | local_table_1
+ fkey_2  | local_table_3 | local_table_1
+ fkey_3  | local_table_1 | local_table_3
+ fkey_4  | local_table_1 | local_table_4
+ fkey_5  | local_table_4 | local_table_3
+ fkey_6  | local_table_4 | local_table_4
+(6 rows)
+
+  SELECT create_citus_local_table('local_table_1', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- show that we do parallel execution
+  show citus.multi_shard_modify_mode;
+ citus.multi_shard_modify_mode
+---------------------------------------------------------------------
+ parallel
+(1 row)
+
+  SELECT conname, conrelid::regclass::text, confrelid::regclass::text
+  FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$'
+  ORDER BY 1,2,3;
+ conname |   conrelid    |   confrelid
+---------------------------------------------------------------------
+ fkey_1  | local_table_2 | local_table_1
+ fkey_2  | local_table_3 | local_table_1
+ fkey_3  | local_table_1 | local_table_3
+ fkey_4  | local_table_1 | local_table_4
+ fkey_5  | local_table_4 | local_table_3
+ fkey_6  | local_table_4 | local_table_4
+(6 rows)
+
+  SELECT COUNT(*)=4 FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT COUNT(*)=6 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+  SELECT COUNT(*)=4 FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  CREATE TABLE partitioned_table (col_1 INT REFERENCES local_table_1 (col_1)) PARTITION BY RANGE (col_1);
+  -- now that we introduced a partitioned table into our foreign key subgraph,
+  -- create_citus_local_table(cascade_via_foreign_keys) would fail for
+  -- partitioned_table as create_citus_local_table doesn't support partitioned tables
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+ERROR:  cannot create citus local table "partitioned_table", only regular tables and foreign tables are supported for citus local table creation
+ROLLBACK;
+BEGIN;
+  DROP TABLE local_table_2;
+  -- show that create_citus_local_table(cascade_via_foreign_keys) works fine after
+  -- dropping one of the relations from foreign key graph
+  SELECT create_citus_local_table('local_table_1', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  -- split local_table_2 from foreign key subgraph
+  ALTER TABLE local_table_1 DROP CONSTRAINT local_table_1_col_1_key CASCADE;
+  -- now that local_table_2 does not have any foreign keys, cascade_via_foreign_keys=true
+  -- is not needed but show that it still works fine
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- show citus tables in current schema
+  SELECT tablename FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade'
+  ORDER BY 1;
+   tablename
+---------------------------------------------------------------------
+ local_table_2
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  -- split local_table_2 from foreign key subgraph
+  ALTER TABLE local_table_1 DROP CONSTRAINT local_table_1_col_1_key CASCADE;
+  -- add a self reference on local_table_2
+  ALTER TABLE local_table_2 ADD CONSTRAINT fkey_self FOREIGN KEY(col_1) REFERENCES local_table_2(col_1);
+  -- now that local_table_2 does not have any
+  -- foreign key relationships with other tables but a self
+  -- referencing foreign key, cascade_via_foreign_keys=true
+  -- is not needed but show that it still works fine
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- show citus tables in current schema
+  SELECT tablename FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade'
+  ORDER BY 1;
+   tablename
+---------------------------------------------------------------------
+ local_table_2
+(1 row)
+
+ROLLBACK;
+CREATE TABLE distributed_table(col INT);
+SELECT create_distributed_Table('distributed_table', 'col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+BEGIN;
+  SELECT * FROM distributed_table;
+ col
+---------------------------------------------------------------------
+(0 rows)
+
+  -- succeeds as create_citus_local_table would also prefer parallel
+  -- execution like above select
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  set citus.multi_shard_modify_mode to 'sequential';
+  -- sequetial execution also works fine
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+-- test behaviour when outside of transaction block
+SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- cleanup at exit
+DROP SCHEMA create_citus_local_table_cascade CASCADE;

src/test/regress/expected/multi_extension.out

@@ -446,6 +446,7 @@ SELECT * FROM print_extension_changes();
                 previous_object                |                                current_object
 ---------------------------------------------------------------------
  function citus_total_relation_size(regclass)  |
+ function create_citus_local_table(regclass)   |
  function undistribute_table(regclass)         |
  function upgrade_to_reference_table(regclass) |
                                                | access method columnar
@@ -454,6 +455,7 @@ SELECT * FROM print_extension_changes();
                                                | function citus_internal.columnar_ensure_objects_exist()
                                                | function citus_total_relation_size(regclass,boolean)
                                                | function columnar.columnar_handler(internal)
+                                               | function create_citus_local_table(regclass,boolean)
                                                | function undistribute_table(regclass,boolean)
                                                | schema columnar
                                                | sequence columnar.storageid_seq
@@ -461,7 +463,7 @@ SELECT * FROM print_extension_changes();
                                                | table columnar.columnar_stripes
                                                | table columnar.options
                                                | view citus_tables
-(16 rows)
+(18 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/multi_extension_0.out

@@ -446,10 +446,12 @@ SELECT * FROM print_extension_changes();
                 previous_object                |                     current_object
 ---------------------------------------------------------------------
  function citus_total_relation_size(regclass)  |
+ function create_citus_local_table(regclass)   |
  function undistribute_table(regclass)         |
  function upgrade_to_reference_table(regclass) |
                                                | function citus_internal.columnar_ensure_objects_exist()
                                                | function citus_total_relation_size(regclass,boolean)
+                                               | function create_citus_local_table(regclass,boolean)
                                                | function undistribute_table(regclass,boolean)
                                                | schema columnar
                                                | sequence columnar.storageid_seq
@@ -457,7 +459,7 @@ SELECT * FROM print_extension_changes();
                                                | table columnar.columnar_stripes
                                                | table columnar.options
                                                | view citus_tables
-(12 rows)
+(14 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/single_node.out

@@ -1085,6 +1085,18 @@ NOTICE:  Renaming the new table to single_node.citus_local_table_1
 
 (1 row)
 
+CREATE TABLE local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE local_table_3 (col_1 INT UNIQUE);
+ALTER TABLE local_table_2 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_7 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_8 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
 CREATE PROCEDURE call_delegation(x int) LANGUAGE plpgsql AS $$
 BEGIN
 	 INSERT INTO test (x) VALUES ($1);

src/test/regress/expected/undistribute_table_cascade_mx.out

@@ -46,6 +46,13 @@ SELECT create_citus_local_table('citus_local_table_1');
 
 (1 row)
 
+CREATE TABLE citus_local_table_2 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_2');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
 CREATE TABLE partitioned_table_1 (col_1 INT UNIQUE, col_2 INT) PARTITION BY RANGE (col_1);
 CREATE TABLE partitioned_table_1_100_200 PARTITION OF partitioned_table_1 FOR VALUES FROM (100) TO (200);
 CREATE TABLE partitioned_table_1_200_300 PARTITION OF partitioned_table_1 FOR VALUES FROM (200) TO (300);
@@ -60,6 +67,7 @@ ALTER TABLE reference_table_1 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_2) REFERENC
 ALTER TABLE distributed_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
 ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
 ALTER TABLE partitioned_table_1 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES citus_local_table_2(col_1);
 SELECT undistribute_table('partitioned_table_1', cascade_via_foreign_keys=>true);
  undistribute_table
 ---------------------------------------------------------------------
@@ -77,5 +85,25 @@ $$);
  (localhost,57638,t,0)
 (2 rows)
 
+-- drop parititoned table as create_citus_local_table doesn't support partitioned tables
+DROP TABLE partitioned_table_1;
+SELECT create_citus_local_table('citus_local_table_1', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- both workers should print 4 as we converted all tables except
+-- partitioned table in this schema to a citus local table
+SELECT run_command_on_workers(
+$$
+SELECT count(*) FROM pg_catalog.pg_tables WHERE schemaname='undistribute_table_cascade_mx'
+$$);
+ run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,4)
+ (localhost,57638,t,4)
+(2 rows)
+
 -- cleanup at exit
 DROP SCHEMA undistribute_table_cascade_mx CASCADE;

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -76,7 +76,7 @@ ORDER BY 1;
  function coord_combine_agg(oid,cstring,anyelement)
  function coord_combine_agg_ffunc(internal,oid,cstring,anyelement)
  function coord_combine_agg_sfunc(internal,oid,cstring,anyelement)
- function create_citus_local_table(regclass)
+ function create_citus_local_table(regclass,boolean)
  function create_distributed_function(regprocedure,text,text)
  function create_distributed_table(regclass,text,citus.distribution_type,text)
  function create_intermediate_result(text,text)

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -72,7 +72,7 @@ ORDER BY 1;
  function coord_combine_agg(oid,cstring,anyelement)
  function coord_combine_agg_ffunc(internal,oid,cstring,anyelement)
  function coord_combine_agg_sfunc(internal,oid,cstring,anyelement)
- function create_citus_local_table(regclass)
+ function create_citus_local_table(regclass,boolean)
  function create_distributed_function(regprocedure,text,text)
  function create_distributed_table(regclass,text,citus.distribution_type,text)
  function create_intermediate_result(text,text)

src/test/regress/multi_schedule

@@ -323,7 +323,9 @@ test: replicate_reference_tables_to_coordinator
 test: coordinator_shouldhaveshards
 test: local_shard_utility_command_execution
 test: citus_local_tables
-test: multi_row_router_insert mixed_relkind_tests undistribute_table_cascade
+test: multi_row_router_insert mixed_relkind_tests
+test: undistribute_table_cascade
+test: create_citus_local_table_cascade
 
 test: remove_coordinator
 

src/test/regress/sql/create_citus_local_table_cascade.sql

@@ -0,0 +1,147 @@
+\set VERBOSITY terse
+
+SET citus.next_shard_id TO 1516000;
+SET citus.shard_replication_factor TO 1;
+
+CREATE SCHEMA create_citus_local_table_cascade;
+SET search_path TO create_citus_local_table_cascade;
+
+SET client_min_messages to ERROR;
+
+-- ensure that coordinator is added to pg_dist_node
+SELECT 1 FROM master_add_node('localhost', :master_port, groupId => 0);
+
+CREATE TABLE local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE local_table_3 (col_1 INT UNIQUE);
+CREATE TABLE local_table_4 (col_1 INT UNIQUE);
+
+--                                        _
+--                                       | |
+--                                       | v
+-- local_table_2 -> local_table_1 -> local_table_4
+--                      ^ |               |
+--                      | v               |
+--                  local_table_3 <--------
+ALTER TABLE local_table_2 ADD CONSTRAINT fkey_1 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES local_table_3(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES local_table_4(col_1);
+ALTER TABLE local_table_4 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES local_table_3(col_1);
+ALTER TABLE local_table_4 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES local_table_4(col_1);
+
+-- show that all of below fails as we didn't provide cascade_via_foreign_keys=true
+SELECT create_citus_local_table('local_table_1');
+SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>false);
+
+-- In each of below two transaction blocks, show that we preserve foreign keys.
+-- Also show that we converted all local_table_xxx tables in current schema
+-- to citus local tables after create_citus_local_table (cascade).
+-- So in each transaction, both selects should return true.
+
+BEGIN;
+  SELECT conname, conrelid::regclass::text, confrelid::regclass::text
+  FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$'
+  ORDER BY 1,2,3;
+
+  SELECT create_citus_local_table('local_table_1', cascade_via_foreign_keys=>true);
+
+  -- show that we do parallel execution
+  show citus.multi_shard_modify_mode;
+
+  SELECT conname, conrelid::regclass::text, confrelid::regclass::text
+  FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$'
+  ORDER BY 1,2,3;
+
+  SELECT COUNT(*)=4 FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade';
+ROLLBACK;
+
+BEGIN;
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+
+  SELECT COUNT(*)=6 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_citus_local_table_cascade') AND
+        conname ~ '^fkey\_\d+$';
+
+  SELECT COUNT(*)=4 FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade';
+ROLLBACK;
+
+BEGIN;
+  CREATE TABLE partitioned_table (col_1 INT REFERENCES local_table_1 (col_1)) PARTITION BY RANGE (col_1);
+  -- now that we introduced a partitioned table into our foreign key subgraph,
+  -- create_citus_local_table(cascade_via_foreign_keys) would fail for
+  -- partitioned_table as create_citus_local_table doesn't support partitioned tables
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+ROLLBACK;
+
+BEGIN;
+  DROP TABLE local_table_2;
+  -- show that create_citus_local_table(cascade_via_foreign_keys) works fine after
+  -- dropping one of the relations from foreign key graph
+  SELECT create_citus_local_table('local_table_1', cascade_via_foreign_keys=>true);
+ROLLBACK;
+
+BEGIN;
+  -- split local_table_2 from foreign key subgraph
+  ALTER TABLE local_table_1 DROP CONSTRAINT local_table_1_col_1_key CASCADE;
+
+  -- now that local_table_2 does not have any foreign keys, cascade_via_foreign_keys=true
+  -- is not needed but show that it still works fine
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+
+  -- show citus tables in current schema
+  SELECT tablename FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade'
+  ORDER BY 1;
+ROLLBACK;
+
+BEGIN;
+  -- split local_table_2 from foreign key subgraph
+  ALTER TABLE local_table_1 DROP CONSTRAINT local_table_1_col_1_key CASCADE;
+
+  -- add a self reference on local_table_2
+  ALTER TABLE local_table_2 ADD CONSTRAINT fkey_self FOREIGN KEY(col_1) REFERENCES local_table_2(col_1);
+
+  -- now that local_table_2 does not have any
+  -- foreign key relationships with other tables but a self
+  -- referencing foreign key, cascade_via_foreign_keys=true
+  -- is not needed but show that it still works fine
+  SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+
+  -- show citus tables in current schema
+  SELECT tablename FROM pg_dist_partition, pg_tables
+  WHERE tablename=logicalrelid::regclass::text AND
+        schemaname='create_citus_local_table_cascade'
+  ORDER BY 1;
+ROLLBACK;
+
+CREATE TABLE distributed_table(col INT);
+SELECT create_distributed_Table('distributed_table', 'col');
+
+BEGIN;
+  SELECT * FROM distributed_table;
+  -- succeeds as create_citus_local_table would also prefer parallel
+  -- execution like above select
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ROLLBACK;
+
+BEGIN;
+  set citus.multi_shard_modify_mode to 'sequential';
+  -- sequetial execution also works fine
+  SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+ROLLBACK;
+
+-- test behaviour when outside of transaction block
+SELECT create_citus_local_table('local_table_4', cascade_via_foreign_keys=>true);
+
+-- cleanup at exit
+DROP SCHEMA create_citus_local_table_cascade CASCADE;

src/test/regress/sql/single_node.sql

@@ -598,6 +598,16 @@ ALTER TABLE partitioned_table_1 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERE
 
 SELECT undistribute_table('partitioned_table_1', cascade_via_foreign_keys=>true);
 
+CREATE TABLE local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE local_table_3 (col_1 INT UNIQUE);
+
+ALTER TABLE local_table_2 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_3 ADD CONSTRAINT fkey_7 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+ALTER TABLE local_table_1 ADD CONSTRAINT fkey_8 FOREIGN KEY (col_1) REFERENCES local_table_1(col_1);
+
+SELECT create_citus_local_table('local_table_2', cascade_via_foreign_keys=>true);
+
 CREATE PROCEDURE call_delegation(x int) LANGUAGE plpgsql AS $$
 BEGIN
 	 INSERT INTO test (x) VALUES ($1);

src/test/regress/sql/undistribute_table_cascade_mx.sql

@@ -25,6 +25,9 @@ SELECT create_distributed_table('distributed_table_1', 'col_1');
 CREATE TABLE citus_local_table_1 (col_1 INT UNIQUE);
 SELECT create_citus_local_table('citus_local_table_1');
 
+CREATE TABLE citus_local_table_2 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_2');
+
 CREATE TABLE partitioned_table_1 (col_1 INT UNIQUE, col_2 INT) PARTITION BY RANGE (col_1);
 CREATE TABLE partitioned_table_1_100_200 PARTITION OF partitioned_table_1 FOR VALUES FROM (100) TO (200);
 CREATE TABLE partitioned_table_1_200_300 PARTITION OF partitioned_table_1 FOR VALUES FROM (200) TO (300);
@@ -35,6 +38,7 @@ ALTER TABLE reference_table_1 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_2) REFERENC
 ALTER TABLE distributed_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
 ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
 ALTER TABLE partitioned_table_1 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES citus_local_table_2(col_1);
 
 SELECT undistribute_table('partitioned_table_1', cascade_via_foreign_keys=>true);
 
@@ -44,5 +48,16 @@ $$
 SELECT count(*) FROM pg_catalog.pg_tables WHERE schemaname='undistribute_table_cascade_mx'
 $$);
 
+-- drop parititoned table as create_citus_local_table doesn't support partitioned tables
+DROP TABLE partitioned_table_1;
+SELECT create_citus_local_table('citus_local_table_1', cascade_via_foreign_keys=>true);
+
+-- both workers should print 4 as we converted all tables except
+-- partitioned table in this schema to a citus local table
+SELECT run_command_on_workers(
+$$
+SELECT count(*) FROM pg_catalog.pg_tables WHERE schemaname='undistribute_table_cascade_mx'
+$$);
+
 -- cleanup at exit
 DROP SCHEMA undistribute_table_cascade_mx CASCADE;