Add hacky checks for unsupported queries

src/backend/distributed/planner/multi_router_planner.c

@@ -2361,6 +2361,22 @@ RouterSelectQuery(Query *originalQuery, RelationRestrictionContext *restrictionC
 
 	if (prunedRelationShardList == NULL)
 	{
+		if (UpdateFromQuery(originalQuery))
+		{
+			StringInfo errorMessage = makeStringInfo();
+			StringInfo errorHint = makeStringInfo();
+			appendStringInfo(errorMessage,
+							 "cannot run UPDATE command which targets multiple "
+							 "shards");
+
+			appendStringInfo(errorHint, "Make sure the value for partition column "
+										"falls into a single shard.");
+
+			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							errmsg("%s", errorMessage->data),
+							errhint("%s", errorHint->data)));
+		}
+
 		return false;
 	}
 
@@ -2376,6 +2392,22 @@ RouterSelectQuery(Query *originalQuery, RelationRestrictionContext *restrictionC
 		/* no shard is present or all shards are pruned out case will be handled later */
 		if (prunedShardList == NIL)
 		{
+			if (UpdateFromQuery(originalQuery))
+			{
+				StringInfo errorMessage = makeStringInfo();
+				StringInfo errorHint = makeStringInfo();
+				appendStringInfo(errorMessage,
+								 "cannot run UPDATE command which targets no "
+								 "shards");
+
+				appendStringInfo(errorHint, "Make sure the value for partition column "
+											"falls into a single shard.");
+
+				ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+								errmsg("%s", errorMessage->data),
+								errhint("%s", errorHint->data)));
+			}
+
 			continue;
 		}
 

src/test/regress/expected/multi_modifications.out

@@ -699,5 +699,33 @@ SELECT * FROM summary_table WHERE id = 1;
   1 |       100 | 150.0000000000000000
 (1 row)
 
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1 AND id = 4
+	) average_query
+WHERE id = 1 AND id = 4;
+ERROR:  cannot run UPDATE command which targets no shards
+HINT:  Make sure the value for partition column falls into a single shard.
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1 AND id = 4
+	) average_query
+WHERE id = 1;
+ERROR:  cannot run UPDATE command which targets no shards
+HINT:  Make sure the value for partition column falls into a single shard.
+SELECT * FROM summary_table WHERE id = 1;
+ id | min_value |    average_value     
+----+-----------+----------------------
+  1 |       100 | 150.0000000000000000
+(1 row)
+
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1
+	) average_query
+WHERE id = 1 AND id = 4;
+ERROR:  cannot run UPDATE command which targets no shards
+HINT:  Make sure the value for partition column falls into a single shard.
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table) average_query;
+ERROR:  cannot run UPDATE command which targets multiple shards
+HINT:  Make sure the value for partition column falls into a single shard.
 DROP TABLE raw_table;
 DROP TABLE summary_table;

src/test/regress/sql/multi_modifications.sql

@@ -456,5 +456,25 @@ UPDATE summary_table SET min_value = 100
 
 SELECT * FROM summary_table WHERE id = 1;
 
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1 AND id = 4
+	) average_query
+WHERE id = 1 AND id = 4;
+
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1 AND id = 4
+	) average_query
+WHERE id = 1;
+
+SELECT * FROM summary_table WHERE id = 1;
+
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table WHERE id = 1
+	) average_query
+WHERE id = 1 AND id = 4;
+
+UPDATE summary_table SET average_value = average_query.average FROM (
+	SELECT avg(value) AS average FROM raw_table) average_query;
+
 DROP TABLE raw_table;
 DROP TABLE summary_table;