Merge fd099d1ca4 into 4cd8bb1b67

src/backend/distributed/deparser/citus_ruleutils.c

@@ -1715,3 +1715,73 @@ RoleSpecString(RoleSpec *spec, bool withQuoteIdentifier)
 		}
 	}
 }
+
+
+/*
+ * list_sort comparator to sort target list by paramid (in MULTIEXPR)
+ * Intended for indirection management: UPDATE SET () = (SELECT )
+ */
+int
+target_list_cmp(const ListCell *a, const ListCell *b)
+{
+	TargetEntry *tleA = lfirst(a);
+	TargetEntry *tleB = lfirst(b);
+
+	if (IsA(tleA->expr, Param) && IsA(tleB->expr, Param))
+	{
+		int la = ((Param *) tleA->expr)->paramid;
+		int lb = ((Param *) tleB->expr)->paramid;
+		return (la > lb) - (la < lb);
+	}
+	else if ((IsA(tleA->expr, Param) && IsA(tleB->expr, SubLink)) ||
+			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, Param)) ||
+			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, SubLink)))
+	{
+		return -1;
+	}
+	else
+	{
+		elog(ERROR, "unexpected nodes");
+	}
+}
+
+
+/*
+ * Recursively search an expression for a Param and return its paramid
+ * Intended for indirection management: UPDATE SET () = (SELECT )
+ * Does not cover all options but those supported by Citus.
+ */
+int
+GetParamId(Node *expr)
+{
+	int paramid = 0;
+
+	if (expr == NULL)
+	{
+		return paramid;
+	}
+
+	/* If it's a Param, return its attnum */
+	if (IsA(expr, Param))
+	{
+		Param *param = (Param *) expr;
+		paramid = param->paramid;
+	}
+	/* If it's a FuncExpr, search in arguments */
+	else if (IsA(expr, FuncExpr))
+	{
+		FuncExpr *func = (FuncExpr *) expr;
+		ListCell *lc;
+
+		foreach(lc, func->args)
+		{
+			paramid = GetParamId((Node *) lfirst(lc));
+			if (paramid != 0)
+			{
+				break; /* Stop at the first valid paramid */
+			}
+		}
+	}
+
+	return paramid;
+}

src/backend/distributed/deparser/ruleutils_15.c

@@ -3509,6 +3509,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
+	AttrNumber	previous_attnum = InvalidAttrNumber;
+	int 		paramid_increment = 0;
+
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3631,11 +3634,48 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
+			AttrNumber attnum = InvalidAttrNumber;
+			if (IsA(expr, Param))
+			{
+				Param *param = (Param *) expr;
+				attnum = param->paramid + paramid_increment;
+			}
+			else if (IsA(expr, FuncExpr))
+			{
+				FuncExpr *func = (FuncExpr *) expr;
+				ListCell *lc;
+
+				/* Iterate through the arguments of the FuncExpr */
+				foreach(lc, func->args)
+				{
+					Node *arg = (Node *) lfirst(lc);
+
+					/* Check if the argument is a PARAM node */
+					if (IsA(arg, Param))
+					{
+						Param *param = (Param *) arg;
+						attnum = param->paramid + paramid_increment;
+
+						break;  /* Exit loop once we find the PARAM node */
+					}
+				}
+			}
+
+			if (previous_attnum >= attnum)
+				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+								errmsg(
+									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
+								errhint("Sort the columns on the left side by physical order.")));
+
+			previous_attnum = attnum;
+
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
+
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
+			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");

src/backend/distributed/deparser/ruleutils_16.c

