Fix inserts into local tables with distributed subqueries

src/backend/distributed/planner/multi_router_planner.c

@@ -547,14 +547,23 @@ DeferredErrorMessage *
 ModifyQuerySupported(Query *queryTree, Query *originalQuery, bool multiShardQuery,
 					 PlannerRestrictionContext *plannerRestrictionContext)
 {
-	Oid distributedTableId = ExtractFirstDistributedTableId(queryTree);
 	uint32 rangeTableId = 1;
-	Var *partitionColumn = PartitionColumn(distributedTableId, rangeTableId);
 	List *rangeTableList = NIL;
 	ListCell *rangeTableCell = NULL;
 	uint32 queryTableCount = 0;
 	CmdType commandType = queryTree->commandType;
 
+	Oid distributedTableId = ModifyQueryResultRelationId(queryTree);
+	if (!IsDistributedTable(distributedTableId))
+	{
+		return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
+							 "cannot plan modifications of local tables involving "
+							 "distributed tables",
+							 NULL, NULL);
+	}
+
+	Var *partitionColumn = PartitionColumn(distributedTableId, rangeTableId);
+
 	DeferredErrorMessage *deferredError = DeferErrorIfModifyView(queryTree);
 	if (deferredError != NULL)
 	{

src/test/regress/expected/dml_recursive.out

@@ -305,7 +305,7 @@ FROM
 	distributed_table
 WHERE
 	distributed_table.tenant_id = local_table.id;
-ERROR:  relation local_table is not distributed
+ERROR:  cannot plan modifications of local tables involving distributed tables
 RESET client_min_messages;
 DROP SCHEMA recursive_dml_queries CASCADE;
 NOTICE:  drop cascades to 5 other objects

src/test/regress/expected/multi_modifications.out

@@ -1,6 +1,7 @@
 SET citus.shard_count TO 32;
 SET citus.next_shard_id TO 750000;
 SET citus.next_placement_id TO 750000;
+CREATE SCHEMA multi_modifications;
 -- some failure messages that comes from the worker nodes
 -- might change due to parallel executions, so suppress those
 -- using \set VERBOSITY terse
@@ -1297,7 +1298,12 @@ ERROR:  relation pg_namespace is not distributed
 DELETE FROM summary_table WHERE id < (
     SELECT 0 FROM pg_dist_node
 );
+CREATE TABLE multi_modifications.local (a int default 1, b int);
+INSERT INTO multi_modifications.local VALUES (default, (SELECT min(id) FROM summary_table));
+ERROR:  cannot plan modifications of local tables involving distributed tables
 DROP TABLE raw_table;
 DROP TABLE summary_table;
 DROP TABLE reference_raw_table;
 DROP TABLE reference_summary_table;
+DROP SCHEMA multi_modifications CASCADE;
+NOTICE:  drop cascades to table multi_modifications.local

src/test/regress/expected/multi_shard_update_delete.out

@@ -730,7 +730,7 @@ UPDATE events_test_table_local
 SET    value_2 = 5
 FROM   users_test_table
 WHERE  events_test_table_local.user_id = users_test_table.user_id;
-ERROR:  relation events_test_table_local is not distributed
+ERROR:  cannot plan modifications of local tables involving distributed tables
 -- Local tables in a subquery are supported through recursive planning
 UPDATE users_test_table
 SET    value_2 = 5

src/test/regress/sql/multi_modifications.sql

@@ -2,6 +2,8 @@ SET citus.shard_count TO 32;
 SET citus.next_shard_id TO 750000;
 SET citus.next_placement_id TO 750000;
 
+CREATE SCHEMA multi_modifications;
+
 -- some failure messages that comes from the worker nodes
 -- might change due to parallel executions, so suppress those
 -- using \set VERBOSITY terse
@@ -871,7 +873,11 @@ DELETE FROM summary_table WHERE id < (
     SELECT 0 FROM pg_dist_node
 );
 
+CREATE TABLE multi_modifications.local (a int default 1, b int);
+INSERT INTO multi_modifications.local VALUES (default, (SELECT min(id) FROM summary_table));
+
 DROP TABLE raw_table;
 DROP TABLE summary_table;
 DROP TABLE reference_raw_table;
 DROP TABLE reference_summary_table;
+DROP SCHEMA multi_modifications CASCADE;