Repository: citusdata/citus
Pull request #7914: fix issue #7676: wrong handler around MULTIEXPR
Author: login: c2main
Created: 2025-03-03T07:24:05Z | Merged: 2025-03-12T11:01:33Z

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.




Diff:
diff --git a/src/backend/distributed/utils/citus_clauses.c b/src/backend/distributed/utils/citus_clauses.c
index f88b173afe0..7086cec4eab 100644
--- a/src/backend/distributed/utils/citus_clauses.c
+++ b/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;
-		}
+		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,7 +264,9 @@ 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;
+}
diff --git a/src/test/regress/expected/multi_modifications.out b/src/test/regress/expected/multi_modifications.out
index 887003a974b..93f6c8c4508 100644
--- a/src/test/regress/expected/multi_modifications.out
+++ b/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');
diff --git a/src/test/regress/sql/multi_modifications.sql b/src/test/regress/sql/multi_modifications.sql
index 7977325ea65..2a00e799281 100644
--- a/src/test/regress/sql/multi_modifications.sql
+++ b/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');


Reviews:
user: login: colm-mchugh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T09:16:56Z
user: login: c2main
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T09:19:09Z
user: login: colm-mchugh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T09:19:09Z
user: login: colm-mchugh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T09:21:12Z
user: login: c2main
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T09:27:26Z
user: login: colm-mchugh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:04:22Z
user: login: c2main
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T08:21:53Z
user: login: c2main
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T08:31:20Z
user: login: onurctirtir
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T09:11:23Z
user: login: c2main
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T10:35:25Z

Review comments:
diff_hunk: @@ -0,0 +1,24 @@
+-- https://github.com/citusdata/citus/issues/7676
+CREATE TABLE test_ref_multiexpr (
+    id bigint primary key
+  , col_int integer
+  , col_bool bool
+  , col_text text
+  , col_timestamp timestamp
+  );
+select create_reference_table('test_ref_multiexpr');
+
+/* TODO how to ensure in test that 'now()' is correctly pre-executed */
+insert into test_ref_multiexpr values (1, 1, true, 'one', now());
+
+update test_ref_multiexpr
+SET (col_timestamp)
+  = (SELECT now())
+returning id, col_int, col_bool;
+
+update test_ref_multiexpr
+SET (col_bool, col_timestamp)
+  = (SELECT true, now())
+returning id, col_int, col_bool;
+
+DROP TABLE test_ref_multiexpr;
path: src/test/regress/sql/issue_7676.sql
user: login: colm-mchugh
body: Can `multi_modifications.sql` be used for testing, instead of creating a new regress test ?
created_at: 2025-03-05T09:16:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,24 @@
+-- https://github.com/citusdata/citus/issues/7676
+CREATE TABLE test_ref_multiexpr (
+    id bigint primary key
+  , col_int integer
+  , col_bool bool
+  , col_text text
+  , col_timestamp timestamp
+  );
+select create_reference_table('test_ref_multiexpr');
+
+/* TODO how to ensure in test that 'now()' is correctly pre-executed */
+insert into test_ref_multiexpr values (1, 1, true, 'one', now());
+
+update test_ref_multiexpr
+SET (col_timestamp)
+  = (SELECT now())
+returning id, col_int, col_bool;
+
+update test_ref_multiexpr
+SET (col_bool, col_timestamp)
+  = (SELECT true, now())
+returning id, col_int, col_bool;
+
+DROP TABLE test_ref_multiexpr;
path: src/test/regress/sql/issue_7676.sql
user: login: c2main
body: I suppose yes, I will update.
created_at: 2025-03-05T09:19:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -537,3 +538,43 @@ FixFunctionArgumentsWalker(Node *expr, void *context)
 
 	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
 }
+
+
+/*
+ * Recursively search an expression for a Param with sublink
+ */
+static bool
+CheckContainsMultiexprOrSublink(Node *expr)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	/* If it's a Param, return its attnum */
path: src/backend/distributed/utils/citus_clauses.c
user: login: colm-mchugh
body: nit: should the comment read something like `If it's a Param, check for types that cannot be evaluated`
created_at: 2025-03-05T09:19:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -99,15 +100,15 @@ PartiallyEvaluateExpression(Node *expression,
 	}
 
 	NodeTag nodeTag = nodeTag(expression);
-	if (nodeTag == T_Param)
+	if (CheckContainsMultiexprOrSublink(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 */
+		return expression;
+	}
 
+	/* ExecInitExpr cannot handle PARAM_MULTIEXPR and PARAM_SUBLINK but we have guards */
+	else if (nodeTag == T_Param)
+	{
path: src/backend/distributed/utils/citus_clauses.c
user: login: colm-mchugh
body: Could assert at this point that its not a PARAM_MULTIEXPR or PARAM_SUBLINK
created_at: 2025-03-05T09:21:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -537,3 +538,43 @@ FixFunctionArgumentsWalker(Node *expr, void *context)
 
 	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
 }
+
+
+/*
+ * Recursively search an expression for a Param with sublink
+ */
+static bool
+CheckContainsMultiexprOrSublink(Node *expr)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	/* If it's a Param, return its attnum */
path: src/backend/distributed/utils/citus_clauses.c
user: login: c2main
body: oops, I'll update, more comments need an update. For example we traverse for `sublink` *or* `multiexpr`.
Maybe the function can be renamed, here it's very technical name, maybe more `CheckExprExecutorSafe()`, though it maybe should cover more cases, I didn't lookup in Postgres.
created_at: 2025-03-05T09:27:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -537,3 +538,43 @@ FixFunctionArgumentsWalker(Node *expr, void *context)
 
 	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
 }