@@ -3533,13 +3533,21 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	ma_sublinks = NIL;
 	if (query->hasSubLinks)		/* else there can't be any */
 	{
+		bool saw_junk = false;
+		bool need_to_sort_target_list = false;
+		int previous_paramid = 0;
+
 		foreach(l, targetList)
 		{
 			TargetEntry *tle = (TargetEntry *) lfirst(l);
 
+			// elog(WARNING, "TOP node to string: %s", nodeToString(tle->expr));
+			// elog(WARNING, "TOP node type: %d", (int) nodeTag(tle->expr));
+
 			if (tle->resjunk && IsA(tle->expr, SubLink))
 			{
 				SubLink    *sl = (SubLink *) tle->expr;
+				saw_junk = true;
 
 				if (sl->subLinkType == MULTIEXPR_SUBLINK)
 				{
@@ -3547,7 +3555,29 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 					Assert(sl->subLinkId == list_length(ma_sublinks));
 				}
 			}
+			else if (!tle->resjunk)
+			{
+				int paramid = 0;
+				Assert(!saw_junk);
+
+				paramid = GetParamId((Node *) tle->expr);
+				if (paramid < previous_paramid)
+					need_to_sort_target_list = true;
+
+				previous_paramid = paramid;
+			}
 		}
+
+		/*
+		 * reorder the target list on left side of the update:
+		 * SET () = (SELECT )
+		 * reordering the SELECT side only does not work, consider a case like:
+		 * SET (col_1, col3) = (SELECT 1, 3), (col_2) = (SELECT 2)
+		 * Then default order will lead to:
+		 * SET (col_1, col2) = (SELECT 1, 3), (col_3) = (SELECT 2)
+		 */
+		if (need_to_sort_target_list)
+			list_sort(targetList, target_list_cmp);
 	}
 	next_ma_cell = list_head(ma_sublinks);
 	cur_ma_sublink = NULL;
@@ -3658,6 +3688,7 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 
 		get_rule_expr(expr, context, false);
 	}
+	elog(DEBUG4, "rewriten query: %s", buf->data);
 }
 
 /* ----------

src/backend/distributed/deparser/ruleutils_17.c

@@ -3542,6 +3542,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
+	AttrNumber	previous_attnum = InvalidAttrNumber;
+	int 		paramid_increment = 0;
+
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3664,11 +3667,48 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
+			AttrNumber attnum = InvalidAttrNumber;
+			if (IsA(expr, Param))
+			{
+				Param *param = (Param *) expr;
+				attnum = param->paramid + paramid_increment;
+			}
+			else if (IsA(expr, FuncExpr))
+			{
+				FuncExpr *func = (FuncExpr *) expr;
+				ListCell *lc;
+
+				/* Iterate through the arguments of the FuncExpr */
+				foreach(lc, func->args)
+				{
+					Node *arg = (Node *) lfirst(lc);
+
+					/* Check if the argument is a PARAM node */
+					if (IsA(arg, Param))
+					{
+						Param *param = (Param *) arg;
+						attnum = param->paramid + paramid_increment;
+
+						break;  /* Exit loop once we find the PARAM node */
+					}
+				}
+			}
+
+			if (previous_attnum >= attnum)
+				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+								errmsg(
+									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
+								errhint("Sort the columns on the left side by physical order.")));
+
+			previous_attnum = attnum;
+
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
+
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
+			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");

src/include/distributed/citus_ruleutils.h

@@ -60,5 +60,8 @@ extern char * generate_operator_name(Oid operid, Oid arg1, Oid arg2);
 extern List * getOwnedSequences_internal(Oid relid, AttrNumber attnum, char deptype);
 extern void AppendOptionListToString(StringInfo stringData, List *options);
 
+extern int GetParamId(Node *expr);
+extern int target_list_cmp(const ListCell *a, const ListCell *b);
+
 
 #endif /* CITUS_RULEUTILS_H */

src/test/regress/expected/distributed_types.out

@@ -444,6 +444,7 @@ HINT:  Use the column name to insert or update the composite type as a single va
 UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
 ERROR:  inserting or modifying composite type fields is not supported
 HINT:  Use the column name to insert or update the composite type as a single value
