FK from dist to ref is tested for partitioning, MX

src/test/regress/expected/foreign_key_to_reference_table.out

@@ -1581,9 +1581,9 @@ INSERT INTO test_table_2 VALUES (1,1), (2,2), (3,3);
 -- should succeed
 ALTER TABLE test_table_2 ALTER COLUMN value_1 SET DATA TYPE bigint;
 ALTER TABLE test_table_1 ALTER COLUMN id SET DATA TYPE bigint;
--- should fail since there is a bigint out of integer range > (2^32 - 1)
 INSERT INTO test_table_1 VALUES (2147483648,4);
 INSERT INTO test_table_2 VALUES (4,2147483648);
+-- should fail since there is a bigint out of integer range > (2^32 - 1)
 ALTER TABLE test_table_2 ALTER COLUMN value_1 SET DATA TYPE int;
 ERROR:  integer out of range
 CONTEXT:  while executing command on localhost:57637
@@ -1802,6 +1802,59 @@ ERROR:  current transaction is aborted, commands ignored until end of transactio
 ERROR:  current transaction is aborted, commands ignored until end of transaction block
 COMMIT;
 DROP TABLE test_table_1, test_table_2, test_table_3;
+-- NOTE: Postgres does not support foreign keys on partitioned tables currently.
+-- However, we can create foreign keys to/from the partitions themselves.
+-- The following tests chech if we create the foreign constraints in partitions properly.
+CREATE TABLE referenced_table(id int PRIMARY KEY, test_column int);
+CREATE TABLE referencing_table(id int, value_1 int) PARTITION BY RANGE (value_1);
+CREATE TABLE referencing_table_0 PARTITION OF referencing_table FOR VALUES FROM (0) TO (2);
+CREATE TABLE referencing_table_2 PARTITION OF referencing_table FOR VALUES FROM (2) TO (4);
+CREATE TABLE referencing_table_4 PARTITION OF referencing_table FOR VALUES FROM (4) TO (6);
+-- partitioned tables are not supported as reference tables
+select create_reference_table('referencing_table');
+ERROR:  distributing partitioned tables in only supported for hash-distributed tables
+-- partitioned tables are supported as hash distributed table
+SELECT create_reference_table('referenced_table');
+ create_reference_table 
+------------------------
+ 
+(1 row)
+
+SELECT create_distributed_table('referencing_table', 'id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+-- add foreign constraints in between partitions
+ALTER TABLE referencing_table_0 ADD CONSTRAINT pkey PRIMARY KEY (id);
+ALTER TABLE referencing_table_4 ADD CONSTRAINT fkey FOREIGN KEY (id) REFERENCES referencing_table_0;
+-- add foreign constraint from a partition to reference table
+ALTER TABLE referencing_table_4 ADD CONSTRAINT fkey_to_ref FOREIGN KEY (value_1) REFERENCES referenced_table;
+-- should fail since the data will flow to partitioning_test_4 and it has a foreign constraint to partitioning_test_0 on id column
+INSERT INTO referencing_table VALUES (0, 5);
+ERROR:  insert or update on table "referencing_table_4_7000533" violates foreign key constraint "fkey_7000533"
+DETAIL:  Key (id)=(0) is not present in table "referencing_table_0_7000517".
+CONTEXT:  while executing command on localhost:57638
+-- should succeed on partitioning_test_0
+INSERT INTO referencing_table VALUES (0, 1);
+SELECT * FROM referencing_table;
+ id | value_1 
+----+---------
+  0 |       1
+(1 row)
+
+-- should fail since partitioning_test_4 has foreign constraint to referenced_table on value_1 column
+INSERT INTO referencing_table VALUES (0, 5);
+ERROR:  insert or update on table "referencing_table_4_7000533" violates foreign key constraint "fkey_to_ref_7000533"
+DETAIL:  Key (value_1)=(5) is not present in table "referenced_table_7000505".
+CONTEXT:  while executing command on localhost:57638
+INSERT INTO referenced_table VALUES(5,5);
+-- should succeed since both of the foreign constraints are positive
+INSERT INTO referencing_table VALUES (0, 5);
+DROP TABLE referenced_table CASCADE;
+NOTICE:  drop cascades to constraint fkey_to_ref on table referencing_table_4
+DROP TABLE referencing_table;
 DROP SCHEMA fkey_reference_table CASCADE;
 NOTICE:  drop cascades to 3 other objects
 DETAIL:  drop cascades to type foreign_details

src/test/regress/expected/mx_foreign_key_to_reference_table.out

@@ -0,0 +1,90 @@
+CREATE SCHEMA fkey_reference_table;
+SET search_path TO 'fkey_reference_table';
+SET citus.shard_replication_factor TO 1;
+SET citus.shard_count TO 8;
+SET citus.next_shard_id TO 7000000;
+SET citus.next_placement_id TO 7000000;
+SET citus.replication_model TO streaming;
+-- Setup the view so that we can check if the foreign keys are created properly
+CREATE TYPE foreign_details AS (name text, relid text, refd_relid text);
+SELECT run_command_on_workers($$CREATE TYPE foreign_details AS (name text, relid text, refd_relid text)$$);
+      run_command_on_workers       
+-----------------------------------
+ (localhost,57637,t,"CREATE TYPE")
+ (localhost,57638,t,"CREATE TYPE")
+(2 rows)
+
+CREATE VIEW table_fkeys_in_workers AS 
+SELECT
+(json_populate_record(NULL::foreign_details, 
+  json_array_elements_text((run_command_on_workers( $$ 
+    SELECT
+      COALESCE(json_agg(row_to_json(d)), '[]'::json) 
+    FROM
+      (
+        SELECT
+          distinct name,
+          relid::regclass::text,
+          refd_relid::regclass::text
+        FROM
+          table_fkey_cols 
+      )
+      d $$ )).RESULT::json )::json )).* ;
+-- Check if MX can create foreign keys properly on foreign keys from distributed to reference tables
+CREATE TABLE referenced_table(test_column int, test_column2 int, PRIMARY KEY(test_column));
+CREATE TABLE referenced_table2(test_column int, test_column2 int, PRIMARY KEY(test_column2));
+CREATE TABLE referencing_table(id int, ref_id int);
+ALTER TABLE referencing_table ADD CONSTRAINT fkey_ref FOREIGN KEY (id) REFERENCES referenced_table(test_column) ON DELETE CASCADE;
+ALTER TABLE referencing_table ADD CONSTRAINT foreign_key_2 FOREIGN KEY (id) REFERENCES referenced_table2(test_column2) ON DELETE CASCADE;
+SELECT create_reference_table('referenced_table');
+ create_reference_table 
+------------------------
+ 
+(1 row)
+
+SELECT create_reference_table('referenced_table2');
+ create_reference_table 
+------------------------
+ 
+(1 row)
+
+SELECT create_distributed_table('referencing_table', 'id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+SET search_path TO 'fkey_reference_table';
+SELECT * FROM table_fkeys_in_workers WHERE relid LIKE 'fkey_reference_table.%' AND refd_relid LIKE 'fkey_reference_table.%' ORDER BY 1, 2;
+         name          |                     relid                      |                   refd_relid                   
+-----------------------+------------------------------------------------+------------------------------------------------
+ fkey_ref              | fkey_reference_table.referencing_table         | fkey_reference_table.referenced_table
+ fkey_ref              | fkey_reference_table.referencing_table         | fkey_reference_table.referenced_table
+ fkey_ref_7000002      | fkey_reference_table.referencing_table_7000002 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000003      | fkey_reference_table.referencing_table_7000003 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000004      | fkey_reference_table.referencing_table_7000004 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000005      | fkey_reference_table.referencing_table_7000005 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000006      | fkey_reference_table.referencing_table_7000006 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000007      | fkey_reference_table.referencing_table_7000007 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000008      | fkey_reference_table.referencing_table_7000008 | fkey_reference_table.referenced_table_7000000
+ fkey_ref_7000009      | fkey_reference_table.referencing_table_7000009 | fkey_reference_table.referenced_table_7000000
+ foreign_key_2         | fkey_reference_table.referencing_table         | fkey_reference_table.referenced_table2
+ foreign_key_2         | fkey_reference_table.referencing_table         | fkey_reference_table.referenced_table2
+ foreign_key_2_7000002 | fkey_reference_table.referencing_table_7000002 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000003 | fkey_reference_table.referencing_table_7000003 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000004 | fkey_reference_table.referencing_table_7000004 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000005 | fkey_reference_table.referencing_table_7000005 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000006 | fkey_reference_table.referencing_table_7000006 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000007 | fkey_reference_table.referencing_table_7000007 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000008 | fkey_reference_table.referencing_table_7000008 | fkey_reference_table.referenced_table2_7000001
+ foreign_key_2_7000009 | fkey_reference_table.referencing_table_7000009 | fkey_reference_table.referenced_table2_7000001
+(20 rows)
+
+DROP SCHEMA fkey_reference_table CASCADE;
+NOTICE:  drop cascades to 5 other objects
+DETAIL:  drop cascades to type foreign_details
+drop cascades to view table_fkeys_in_workers
+drop cascades to table referenced_table
+drop cascades to table referenced_table2
+drop cascades to table referencing_table
+SET search_path TO DEFAULT;

src/test/regress/multi_mx_schedule

@@ -25,7 +25,7 @@ test: multi_mx_tpch_query12 multi_mx_tpch_query14 multi_mx_tpch_query19
 test: multi_mx_tpch_query3 multi_mx_tpch_query6 multi_mx_tpch_query7
 test: multi_mx_tpch_query7_nested multi_mx_ddl
 test: recursive_dml_queries_mx
-test: multi_mx_repartition_udt_prepare
+test: multi_mx_repartition_udt_prepare mx_foreign_key_to_reference_table
 test: multi_mx_repartition_join_w1 multi_mx_repartition_join_w2 multi_mx_repartition_udt_w1 multi_mx_repartition_udt_w2
 test: multi_mx_metadata 
 test: multi_mx_modifications

src/test/regress/sql/foreign_key_to_reference_table.sql

@@ -795,11 +795,10 @@ INSERT INTO test_table_2 VALUES (1,1), (2,2), (3,3);
 ALTER TABLE test_table_2 ALTER COLUMN value_1 SET DATA TYPE bigint;
 ALTER TABLE test_table_1 ALTER COLUMN id SET DATA TYPE bigint;
 
--- should fail since there is a bigint out of integer range > (2^32 - 1)
 INSERT INTO test_table_1 VALUES (2147483648,4);
 INSERT INTO test_table_2 VALUES (4,2147483648);
+-- should fail since there is a bigint out of integer range > (2^32 - 1)
 ALTER TABLE test_table_2 ALTER COLUMN value_1 SET DATA TYPE int;
-
 SELECT count(*) FROM table_fkeys_in_workers WHERE relid LIKE 'fkey_reference_table.%' AND refd_relid LIKE 'fkey_reference_table.%';
 DROP TABLE test_table_1 CASCADE;
 DROP TABLE test_table_2;
@@ -909,6 +908,40 @@ BEGIN;
   ALTER TABLE test_table_1 ADD COLUMN id int;
 COMMIT;
 DROP TABLE test_table_1, test_table_2, test_table_3;
+-- NOTE: Postgres does not support foreign keys on partitioned tables currently.
+-- However, we can create foreign keys to/from the partitions themselves.
+-- The following tests chech if we create the foreign constraints in partitions properly.
+CREATE TABLE referenced_table(id int PRIMARY KEY, test_column int);
+CREATE TABLE referencing_table(id int, value_1 int) PARTITION BY RANGE (value_1);
+CREATE TABLE referencing_table_0 PARTITION OF referencing_table FOR VALUES FROM (0) TO (2);
+CREATE TABLE referencing_table_2 PARTITION OF referencing_table FOR VALUES FROM (2) TO (4);
+CREATE TABLE referencing_table_4 PARTITION OF referencing_table FOR VALUES FROM (4) TO (6);
+
+-- partitioned tables are not supported as reference tables
+select create_reference_table('referencing_table');
+
+-- partitioned tables are supported as hash distributed table
+SELECT create_reference_table('referenced_table');
+SELECT create_distributed_table('referencing_table', 'id');
+
+-- add foreign constraints in between partitions
+ALTER TABLE referencing_table_0 ADD CONSTRAINT pkey PRIMARY KEY (id);
+ALTER TABLE referencing_table_4 ADD CONSTRAINT fkey FOREIGN KEY (id) REFERENCES referencing_table_0;
+-- add foreign constraint from a partition to reference table
+ALTER TABLE referencing_table_4 ADD CONSTRAINT fkey_to_ref FOREIGN KEY (value_1) REFERENCES referenced_table;
+-- should fail since the data will flow to partitioning_test_4 and it has a foreign constraint to partitioning_test_0 on id column
+INSERT INTO referencing_table VALUES (0, 5);
+-- should succeed on partitioning_test_0
+INSERT INTO referencing_table VALUES (0, 1);
+SELECT * FROM referencing_table;
+-- should fail since partitioning_test_4 has foreign constraint to referenced_table on value_1 column
+INSERT INTO referencing_table VALUES (0, 5);
+INSERT INTO referenced_table VALUES(5,5);
+-- should succeed since both of the foreign constraints are positive
+INSERT INTO referencing_table VALUES (0, 5);
+
+DROP TABLE referenced_table CASCADE;
+DROP TABLE referencing_table;
 
 DROP SCHEMA fkey_reference_table CASCADE;
 SET search_path TO DEFAULT;

src/test/regress/sql/mx_foreign_key_to_reference_table.sql

@@ -0,0 +1,45 @@
+CREATE SCHEMA fkey_reference_table;
+SET search_path TO 'fkey_reference_table';
+SET citus.shard_replication_factor TO 1;
+SET citus.shard_count TO 8;
+SET citus.next_shard_id TO 7000000;
+SET citus.next_placement_id TO 7000000;
+SET citus.replication_model TO streaming;
+
+-- Setup the view so that we can check if the foreign keys are created properly
+CREATE TYPE foreign_details AS (name text, relid text, refd_relid text);
+SELECT run_command_on_workers($$CREATE TYPE foreign_details AS (name text, relid text, refd_relid text)$$);
+
+CREATE VIEW table_fkeys_in_workers AS 
+SELECT
+(json_populate_record(NULL::foreign_details, 
+  json_array_elements_text((run_command_on_workers( $$ 
+    SELECT
+      COALESCE(json_agg(row_to_json(d)), '[]'::json) 
+    FROM
+      (
+        SELECT
+          distinct name,
+          relid::regclass::text,
+          refd_relid::regclass::text
+        FROM
+          table_fkey_cols 
+      )
+      d $$ )).RESULT::json )::json )).* ;
+
+-- Check if MX can create foreign keys properly on foreign keys from distributed to reference tables
+CREATE TABLE referenced_table(test_column int, test_column2 int, PRIMARY KEY(test_column));
+CREATE TABLE referenced_table2(test_column int, test_column2 int, PRIMARY KEY(test_column2));
+CREATE TABLE referencing_table(id int, ref_id int);
+ALTER TABLE referencing_table ADD CONSTRAINT fkey_ref FOREIGN KEY (id) REFERENCES referenced_table(test_column) ON DELETE CASCADE;
+ALTER TABLE referencing_table ADD CONSTRAINT foreign_key_2 FOREIGN KEY (id) REFERENCES referenced_table2(test_column2) ON DELETE CASCADE;
+
+SELECT create_reference_table('referenced_table');
+SELECT create_reference_table('referenced_table2');
+SELECT create_distributed_table('referencing_table', 'id');
+
+SET search_path TO 'fkey_reference_table';
+SELECT * FROM table_fkeys_in_workers WHERE relid LIKE 'fkey_reference_table.%' AND refd_relid LIKE 'fkey_reference_table.%' ORDER BY 1, 2;
+
+DROP SCHEMA fkey_reference_table CASCADE;
+SET search_path TO DEFAULT;