Enable reference/distributed table creation from citus local tables

src/backend/distributed/commands/alter_table.c

@@ -29,6 +29,7 @@
 #include "fmgr.h"
 
 #include "access/hash.h"
+#include "access/xact.h"
 #include "catalog/dependency.h"
 #include "catalog/pg_am.h"
 #include "columnar/cstore.h"

src/backend/distributed/commands/create_distributed_table.c

@@ -114,6 +114,9 @@ static void EnsureLocalTableEmptyIfNecessary(Oid relationId, char distributionMe
 static bool ShouldLocalTableBeEmpty(Oid relationId, char distributionMethod, bool
 									viaDeprecatedAPI);
 static void EnsureCitusTableCanBeCreated(Oid relationOid);
+static List * GetFKeyCreationCommandsRelationInvolved(Oid relationId);
+static Oid DropFKeysAndUndistributeTable(Oid relationId);
+static void DropFKeysRelationInvolved(Oid relationId);
 static bool LocalTableEmpty(Oid tableId);
 static void CopyLocalDataIntoShards(Oid relationId);
 static List * TupleDescColumnNameList(TupleDesc tupleDescriptor);
@@ -209,13 +212,14 @@ create_distributed_table(PG_FUNCTION_ARGS)
 	 * backends manipulating this relation.
 	 */
 	Relation relation = try_relation_open(relationId, ExclusiveLock);
-
 	if (relation == NULL)
 	{
 		ereport(ERROR, (errmsg("could not create distributed table: "
 							   "relation does not exist")));
 	}
 
+	relation_close(relation, NoLock);
+
 	char *distributionColumnName = text_to_cstring(distributionColumnText);
 	Var *distributionColumn = BuildDistributionKeyFromColumnName(relation,
 																 distributionColumnName);
@@ -227,8 +231,6 @@ create_distributed_table(PG_FUNCTION_ARGS)
 	CreateDistributedTable(relationId, distributionColumn, distributionMethod,
 						   ShardCount, colocateWithTableName, viaDeprecatedAPI);
 
-	relation_close(relation, NoLock);
-
 	PG_RETURN_VOID();
 }
 
@@ -260,7 +262,14 @@ create_reference_table(PG_FUNCTION_ARGS)
 	 * sense of this table until we've committed, and we don't want multiple
 	 * backends manipulating this relation.
 	 */
-	Relation relation = relation_open(relationId, ExclusiveLock);
+	Relation relation = try_relation_open(relationId, ExclusiveLock);
+	if (relation == NULL)
+	{
+		ereport(ERROR, (errmsg("could not create reference table: "
+							   "relation does not exist")));
+	}
+
+	relation_close(relation, NoLock);
 
 	List *workerNodeList = ActivePrimaryNodeList(ShareLock);
 	int workerCount = list_length(workerNodeList);
@@ -277,9 +286,6 @@ create_reference_table(PG_FUNCTION_ARGS)
 
 	CreateDistributedTable(relationId, distributionColumn, DISTRIBUTE_BY_NONE,
 						   ShardCount, colocateWithTableName, viaDeprecatedAPI);
-
-	relation_close(relation, NoLock);
-
 	PG_RETURN_VOID();
 }
 