+HINT:  Use the column name to insert or update the composite type as a single value
 -- below are supported as we don't do any field indirection
 INSERT INTO field_indirection_test_2 (ct2_col, int_col, ct1_col)
 VALUES ('(1, "text1", 2)', 3, '(4, 5)'), ('(6, "text2", 7)', 8, '(9, 10)');

src/test/regress/expected/indirections.out

@@ -0,0 +1,447 @@
+SET citus.shard_count TO 2;
+SET citus.next_shard_id TO 750000;
+SET citus.next_placement_id TO 750000;
+CREATE SCHEMA indirections;
+SET search_path TO indirections;
+-- specific tests related to get_update_query_targetlist_def
+-- we test only queries with sublinks, like:
+-- ( ... SET (...) = (SELECT ...))
+-- Reference tables
+CREATE TABLE test_ref_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE test_ref_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection_new');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- Distributed tables
+CREATE TABLE test_dist_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE test_dist_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- Local tables required ?
+-- those should work:
+INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_dist_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_ref_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_dist_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_ref_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection_new SELECT * FROM test_ref_indirection;
+INSERT INTO test_dist_indirection_new SELECT * FROM test_dist_indirection;
+SELECT * FROM test_ref_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_dist_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_ref_indirection_new ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_dist_indirection_new ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+-- now UPDATEs
+UPDATE test_ref_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1971 |       2 | 2 ok
+  2 | f        | 01-01-1971 |       2 | 2 ok
+  3 | f        | 01-01-1971 |       2 | 2 ok
+  4 | f        | 01-01-1971 |       2 | 2 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1971 |       2 | 2 ok
+  2 | f        | 01-01-1971 |       2 | 2 ok
+  3 | f        | 01-01-1971 |       2 | 2 ok
+  4 | f        | 01-01-1971 |       2 | 2 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1971 |       3 | 3 ok
+  2 | t        | 01-01-1971 |       3 | 3 ok
+  3 | t        | 01-01-1971 |       3 | 3 ok
+  4 | t        | 01-01-1971 |       3 | 3 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1971 |       3 | 3 ok
+  2 | t        | 01-01-1971 |       3 | 3 ok
+  3 | t        | 01-01-1971 |       3 | 3 ok
+  4 | t        | 01-01-1971 |       3 | 3 ok
+(4 rows)
+
+-- but those should work since 13.X
+UPDATE test_ref_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1972 |       4 | 4 ok
+  2 | f        | 01-01-1972 |       4 | 4 ok
+  3 | f        | 01-01-1972 |       4 | 4 ok
+  4 | f        | 01-01-1972 |       4 | 4 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1972 |       4 | 4 ok
+  2 | f        | 01-01-1972 |       4 | 4 ok
+  3 | f        | 01-01-1972 |       4 | 4 ok
+  4 | f        | 01-01-1972 |       4 | 4 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1972 |       5 | 5 ok
+  2 | t        | 01-01-1972 |       5 | 5 ok
+  3 | t        | 01-01-1972 |       5 | 5 ok
+  4 | t        | 01-01-1972 |       5 | 5 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1972 |       5 | 5 ok
+  2 | t        | 01-01-1972 |       5 | 5 ok
+  3 | t        | 01-01-1972 |       5 | 5 ok
+  4 | t        | 01-01-1972 |       5 | 5 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1973 |       6 | 6 ok
+  2 | f        | 01-01-1973 |       6 | 6 ok
+  3 | f        | 01-01-1973 |       6 | 6 ok
+  4 | f        | 01-01-1973 |       6 | 6 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1973 |       6 | 6 ok
+  2 | f        | 01-01-1973 |       6 | 6 ok
+  3 | f        | 01-01-1973 |       6 | 6 ok
+  4 | f        | 01-01-1973 |       6 | 6 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1974 |       7 | 7 ok
+  2 | t        | 01-01-1974 |       7 | 7 ok
+  3 | t        | 01-01-1974 |       7 | 7 ok
+  4 | t        | 01-01-1974 |       7 | 7 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1974 |       7 | 7 ok
+  2 | t        | 01-01-1974 |       7 | 7 ok
+  3 | t        | 01-01-1974 |       7 | 7 ok
+  4 | t        | 01-01-1974 |       7 | 7 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 8 ok
+  2 | f        | 01-01-1975 |       8 | 8 ok
+  3 | f        | 01-01-1975 |       8 | 8 ok
+  4 | f        | 01-01-1975 |       8 | 8 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 8 ok
+  2 | f        | 01-01-1975 |       8 | 8 ok
+  3 | f        | 01-01-1975 |       8 | 8 ok
+  4 | f        | 01-01-1975 |       8 | 8 ok
+(4 rows)
+
+--
+-- more restrictive ones, just in case we miss a wrong value
+--
+-- those should work
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1975 |       8 | 9 ok
+  2 | t        | 01-01-1975 |       8 | 9 ok
+  3 | t        | 01-01-1975 |       8 | 9 ok
+  4 | t        | 01-01-1975 |       8 | 9 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1975 |       8 | 9 ok
+  2 | t        | 01-01-1975 |       8 | 9 ok
+  3 | t        | 01-01-1975 |       8 | 9 ok
+  4 | t        | 01-01-1975 |       8 | 9 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
+WHERE id = 1
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 10 ok
+(1 row)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
+WHERE id = 1
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 10 ok
+(1 row)
+
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1975 |       8 | 11 ok
+  2 | t        | 01-01-1975 |       8 | 11 ok
+  3 | t        | 01-01-1975 |       8 | 11 ok
+  4 | t        | 01-01-1975 |       8 | 11 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1975 |       8 | 11 ok
+  2 | t        | 01-01-1975 |       8 | 11 ok
+  3 | t        | 01-01-1975 |       8 | 11 ok
+  4 | t        | 01-01-1975 |       8 | 11 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
+WHERE id = 2
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  2 | f        | 01-01-1975 |       8 | 12 ok
+(1 row)
+
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
+WHERE id = 2
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  2 | f        | 01-01-1975 |       8 | 12 ok
+(1 row)
+
+-- several updates in CTE shoult not work
+with qq3 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT '13', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT '14', false)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  3 | t        | 01-01-1975 |       8 | 13
+  4 | f        | 01-01-1975 |       8 | 14
+(2 rows)
+
+with qq3 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT '13', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT '14', false)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  3 | t        | 01-01-1975 |       8 | 13
+  4 | f        | 01-01-1975 |       8 | 14
+(2 rows)
+
+DROP TABLE test_dist_indirection;
+DROP TABLE test_dist_indirection_new;
+DROP TABLE test_ref_indirection;
+DROP TABLE test_ref_indirection_new;
+-- https://github.com/citusdata/citus/issues/4092
+CREATE TABLE update_test (
+    a   INT DEFAULT 10,
+    b   INT,
+    c   TEXT
+);
+SELECT create_reference_table('indirections.update_test');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+UPDATE update_test
+SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
+WHERE a = 100 AND b = 20;
+-- https://github.com/citusdata/citus/pull/5692
+set client_min_messages to ERROR;
+DROP SCHEMA indirections CASCADE;