+
+
+/*
+ * Recursively search an expression for a Param with sublink
+ */
+static bool
+CheckContainsMultiexprOrSublink(Node *expr)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	/* If it's a Param, return its attnum */
path: src/backend/distributed/utils/citus_clauses.c
user: login: colm-mchugh
body: Rename sounds good - from a quick look at `execExpr.c` (in Postgres source code) there's no other obvious cases that could be included, so the PR's function implementation is good as is - we can add cases as needed.
created_at: 2025-03-05T10:04:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -537,3 +538,43 @@ FixFunctionArgumentsWalker(Node *expr, void *context)
 
 	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
 }
+
+
+/*
+ * Recursively search an expression for a Param with sublink
+ */
+static bool
+CheckContainsMultiexprOrSublink(Node *expr)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	/* If it's a Param, return its attnum */
path: src/backend/distributed/utils/citus_clauses.c
user: login: c2main
body: I renamed and updated comments, also changed behavior (return true where is was false, and reverse)
created_at: 2025-03-07T08:21:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -99,15 +100,15 @@ PartiallyEvaluateExpression(Node *expression,
 	}
 
 	NodeTag nodeTag = nodeTag(expression);
-	if (nodeTag == T_Param)
+	if (CheckContainsMultiexprOrSublink(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 */
+		return expression;
+	}
 
+	/* ExecInitExpr cannot handle PARAM_MULTIEXPR and PARAM_SUBLINK but we have guards */
+	else if (nodeTag == T_Param)
+	{
path: src/backend/distributed/utils/citus_clauses.c
user: login: c2main
body: I did that
created_at: 2025-03-07T08:31:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -812,6 +812,42 @@ 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 TABLE test_ref_multiexpr (
+    id bigint primary key
+  , col_int integer
+  , col_bool bool
+  , col_text text
+  , col_timestamp timestamp
+  );
+SELECT create_reference_table('test_ref_multiexpr');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+/* TODO how to ensure in test that 'now()' is correctly pre-executed */
path: src/test/regress/expected/multi_modifications.out
user: login: onurctirtir
body: Below update command reproduces the same bug on main, so we can replace this with a more reliable test:
```sql
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);
set client_min_messages to 'DEBUG4';

update bool_test
SET (col_bool)
  = (SELECT citus_is_coordinator_stable())
returning id, col_bool;
```

If we're evaluating this in the coordinator -as expected-, then we should see "t".
created_at: 2025-03-07T09:11:23Z
author_association: MEMBER
diff_hunk: @@ -812,6 +812,42 @@ 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 TABLE test_ref_multiexpr (
+    id bigint primary key
+  , col_int integer
+  , col_bool bool
+  , col_text text
+  , col_timestamp timestamp
+  );
+SELECT create_reference_table('test_ref_multiexpr');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+/* TODO how to ensure in test that 'now()' is correctly pre-executed */
path: src/test/regress/expected/multi_modifications.out
user: login: c2main
body: ah great, thanks!
created_at: 2025-03-07T10:35:24Z
author_association: CONTRIBUTOR

PR comments:
user: login: colm-mchugh
created_at: 2025-03-05T09:22:06Z
body: > 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.
> 
> #7676
> 
> Fixed by adding a new function to check sublink or multiexpr in the tree.

Thanks for submitting this. Can you rebase the PR onto the release-13.0 branch ? We will consider it for the next patch release.

author_association: CONTRIBUTOR
user: login: c2main
created_at: 2025-03-07T08:34:05Z
body: > > 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.
> > #7676
> > Fixed by adding a new function to check sublink or multiexpr in the tree.
> 
> Thanks for submitting this. Can you rebase the PR onto the release-13.0 branch ? We will consider it for the next patch release.

I did address all the comments, and rebased. Not sure who is supposed to resolve conversation in github/citus workflow so I let them "unresolved"...
author_association: CONTRIBUTOR
user: login: onurctirtir
created_at: 2025-03-07T09:17:06Z
body: Also, before merging, please make sure to add a line that starts with `DESCRIPTION: Fixes a bug ..` into the PR description so that our release tooling can capture this into the release changelog.

You can check https://github.com/citusdata/citus/blob/main/CHANGELOG.md to see the expected format for changelog entries.
author_association: MEMBER
user: login: colm-mchugh
created_at: 2025-03-08T12:34:45Z
body: The PR looks good to me, but we need to fix [this failing check ](https://github.com/citusdata/citus/actions/runs/13719449855/job/38395082349?pr=7914)- do you know what's required to fix this @onurctirtir @m3hm3t ?
author_association: CONTRIBUTOR
user: login: c2main
created_at: 2025-03-11T09:37:36Z
body: @colm-mchugh I though I had a clean rebase, I can do it again to prevent the merge commit if you wish.
author_association: CONTRIBUTOR
user: login: colm-mchugh
created_at: 2025-03-12T09:27:56Z
body: > @colm-mchugh I though I had a clean rebase, I can do it again to prevent the merge commit if you wish.

No need to I think, I just wanted to rebase to check the build status, to see if it could be merged to get into the next patch release. But the upload-coverage check is still red, and I'm not yet sure what is the problem.. In any case we will aim to get this merged soon.
author_association: CONTRIBUTOR