@@ -338,6 +344,25 @@ void
 CreateDistributedTable(Oid relationId, Var *distributionColumn, char distributionMethod,
 					   int shardCount, char *colocateWithTableName, bool viaDeprecatedAPI)
 {
+	/*
+	 * EnsureTableNotDistributed errors out when relation is a citus table but
+	 * we don't want to ask user to first undistribute their citus local tables
+	 * when creating reference or distributed tables from them.
+	 * For this reason, here we undistribute citus local tables beforehand.
+	 * But since UndistributeTable does not support undistributing relations
+	 * involved in foreign key relationships, we first drop foreign keys that
+	 * given relation is involved, then we undistribute the relation and finally
+	 * we re-create dropped foreign keys at the end of this function.
+	 */
+	List *fKeyCreationCommandsRelationInvolved = NIL;
+	if (IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+	{
+		/* store foreign key creation commands that relation is involved */
+		fKeyCreationCommandsRelationInvolved =
+			GetFKeyCreationCommandsRelationInvolved(relationId);
+		relationId = DropFKeysAndUndistributeTable(relationId);
+	}
+
 	/*
 	 * distributed tables might have dependencies on different objects, since we create
 	 * shards for a distributed table via multiple sessions these objects will be created
@@ -440,6 +465,85 @@ CreateDistributedTable(Oid relationId, Var *distributionColumn, char distributio
 			CopyLocalDataIntoShards(relationId);
 		}
 	}
+
+	/* now recreate foreign keys that we dropped beforehand */
+	ExecuteAndLogDDLCommandList(fKeyCreationCommandsRelationInvolved);
+}
+
+
+/*
+ * GetFKeyCreationCommandsRelationInvolved returns a list of DDL commands to
+ * recreate the foreign keys that relation with relationId is involved.
+ */
+static List *
+GetFKeyCreationCommandsRelationInvolved(Oid relationId)
+{
+	int referencingFKeysFlag = INCLUDE_REFERENCING_CONSTRAINTS |
+							   INCLUDE_ALL_TABLE_TYPES;
+	List *referencingFKeyCreationCommands =
+		GetForeignConstraintCommandsInternal(relationId, referencingFKeysFlag);
+
+	/* already captured self referencing foreign keys, so use EXCLUDE_SELF_REFERENCES */
+	int referencedFKeysFlag = INCLUDE_REFERENCED_CONSTRAINTS |
+							  EXCLUDE_SELF_REFERENCES |
+							  INCLUDE_ALL_TABLE_TYPES;
+	List *referencedFKeyCreationCommands =
+		GetForeignConstraintCommandsInternal(relationId, referencedFKeysFlag);
+	return list_concat(referencingFKeyCreationCommands, referencedFKeyCreationCommands);
+}
+
+
+/*
+ * DropFKeysAndUndistributeTable drops all foreign keys that relation with
+ * relationId is involved then undistributes it.
+ * Note that as UndistributeTable changes relationId of relation, this
+ * function also returns new relationId of relation.
+ * Also note that callers are responsible for storing & recreating foreign
+ * keys to be dropped if needed.
+ */
+static Oid
+DropFKeysAndUndistributeTable(Oid relationId)
+{
+	DropFKeysRelationInvolved(relationId);
+
+	/* store them before calling UndistributeTable as it changes relationId */
+	char *relationName = get_rel_name(relationId);
+	Oid schemaId = get_rel_namespace(relationId);
+
+	TableConversionParameters params = {
+		.relationId = relationId,
+		.cascadeViaForeignKeys = false
+	};
+	UndistributeTable(&params);
+
+	Oid newRelationId = get_relname_relid(relationName, schemaId);
+
+	/*
+	 * We don't expect this to happen but to be on the safe side let's error
+	 * out here.
+	 */
+	EnsureRelationExists(newRelationId);
+
+	return newRelationId;
+}
+
+
+/*
+ * DropFKeysRelationInvolved drops all foreign keys that relation with
+ * relationId is involved.
+ */
+static void
+DropFKeysRelationInvolved(Oid relationId)
+{
+	int referencingFKeysFlag = INCLUDE_REFERENCING_CONSTRAINTS |
+							   INCLUDE_ALL_TABLE_TYPES;
+	DropRelationForeignKeys(relationId, referencingFKeysFlag);
+
+	/* already captured self referencing foreign keys, so use EXCLUDE_SELF_REFERENCES */
+	int referencedFKeysFlag = INCLUDE_REFERENCED_CONSTRAINTS |
+							  EXCLUDE_SELF_REFERENCES |
+							  INCLUDE_ALL_TABLE_TYPES;
+	DropRelationForeignKeys(relationId, referencedFKeysFlag);
 }
 
 

src/test/regress/expected/create_ref_dist_from_citus_local.out

@@ -0,0 +1,369 @@
+\set VERBOSITY terse
+SET citus.next_shard_id TO 1800000;
+SET citus.shard_replication_factor TO 1;
+CREATE SCHEMA create_ref_dist_from_citus_local;
+SET search_path TO create_ref_dist_from_citus_local;
+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 citus_local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_3 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_4 (col_1 INT UNIQUE);
+ALTER TABLE citus_local_table_2 ADD CONSTRAINT fkey_1 FOREIGN KEY (col_1) REFERENCES citus_local_table_1(col_1);
+ALTER TABLE citus_local_table_3 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_1) REFERENCES citus_local_table_1(col_1);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES citus_local_table_3(col_1);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES citus_local_table_4(col_1);
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES citus_local_table_3(col_1);
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES citus_local_table_4(col_1);
+SELECT create_citus_local_table('citus_local_table_1', cascade_via_foreign_keys=>true);
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE reference_table_1(col_1 INT UNIQUE, col_2 INT UNIQUE);
+CREATE TABLE reference_table_2(col_1 INT UNIQUE, col_2 INT UNIQUE);
+SELECT create_reference_table('reference_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_reference_table('reference_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_7 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
+ALTER TABLE reference_table_2 ADD CONSTRAINT fkey_8 FOREIGN KEY (col_1) REFERENCES citus_local_table_2(col_1);
+CREATE TABLE distributed_table_1(col_1 INT UNIQUE, col_2 INT);
+CREATE TABLE partitioned_dist_table_1 (col_1 INT UNIQUE, col_2 INT) PARTITION BY RANGE (col_1);
+SELECT create_distributed_table('distributed_table_1', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('partitioned_dist_table_1', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE partitioned_dist_table_1 ADD CONSTRAINT fkey_9 FOREIGN KEY (col_1) REFERENCES distributed_table_1(col_1);
+ALTER TABLE distributed_table_1 ADD CONSTRAINT fkey_10 FOREIGN KEY (col_1) REFERENCES reference_table_2(col_2);
+ALTER TABLE partitioned_dist_table_1 ADD CONSTRAINT fkey_11 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
+-- As we will heavily rely on this feature after implementing automatic
+-- convertion of postgres tables to citus local tables, let's have a
+-- complex foreign key graph to see everything is fine.
+--
+--  distributed_table_1  <----------------    partitioned_dist_table_1
+--           |                                           |
+--           v                                           v
+--  reference_table_2                 _          reference_table_1
+--           |                       | |                 ^
+--           v                       | v                 |
+-- citus_local_table_2 -> citus_local_table_1 -> citus_local_table_4
+--                                   ^ |                 |
+--                                   | v                 |
+--                           citus_local_table_3 <--------
+-- Now print metadata after each of create_reference/distributed_table
+-- operations to show that everything is fine. Also show that we
+-- preserve foreign keys.
+BEGIN;
+  SELECT create_reference_table('citus_local_table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+        tablename         | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table_1      | n          | t
+ citus_local_table_2      | n          | c
+ citus_local_table_3      | n          | c
+ citus_local_table_4      | n          | c
+ distributed_table_1      | h          | c
+ partitioned_dist_table_1 | h          | c
+ reference_table_1        | n          | t
+ reference_table_2        | n          | t
+(8 rows)
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  SELECT create_reference_table('citus_local_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+        tablename         | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table_1      | n          | c
+ citus_local_table_2      | n          | t
+ citus_local_table_3      | n          | c
+ citus_local_table_4      | n          | c
+ distributed_table_1      | h          | c
+ partitioned_dist_table_1 | h          | c
+ reference_table_1        | n          | t
+ reference_table_2        | n          | t
+(8 rows)
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+-- those two errors out as they reference to citus local tables but
+-- distributed tables cannot reference to postgres or citus local tables
+SELECT create_distributed_table('citus_local_table_1', 'col_1');
+ERROR:  cannot create foreign key constraint since relations are not colocated or not referencing a reference table
+SELECT create_distributed_table('citus_local_table_4', 'col_1');
+ERROR:  cannot create foreign key constraint since relations are not colocated or not referencing a reference table
+BEGIN;
+  SELECT create_reference_table('citus_local_table_2');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- this would error out
+  SELECT create_reference_table('citus_local_table_2');
+ERROR:  table "citus_local_table_2" is already distributed
+ROLLBACK;
+-- test with a standalone table
+CREATE TABLE citus_local_table_5 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_5');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- this would error out
+  SELECT create_reference_table('citus_local_table_5');
+ERROR:  table "citus_local_table_5" is already distributed
+ROLLBACK;
+BEGIN;
+  SELECT create_reference_table('citus_local_table_5');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  SELECT create_reference_table('citus_local_table_5');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+        tablename         | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table_1      | n          | c
+ citus_local_table_2      | n          | c
+ citus_local_table_3      | n          | c
+ citus_local_table_4      | n          | c
+ citus_local_table_5      | h          | c
+ distributed_table_1      | h          | c
+ partitioned_dist_table_1 | h          | c
+ reference_table_1        | n          | t
+ reference_table_2        | n          | t
+(9 rows)
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+        tablename         | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table_1      | n          | c
+ citus_local_table_2      | n          | c
+ citus_local_table_3      | n          | c
+ citus_local_table_4      | n          | c
+ citus_local_table_5      | h          | c
+ distributed_table_1      | h          | c
+ partitioned_dist_table_1 | h          | c
+ reference_table_1        | n          | t
+ reference_table_2        | n          | t
+(9 rows)
+
+  SELECT COUNT(*)=12 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  -- define a self reference and a foreign key to reference table
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_13 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+        tablename         | partmethod | repmodel
+---------------------------------------------------------------------
+ citus_local_table_1      | n          | c
+ citus_local_table_2      | n          | c
+ citus_local_table_3      | n          | c
+ citus_local_table_4      | n          | c
+ citus_local_table_5      | h          | c
+ distributed_table_1      | h          | c
+ partitioned_dist_table_1 | h          | c
+ reference_table_1        | n          | t
+ reference_table_2        | n          | t
+(9 rows)
+
+  SELECT COUNT(*)=13 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ ?column?
+---------------------------------------------------------------------
+ t
+(1 row)
+
+ROLLBACK;
+CREATE TABLE citus_local_table_6 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_6');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  -- errors out as foreign keys from distributed tables to citus
+  -- local tables are not supported
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ERROR:  cannot create foreign key constraint since relations are not colocated or not referencing a reference table
+ROLLBACK;
+BEGIN;
+  -- errors out as foreign keys from citus local tables to distributed
+  -- tables are not supported
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_distributed_table('citus_local_table_6', 'col_1');
+ERROR:  cannot create foreign key constraint since foreign keys from reference tables and citus local tables to distributed tables are not supported
+ROLLBACK;
+-- have some more tests with foreign keys between citus local
+-- and reference tables
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_reference_table('citus_local_table_5');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_reference_table('citus_local_table_6');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  CREATE FUNCTION update_value() RETURNS trigger AS $update_value$
+  BEGIN
+      NEW.value := value+1 ;
+      RETURN NEW;
+  END;
+  $update_value$ LANGUAGE plpgsql;
+  CREATE TRIGGER update_value_dist
+  AFTER INSERT ON citus_local_table_6
+  FOR EACH ROW EXECUTE FUNCTION update_value();
+  -- show that we error out as we don't supprt triggers on distributed tables
+  SELECT create_distributed_table('citus_local_table_6', 'col_1');
+ERROR:  cannot distribute relation "citus_local_table_6" because it has triggers
+ROLLBACK;
+-- make sure that creating append / range distributed tables is also ok
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1', 'range');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+BEGIN;
+  ALTER TABLE citus_local_table_5 DROP CONSTRAINT citus_local_table_5_col_1_key;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1', 'append');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ROLLBACK;
+-- cleanup at exit
+DROP SCHEMA create_ref_dist_from_citus_local CASCADE;

src/test/regress/multi_schedule

@@ -324,7 +324,7 @@ 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
+test: multi_row_router_insert mixed_relkind_tests create_ref_dist_from_citus_local
 test: undistribute_table_cascade
 test: create_citus_local_table_cascade
 

src/test/regress/sql/create_ref_dist_from_citus_local.sql

@@ -0,0 +1,217 @@
+\set VERBOSITY terse
+
+SET citus.next_shard_id TO 1800000;
+SET citus.shard_replication_factor TO 1;
+
+CREATE SCHEMA create_ref_dist_from_citus_local;
+SET search_path TO create_ref_dist_from_citus_local;
+
+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 citus_local_table_1 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_2 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_3 (col_1 INT UNIQUE);
+CREATE TABLE citus_local_table_4 (col_1 INT UNIQUE);
+ALTER TABLE citus_local_table_2 ADD CONSTRAINT fkey_1 FOREIGN KEY (col_1) REFERENCES citus_local_table_1(col_1);
+ALTER TABLE citus_local_table_3 ADD CONSTRAINT fkey_2 FOREIGN KEY (col_1) REFERENCES citus_local_table_1(col_1);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_3 FOREIGN KEY (col_1) REFERENCES citus_local_table_3(col_1);
+ALTER TABLE citus_local_table_1 ADD CONSTRAINT fkey_4 FOREIGN KEY (col_1) REFERENCES citus_local_table_4(col_1);
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_5 FOREIGN KEY (col_1) REFERENCES citus_local_table_3(col_1);
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_6 FOREIGN KEY (col_1) REFERENCES citus_local_table_4(col_1);
+
+SELECT create_citus_local_table('citus_local_table_1', cascade_via_foreign_keys=>true);
+
+CREATE TABLE reference_table_1(col_1 INT UNIQUE, col_2 INT UNIQUE);
+CREATE TABLE reference_table_2(col_1 INT UNIQUE, col_2 INT UNIQUE);
+
+SELECT create_reference_table('reference_table_1');
+SELECT create_reference_table('reference_table_2');
+
+ALTER TABLE citus_local_table_4 ADD CONSTRAINT fkey_7 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
+ALTER TABLE reference_table_2 ADD CONSTRAINT fkey_8 FOREIGN KEY (col_1) REFERENCES citus_local_table_2(col_1);
+
+CREATE TABLE distributed_table_1(col_1 INT UNIQUE, col_2 INT);
+CREATE TABLE partitioned_dist_table_1 (col_1 INT UNIQUE, col_2 INT) PARTITION BY RANGE (col_1);
+
+SELECT create_distributed_table('distributed_table_1', 'col_1');
+SELECT create_distributed_table('partitioned_dist_table_1', 'col_1');
+
+ALTER TABLE partitioned_dist_table_1 ADD CONSTRAINT fkey_9 FOREIGN KEY (col_1) REFERENCES distributed_table_1(col_1);
+ALTER TABLE distributed_table_1 ADD CONSTRAINT fkey_10 FOREIGN KEY (col_1) REFERENCES reference_table_2(col_2);
+ALTER TABLE partitioned_dist_table_1 ADD CONSTRAINT fkey_11 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_2);
+
+-- As we will heavily rely on this feature after implementing automatic
+-- convertion of postgres tables to citus local tables, let's have a
+-- complex foreign key graph to see everything is fine.
+--
+--  distributed_table_1  <----------------    partitioned_dist_table_1
+--           |                                           |
+--           v                                           v
+--  reference_table_2                 _          reference_table_1
+--           |                       | |                 ^
+--           v                       | v                 |
+-- citus_local_table_2 -> citus_local_table_1 -> citus_local_table_4
+--                                   ^ |                 |
+--                                   | v                 |
+--                           citus_local_table_3 <--------
+
+-- Now print metadata after each of create_reference/distributed_table
+-- operations to show that everything is fine. Also show that we
+-- preserve foreign keys.
+
+BEGIN;
+  SELECT create_reference_table('citus_local_table_1');
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ROLLBACK;
+
+BEGIN;
+  SELECT create_reference_table('citus_local_table_2');
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ROLLBACK;
+
+-- those two errors out as they reference to citus local tables but
+-- distributed tables cannot reference to postgres or citus local tables
+SELECT create_distributed_table('citus_local_table_1', 'col_1');
+SELECT create_distributed_table('citus_local_table_4', 'col_1');
+
+BEGIN;
+  SELECT create_reference_table('citus_local_table_2');
+  -- this would error out
+  SELECT create_reference_table('citus_local_table_2');
+ROLLBACK;
+
+-- test with a standalone table
+CREATE TABLE citus_local_table_5 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_5');
+
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+  -- this would error out
+  SELECT create_reference_table('citus_local_table_5');
+ROLLBACK;
+
+BEGIN;
+  SELECT create_reference_table('citus_local_table_5');
+ROLLBACK;
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  SELECT create_reference_table('citus_local_table_5');
+ROLLBACK;
+
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+
+  SELECT COUNT(*)=11 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ROLLBACK;
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+
+  SELECT COUNT(*)=12 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ROLLBACK;
+
+BEGIN;
+  -- define a self reference and a foreign key to reference table
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_5(col_1);
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_13 FOREIGN KEY (col_1) REFERENCES reference_table_1(col_1);
+
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+
+  SELECT logicalrelid::text AS tablename, partmethod, repmodel FROM pg_dist_partition
+  WHERE logicalrelid::text IN (SELECT tablename FROM pg_tables WHERE schemaname='create_ref_dist_from_citus_local')
+  ORDER BY tablename;
+
+  SELECT COUNT(*)=13 FROM pg_constraint
+  WHERE connamespace = (SELECT oid FROM pg_namespace WHERE nspname='create_ref_dist_from_citus_local') AND
+        conname ~ '^fkey\_\d+$';
+ROLLBACK;
+
+CREATE TABLE citus_local_table_6 (col_1 INT UNIQUE);
+SELECT create_citus_local_table('citus_local_table_6');
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  -- errors out as foreign keys from distributed tables to citus
+  -- local tables are not supported
+  SELECT create_distributed_table('citus_local_table_5', 'col_1');
+ROLLBACK;
+
+BEGIN;
+  -- errors out as foreign keys from citus local tables to distributed
+  -- tables are not supported
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_distributed_table('citus_local_table_6', 'col_1');
+ROLLBACK;
+
+-- have some more tests with foreign keys between citus local
+-- and reference tables
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_reference_table('citus_local_table_5');
+ROLLBACK;
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 ADD CONSTRAINT fkey_12 FOREIGN KEY (col_1) REFERENCES citus_local_table_6(col_1);
+  SELECT create_reference_table('citus_local_table_6');
+ROLLBACK;
+
+BEGIN;
+  CREATE FUNCTION update_value() RETURNS trigger AS $update_value$
+  BEGIN
+      NEW.value := value+1 ;
+      RETURN NEW;
+  END;
+  $update_value$ LANGUAGE plpgsql;
+
+  CREATE TRIGGER update_value_dist
+  AFTER INSERT ON citus_local_table_6
+  FOR EACH ROW EXECUTE PROCEDURE update_value();
+
+  -- show that we error out as we don't supprt triggers on distributed tables
+  SELECT create_distributed_table('citus_local_table_6', 'col_1');
+ROLLBACK;
+
+-- make sure that creating append / range distributed tables is also ok
+BEGIN;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1', 'range');
+ROLLBACK;
+
+BEGIN;
+  ALTER TABLE citus_local_table_5 DROP CONSTRAINT citus_local_table_5_col_1_key;
+  SELECT create_distributed_table('citus_local_table_5', 'col_1', 'append');
+ROLLBACK;
+
+-- cleanup at exit
+DROP SCHEMA create_ref_dist_from_citus_local CASCADE;