fix issue #7676: wrong handler around MULTIEXPR (#7914)

DESCRIPTION: Fixes a bug with `UPDATE SET (...) = (SELECT
some_func(),... )` (#7676)

Citus was checking for presence of sublink, but forgot to manage
multiexpr while evaluating clauses during planning. At this stage (citus
planner), it's not always possible to call PostgreSQL code because the
tree is not yet ready for PostgreSQL pure executor.

Fixes https://github.com/citusdata/citus/issues/7676.

Fixed by adding a new function to check sublink or multiexpr in the
tree.

---------

Co-authored-by: Colm <colmmchugh@microsoft.com>

src/backend/distributed/utils/citus_clauses.c

@@ -41,6 +41,7 @@ static bool ShouldEvaluateExpression(Expr *expression);
 static bool ShouldEvaluateFunctions(CoordinatorEvaluationContext *evaluationContext);
 static void FixFunctionArguments(Node *expr);
 static bool FixFunctionArgumentsWalker(Node *expr, void *context);
+static bool CheckExprExecutorSafe(Node *expr);
 
 
 /*
@@ -99,15 +100,18 @@ PartiallyEvaluateExpression(Node *expression,
 	}
 
 	NodeTag nodeTag = nodeTag(expression);
-	if (nodeTag == T_Param)
+
+	/* ExecInitExpr cannot handle some expressions (PARAM_MULTIEXPR and PARAM_SUBLINK) */
+	if (!CheckExprExecutorSafe(expression))
 	{
-		Param *param = (Param *) expression;
-		if (param->paramkind == PARAM_SUBLINK)
-		{
-			/* ExecInitExpr cannot handle PARAM_SUBLINK */
 		return expression;
 	}
 
+	/* ExecInitExpr cannot handle PARAM_MULTIEXPR and PARAM_SUBLINK but we have guards */
+	else if (nodeTag == T_Param)
+	{
+		Assert(((Param *) expression)->paramkind != PARAM_MULTIEXPR &&
+			   ((Param *) expression)->paramkind != PARAM_SUBLINK);
 		return (Node *) citus_evaluate_expr((Expr *) expression,
 											exprType(expression),
 											exprTypmod(expression),
@@ -260,8 +264,10 @@ ShouldEvaluateExpression(Expr *expression)
 		}
 
 		default:
+		{
 			return false;
 		}
+	}
 }
 
 
@@ -537,3 +543,48 @@ FixFunctionArgumentsWalker(Node *expr, void *context)
 
 	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
 }
+
+
+/*
+ * Recursively explore an expression to ensure it can be used in the PostgreSQL
+ * ExecInitExpr.
+ * Currently only search for PARAM_MULTIEXPR or PARAM_SUBLINK.
+ */
+static bool
+CheckExprExecutorSafe(Node *expr)
+{
+	if (expr == NULL)
+	{
+		return true;
+	}
+
+	/*
+	 * If it's a Param, we're done traversing the tree.
+	 * Just check if it contins a sublink or multiexpr.
+	 */
+	else if (IsA(expr, Param))
+	{
+		Param *param = (Param *) expr;
+		if (param->paramkind == PARAM_MULTIEXPR ||
+			param->paramkind == PARAM_SUBLINK)
+		{
+			return false;
+		}
+	}
+
+	/* If it's a FuncExpr, search in arguments */
+	else if (IsA(expr, FuncExpr))
+	{
+		FuncExpr *func = (FuncExpr *) expr;
+		ListCell *lc;
+
+		foreach(lc, func->args)
+		{
+			if (!CheckExprExecutorSafe((Node *) lfirst(lc)))
+			{
+				return false;
+			}
+		}
+	}
+	return true;
+}

src/test/regress/expected/multi_modifications.out

@@ -812,6 +812,33 @@ SELECT * FROM app_analytics_events ORDER BY id;
 (2 rows)
 
 DROP TABLE app_analytics_events;
+-- test function call in UPDATE SET
+-- https://github.com/citusdata/citus/issues/7676
+CREATE FUNCTION citus_is_coordinator_stable() returns bool as $$
+  select citus_is_coordinator();
+$$ language sql stable;
+CREATE TABLE bool_test (
+  id bigint primary key,
+  col_bool bool
+  );
+SELECT create_reference_table('bool_test');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO bool_test values (1, true);
+UPDATE bool_test
+SET (col_bool)
+  = (SELECT citus_is_coordinator_stable())
+RETURNING id, col_bool;
+ id | col_bool
+---------------------------------------------------------------------
+  1 | t
+(1 row)
+
+DROP TABLE bool_test;
+DROP FUNCTION citus_is_coordinator_stable();
 -- Test multi-row insert with serial in a non-partition column
 CREATE TABLE app_analytics_events (id int, app_id serial, name text);
 SELECT create_distributed_table('app_analytics_events', 'id');

src/test/regress/sql/multi_modifications.sql

@@ -505,6 +505,28 @@ VALUES (104, 'Wayz'), (105, 'Mynt') RETURNING *;
 SELECT * FROM app_analytics_events ORDER BY id;
 DROP TABLE app_analytics_events;
 
+-- test function call in UPDATE SET
+-- https://github.com/citusdata/citus/issues/7676
+CREATE FUNCTION citus_is_coordinator_stable() returns bool as $$
+  select citus_is_coordinator();
+$$ language sql stable;
+
+CREATE TABLE bool_test (
+  id bigint primary key,
+  col_bool bool
+  );
+SELECT create_reference_table('bool_test');
+
+INSERT INTO bool_test values (1, true);
+
+UPDATE bool_test
+SET (col_bool)
+  = (SELECT citus_is_coordinator_stable())
+RETURNING id, col_bool;
+
+DROP TABLE bool_test;
+DROP FUNCTION citus_is_coordinator_stable();
+
 -- Test multi-row insert with serial in a non-partition column
 CREATE TABLE app_analytics_events (id int, app_id serial, name text);
 SELECT create_distributed_table('app_analytics_events', 'id');