Add support for UPDATE ... WHERE IN

src/backend/distributed/planner/multi_router_planner.c

@@ -292,6 +292,11 @@ UpdateFromQuery(Query *query)
 	bool hasSubquery = false;
 	ListCell *rangeTableCell = NULL;
 
+	if (query->hasSubLinks)
+	{
+		hasSubquery = true;
+	}
+
 	foreach(rangeTableCell, rangeTableList)
 	{
 		RangeTblEntry *rangeTableEntry = (RangeTblEntry *) lfirst(rangeTableCell);
@@ -1275,7 +1280,7 @@ ModifyQuerySupported(Query *queryTree)
 	/*
 	 * Reject subqueries which are in SELECT or WHERE clause.
 	 */
-	if (queryTree->hasSubLinks == true)
+	if (queryTree->hasSubLinks == true && !UpdateFromQuery(queryTree))
 	{
 		return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
 							 "subqueries are not supported in distributed modifications",

src/test/regress/expected/multi_modifications.out

@@ -659,7 +659,7 @@ INSERT INTO app_analytics_events (app_id, name) VALUES (103, 'Mynt') RETURNING *
 DROP TABLE app_analytics_events;
 -- test UPDATE ... FROM
 CREATE TABLE raw_table (id bigint, value bigint);
-CREATE TABLE summary_table (id bigint, average_value numeric);
+CREATE TABLE summary_table (id bigint, min_value numeric, average_value numeric);
 SELECT create_distributed_table('raw_table', 'id');
  create_distributed_table 
 --------------------------
@@ -676,9 +676,9 @@ INSERT INTO raw_table VALUES (1, 100);
 INSERT INTO raw_table VALUES (1, 200);
 INSERT INTO summary_table VALUES (1, NULL);
 SELECT * FROM summary_table WHERE id = 1;
- id | average_value 
-----+---------------
-  1 |              
+ id | min_value | average_value 
+----+-----------+---------------
+  1 |           |              
 (1 row)
 
 UPDATE summary_table SET average_value = average_query.average FROM (
@@ -686,9 +686,17 @@ UPDATE summary_table SET average_value = average_query.average FROM (
 	) average_query
 WHERE id = 1;
 SELECT * FROM summary_table WHERE id = 1;
- id |    average_value     
-----+----------------------
-  1 | 150.0000000000000000
+ id | min_value |    average_value     
+----+-----------+----------------------
+  1 |           | 150.0000000000000000
+(1 row)
+
+UPDATE summary_table SET min_value = 100
+	WHERE id IN (SELECT id FROM raw_table WHERE id = 1 and value > 100) AND id = 1;
+SELECT * FROM summary_table WHERE id = 1;
+ id | min_value |    average_value     
+----+-----------+----------------------
+  1 |       100 | 150.0000000000000000
 (1 row)
 
 DROP TABLE raw_table;

src/test/regress/sql/multi_modifications.sql

@@ -432,7 +432,7 @@ DROP TABLE app_analytics_events;
 
 -- test UPDATE ... FROM
 CREATE TABLE raw_table (id bigint, value bigint);
-CREATE TABLE summary_table (id bigint, average_value numeric);
+CREATE TABLE summary_table (id bigint, min_value numeric, average_value numeric);
 
 SELECT create_distributed_table('raw_table', 'id');
 SELECT create_distributed_table('summary_table', 'id');
@@ -451,5 +451,10 @@ WHERE id = 1;
 
 SELECT * FROM summary_table WHERE id = 1;
 
+UPDATE summary_table SET min_value = 100
+	WHERE id IN (SELECT id FROM raw_table WHERE id = 1 and value > 100) AND id = 1;
+
+SELECT * FROM summary_table WHERE id = 1;
+
 DROP TABLE raw_table;
 DROP TABLE summary_table;