Enable repartitioned INSERT/SELECT ON CONFLICT.

src/backend/distributed/executor/insert_select_executor.c

@@ -874,7 +874,7 @@ PartitionColumnIndex(List *insertTargetList, Var *partitionColumn)
 static bool
 IsRedistributablePlan(Plan *selectPlan, bool hasReturning, bool hasOnConflict)
 {
-	if (hasReturning || hasOnConflict)
+	if (hasReturning)
 	{
 		return false;
 	}

src/test/regress/expected/insert_select_repartition.out

@@ -68,7 +68,7 @@ INSERT INTO source_table VALUES (4, (30, 'e'), 5, (30, 'e'));
 INSERT INTO source_table VALUES (5, (31, 'f'), 6, (31, 'f'));
 INSERT INTO source_table VALUES (6, (32, 'g'), 50, (8, 'g'));
 -- target
-CREATE TABLE target_table(f1 int DEFAULT 0, value int, key composite_key_type);
+CREATE TABLE target_table(f1 int DEFAULT 0, value int, key composite_key_type PRIMARY KEY);
 SELECT create_distributed_table('target_table', 'key', 'range');
 NOTICE:  using statement-based replication
 DETAIL:  Streaming replication is supported only for hash-distributed tables.
@@ -154,6 +154,32 @@ SELECT * FROM target_table ORDER BY key;
   0 |       | (31,f)
 (7 rows)
 
+-- ON CONFLICT
+SET client_min_messages TO DEBUG2;
+INSERT INTO target_table(key)
+SELECT mapped_key AS key_renamed FROM source_table
+WHERE (mapped_key).f1 % 2 = 1
+ON CONFLICT (key) DO UPDATE SET f1=1;
+DEBUG:  cannot perform distributed INSERT INTO ... SELECT because the partition columns in the source table and subquery do not match
+DETAIL:  The target table's partition column should correspond to a partition column in the subquery.
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  performing repartitioned INSERT ... SELECT
+DEBUG:  partitioning SELECT query by column index 1 with name 'key_renamed'
+DEBUG:  distributed statement: INSERT INTO insert_select_repartition.target_table_4213591 AS citus_table_alias (f1, key) SELECT f1, key_renamed FROM read_intermediate_results('{repartitioned_results_from_4213589_to_0}'::text[], 'text'::citus_copy_format) intermediate_result(f1 integer, key_renamed insert_select_repartition.composite_key_type) ON CONFLICT(key) DO UPDATE SET f1 = 1
+DEBUG:  distributed statement: INSERT INTO insert_select_repartition.target_table_4213592 AS citus_table_alias (f1, key) SELECT f1, key_renamed FROM read_intermediate_results('{repartitioned_results_from_4213589_to_1,repartitioned_results_from_4213590_to_1}'::text[], 'text'::citus_copy_format) intermediate_result(f1 integer, key_renamed insert_select_repartition.composite_key_type) ON CONFLICT(key) DO UPDATE SET f1 = 1
+RESET client_min_messages;
+SELECT * FROM target_table ORDER BY key;
+ f1 | value |  key
+---------------------------------------------------------------------
+  0 |       | (0,a)
+  1 |       | (3,c)
+  0 |       | (8,g)
+  0 |       | (26,b)
+  1 |       | (27,d)
+  0 |       | (30,e)
+  1 |       | (31,f)
+(7 rows)
+
 -- missing value for distribution column
 INSERT INTO target_table(value) SELECT value FROM source_table;
 ERROR:  the partition column of table insert_select_repartition.target_table should have a value

src/test/regress/sql/insert_select_repartition.sql

@@ -43,7 +43,7 @@ INSERT INTO source_table VALUES (5, (31, 'f'), 6, (31, 'f'));
 INSERT INTO source_table VALUES (6, (32, 'g'), 50, (8, 'g'));
 
 -- target
-CREATE TABLE target_table(f1 int DEFAULT 0, value int, key composite_key_type);
+CREATE TABLE target_table(f1 int DEFAULT 0, value int, key composite_key_type PRIMARY KEY);
 SELECT create_distributed_table('target_table', 'key', 'range');
 CALL public.create_range_partitioned_shards('target_table', '{"(0,a)","(25,a)"}','{"(24,z)","(49,z)"}');
 
@@ -68,8 +68,38 @@ INSERT INTO target_table(key) SELECT mapped_key AS key_renamed FROM source_table
 RESET client_min_messages;
 SELECT * FROM target_table ORDER BY key;
 
+-- ON CONFLICT
+SET client_min_messages TO DEBUG2;
+INSERT INTO target_table(key)
+SELECT mapped_key AS key_renamed FROM source_table
+WHERE (mapped_key).f1 % 2 = 1
+ON CONFLICT (key) DO UPDATE SET f1=1;
+RESET client_min_messages;
+SELECT * FROM target_table ORDER BY key;
+
 -- missing value for distribution column
 INSERT INTO target_table(value) SELECT value FROM source_table;
+DROP TABLE source_table, target_table;
+
+-- different column types
+CREATE TABLE target_table(col_1 int primary key, col_2 int);
+SELECT create_distributed_table('target_table','col_1');
+INSERT INTO target_table VALUES(1,2),(2,3),(3,4),(4,5),(5,6);
+
+CREATE TABLE source_table(col_1 numeric, col_2 numeric, col_3 numeric);
+SELECT create_distributed_table('source_table','col_1');
+INSERT INTO source_table VALUES(1,1,1),(2,2,2),(3,3,3),(4,4,4),(5,5,5);
+
+SET client_min_messages TO DEBUG2;
+INSERT INTO target_table
+SELECT
+	col_1, col_2
+FROM
+	source_table
+ON CONFLICT(col_1) DO UPDATE SET col_2 = EXCLUDED.col_2;
+RESET client_min_messages;
+
+SELECT * FROM target_table ORDER BY 1;
 
 DROP TABLE source_table, target_table;