Check more carrefuly UPDATE SET ()=(SELECT) query

Produce an ERROR when attribute in the target list have been reordered by PostgreSQL rewritter.
citusdata · Feb 27, 2025 · c1aa700 · c1aa700
1 parent 56736e1
commit c1aa700
Show file tree

Hide file tree

Showing 5 changed files with 223 additions and 104 deletions.
diff --git a/src/backend/distributed/deparser/ruleutils_14.c b/src/backend/distributed/deparser/ruleutils_14.c
@@ -3440,6 +3440,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
@@ -3562,11 +3565,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, " = ");

diff --git a/src/backend/distributed/deparser/ruleutils_15.c b/src/backend/distributed/deparser/ruleutils_15.c
@@ -3505,6 +3505,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
@@ -3627,11 +3630,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, " = ");

diff --git a/src/backend/distributed/deparser/ruleutils_16.c b/src/backend/distributed/deparser/ruleutils_16.c
@@ -3521,6 +3521,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
@@ -3643,11 +3646,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, " = ");