src/test/regress/expected/multi_modifications.out

@@ -330,6 +330,12 @@ UPDATE limit_orders SET (kind, limit_price) = ('buy', 999) WHERE id = 246 RETURN
  246 | GM     |        30 | Mon Jul 02 16:32:15 2007 | buy  |         999
 (1 row)
 
+UPDATE limit_orders SET (kind, limit_price) = (SELECT 'buy'::order_side, 999) WHERE id = 246 RETURNING *;
+ id  | symbol | bidder_id |        placed_at         | kind | limit_price
+---------------------------------------------------------------------
+ 246 | GM     |        30 | Mon Jul 02 16:32:15 2007 | buy  |         999
+(1 row)
+
 -- Test that on unique contraint violations, we fail fast
 \set VERBOSITY terse
 INSERT INTO limit_orders VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sell', 43.67);
@@ -435,6 +441,8 @@ UPDATE limit_orders SET limit_price = 0.00 FROM bidders
 -- should succeed with a CTE
 WITH deleted_orders AS (INSERT INTO limit_orders VALUES (399, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43))
 UPDATE limit_orders SET symbol = 'GM';
+WITH deleted_orders AS (INSERT INTO limit_orders SELECT 400, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43)
+UPDATE limit_orders SET symbol = 'GM';
 SELECT symbol, bidder_id FROM limit_orders WHERE id = 246;
  symbol | bidder_id
 ---------------------------------------------------------------------
