Allow references to columns in UPDATE statements

src/backend/distributed/planner/multi_router_planner.c

@@ -272,7 +272,7 @@ ErrorIfModifyQueryNotSupported(Query *queryTree)
 				continue;
 			}
 
-			if (!IsA(targetEntry->expr, Const))
+			if (contain_mutable_functions((Node *) targetEntry->expr))
 			{
 				hasNonConstTargetEntryExprs = true;
 			}

src/test/regress/expected/multi_expanded_modifications.out

@@ -0,0 +1,30 @@
+CREATE TABLE modifications (
+	key VARCHAR,
+	counter INT,
+	time TIMESTAMPTZ
+);
+SELECT master_create_distributed_table('modifications', 'key', 'hash');
+ master_create_distributed_table 
+---------------------------------
+ 
+(1 row)
+
+SELECT master_create_worker_shards('modifications', 4, 1);
+ master_create_worker_shards 
+-----------------------------
+ 
+(1 row)
+
+INSERT INTO modifications VALUES ('one', 0, null);
+UPDATE modifications SET counter = counter + 1;
+ERROR:  distributed modifications must target exactly one shard
+UPDATE modifications SET counter = counter + 1 WHERE key = 'one';
+UPDATE modifications SET time = now() WHERE key = 'one';
+ERROR:  cannot plan sharded modification containing values which are not constants or constant expressions
+SELECT * FROM modifications;
+ key | counter | time 
+-----+---------+------
+ one |       1 | 
+(1 row)
+
+DROP TABLE modifications;

src/test/regress/multi_schedule

@@ -141,3 +141,9 @@ test: multi_large_shardid
 # multi_drop_extension makes sure we can safely drop and recreate the extension
 # ----------
 test: multi_drop_extension
+
+# ----------
+# multi_expanded_modifications includes some tests that we can reference columns in
+# UPDATEs
+# ----------
+test: multi_expanded_modifications

src/test/regress/sql/multi_expanded_modifications.sql

@@ -0,0 +1,19 @@
+CREATE TABLE modifications (
+	key VARCHAR,
+	counter INT,
+	time TIMESTAMPTZ
+);
+SELECT master_create_distributed_table('modifications', 'key', 'hash');
+SELECT master_create_worker_shards('modifications', 4, 1);
+
+INSERT INTO modifications VALUES ('one', 0, null);
+
+UPDATE modifications SET counter = counter + 1;
+
+UPDATE modifications SET counter = counter + 1 WHERE key = 'one';
+
+UPDATE modifications SET time = now() WHERE key = 'one';
+
+SELECT * FROM modifications;
+
+DROP TABLE modifications;