@@ -927,6 +935,17 @@ SELECT * FROM summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
+--- TODO this one is a silent corruption:
+-- UPDATE summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+-- WHERE id = 2;
+SELECT * FROM summary_table ORDER BY id;
+ id | min_value |    average_value     | count | uniques
+---------------------------------------------------------------------
+  1 |           | 200.0000000000000000 |       |
+  2 |       400 | 450.0000000000000000 |       |
+(2 rows)
+
 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 ORDER BY id;
@@ -1103,6 +1122,17 @@ SELECT * FROM reference_summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
+--- TODO this one is a silent corruption:
+-- UPDATE reference_summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+-- WHERE id = 2;
+SELECT * FROM reference_summary_table ORDER BY id;
+ id | min_value |    average_value     | count | uniques
+---------------------------------------------------------------------
+  1 |           | 200.0000000000000000 |       |
+  2 |       400 | 450.0000000000000000 |       |
+(2 rows)
+
 -- no need partition colum equalities on reference tables
 UPDATE reference_summary_table SET (count) =
 	(SELECT id AS inner_id FROM reference_raw_table WHERE value = 500)

src/test/regress/expected/multi_mx_modifications.out

@@ -91,6 +91,8 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
+INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
+								 'sell', 0.58;
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
 -- commands with mutable functions in their quals
@@ -103,6 +105,22 @@ DELETE FROM limit_orders_mx WHERE id = 246 AND placed_at = current_timestamp::ti
 INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+INSERT INTO limit_orders_mx SELECT * FROM (VALUES
+                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
 -- connect back to the other node
 \c - - - :worker_1_port
 -- commands containing a CTE are supported
@@ -216,9 +234,12 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 ERROR:  duplicate key value violates unique constraint "limit_orders_mx_pkey_1220093"
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
+UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
 ERROR:  modifying the partition value of rows is not allowed
+UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
+ERROR:  modifying the partition value of rows is not allowed
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders
 					WHERE limit_orders_mx.id = 246 AND

src/test/regress/sql/multi_modifications.sql

@@ -234,6 +234,7 @@ SELECT kind, limit_price FROM limit_orders WHERE id = 246;
 
 -- multi-column UPDATE with RETURNING
 UPDATE limit_orders SET (kind, limit_price) = ('buy', 999) WHERE id = 246 RETURNING *;
+UPDATE limit_orders SET (kind, limit_price) = (SELECT 'buy'::order_side, 999) WHERE id = 246 RETURNING *;
 
 -- Test that on unique contraint violations, we fail fast
 \set VERBOSITY terse
@@ -337,6 +338,9 @@ UPDATE limit_orders SET limit_price = 0.00 FROM bidders
 WITH deleted_orders AS (INSERT INTO limit_orders VALUES (399, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43))
 UPDATE limit_orders SET symbol = 'GM';
 
+WITH deleted_orders AS (INSERT INTO limit_orders SELECT 400, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43)
+UPDATE limit_orders SET symbol = 'GM';
+
 SELECT symbol, bidder_id FROM limit_orders WHERE id = 246;
 
 -- updates referencing just a var are supported
@@ -584,6 +588,13 @@ WHERE id = 2;
 
 SELECT * FROM summary_table ORDER BY id;
 
+--- TODO this one is a silent corruption:
+-- UPDATE summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+-- WHERE id = 2;
+
+SELECT * FROM summary_table ORDER BY id;
+
 UPDATE summary_table SET min_value = 100
 	WHERE id IN (SELECT id FROM raw_table WHERE id = 1 and value > 100) AND id = 1;
 
@@ -712,6 +723,13 @@ WHERE id = 2;
 
 SELECT * FROM reference_summary_table ORDER BY id;
 
+--- TODO this one is a silent corruption:
+-- UPDATE reference_summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+-- WHERE id = 2;
+
+SELECT * FROM reference_summary_table ORDER BY id;
+
 -- no need partition colum equalities on reference tables
 UPDATE reference_summary_table SET (count) =
 	(SELECT id AS inner_id FROM reference_raw_table WHERE value = 500)

src/test/regress/sql/multi_mx_modifications.sql

@@ -67,6 +67,8 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
+INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
+								 'sell', 0.58;
 
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
@@ -83,6 +85,17 @@ INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
 
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
+
+INSERT INTO limit_orders_mx SELECT * FROM (VALUES
+                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
+
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
+
+
+
 -- connect back to the other node
 \c - - - :worker_1_port
 
@@ -153,9 +166,11 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
+UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
+UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
 
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders

src/test/regress/sql/multi_update_select.sql

@@ -0,0 +1,228 @@
+SET citus.shard_count TO 2;
+SET citus.next_shard_id TO 750000;
+SET citus.next_placement_id TO 750000;
+
+CREATE SCHEMA indirections;
+SET search_path TO indirections;
+
+-- specific tests related to get_update_query_targetlist_def
+-- we test only queries with sublinks, like:
+-- ( ... SET (...) = (SELECT ...))
+
+-- Reference tables
+CREATE TABLE test_ref_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection');
+
+CREATE TABLE test_ref_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection_new');
+
+-- Distributed tables
+CREATE TABLE test_dist_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+
+CREATE TABLE test_dist_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
+
+-- Local tables required ?
+
+-- those should work:
+INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_dist_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+
+INSERT INTO test_ref_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_dist_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+
+INSERT INTO test_ref_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+
+INSERT INTO test_ref_indirection_new SELECT * FROM test_ref_indirection;
+INSERT INTO test_dist_indirection_new SELECT * FROM test_dist_indirection;
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+SELECT * FROM test_dist_indirection ORDER BY id;
+
+SELECT * FROM test_ref_indirection_new ORDER BY id;
+SELECT * FROM test_dist_indirection_new ORDER BY id;
+
+-- now UPDATEs
+UPDATE test_ref_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
+
+-- but those should work since 13.X
+UPDATE test_ref_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
+RETURNING *;
+
+--
+-- more restrictive ones, just in case we miss a wrong value
+--
+-- those should work
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
+WHERE id = 1
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
+WHERE id = 1
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
+WHERE id = 2
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
+WHERE id = 2
+RETURNING *;
+
+-- several updates in CTE shoult not work
+with qq3 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT '13', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT '14', false)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+with qq3 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT '13', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT '14', false)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+
+DROP TABLE test_dist_indirection;
+DROP TABLE test_dist_indirection_new;
+DROP TABLE test_ref_indirection;
+DROP TABLE test_ref_indirection_new;
+
+-- https://github.com/citusdata/citus/issues/4092
+CREATE TABLE update_test (
+    a   INT DEFAULT 10,
+    b   INT,
+    c   TEXT
+);
+SELECT create_reference_table('indirections.update_test');
+UPDATE update_test
+SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
+WHERE a = 100 AND b = 20;
+
+-- https://github.com/citusdata/citus/pull/5692
+set client_min_messages to ERROR;
+DROP SCHEMA indirections CASCADE;
+