From c97692f382bdabf3fb764c70ed85fb72ffee1bf7 Mon Sep 17 00:00:00 2001 From: Marco Slot Date: Wed, 23 Aug 2017 13:16:22 +0200 Subject: [PATCH] Fix multi-row INSERT with RETURNING on reference tables --- .../distributed/executor/multi_executor.c | 23 +-------- .../distributed/planner/deparse_shard_query.c | 51 ++----------------- .../planner/multi_router_planner.c | 46 ++++++++++++++++- .../distributed/multi_router_planner.h | 1 + .../regress/expected/multi_modifications.out | 11 ++-- src/test/regress/sql/multi_modifications.sql | 4 +- 6 files changed, 60 insertions(+), 76 deletions(-) diff --git a/src/backend/distributed/executor/multi_executor.c b/src/backend/distributed/executor/multi_executor.c index 2200f62d5..d83b7a3e0 100644 --- a/src/backend/distributed/executor/multi_executor.c +++ b/src/backend/distributed/executor/multi_executor.c @@ -22,6 +22,7 @@ #include "distributed/multi_master_planner.h" #include "distributed/multi_planner.h" #include "distributed/multi_router_executor.h" +#include "distributed/multi_router_planner.h" #include "distributed/multi_resowner.h" #include "distributed/multi_server_executor.h" #include "distributed/multi_utility.h" @@ -204,27 +205,7 @@ RouterCreateScan(CustomScan *scan) static bool IsMultiRowInsert(Query *query) { - ListCell *rteCell = NULL; - bool hasValuesRTE = false; - - CmdType commandType = query->commandType; - if (commandType != CMD_INSERT) - { - return false; - } - - foreach(rteCell, query->rtable) - { - RangeTblEntry *rte = (RangeTblEntry *) lfirst(rteCell); - - if (rte->rtekind == RTE_VALUES) - { - hasValuesRTE = true; - break; - } - } - - return hasValuesRTE; + return ExtractDistributedInsertValuesRTE(query) != NULL; } diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index c06ff794c..dc06ff81a 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -31,8 +31,7 @@ #include "utils/lsyscache.h" #include "utils/rel.h" -static RangeTblEntry * ExtractDistributedInsertValuesRTE(Query *query, - Oid distributedTableId); + static void UpdateTaskQueryString(Query *query, Oid distributedTableId, RangeTblEntry *valuesRTE, Task *task); static void ConvertRteToSubqueryWithEmptyResult(RangeTblEntry *rte); @@ -47,8 +46,7 @@ RebuildQueryStrings(Query *originalQuery, List *taskList) { ListCell *taskCell = NULL; Oid relationId = ((RangeTblEntry *) linitial(originalQuery->rtable))->relid; - RangeTblEntry *valuesRTE = ExtractDistributedInsertValuesRTE(originalQuery, - relationId); + RangeTblEntry *valuesRTE = ExtractDistributedInsertValuesRTE(originalQuery); foreach(taskCell, taskList) { @@ -108,50 +106,6 @@ RebuildQueryStrings(Query *originalQuery, List *taskList) } -/* - * ExtractDistributedInsertValuesRTE does precisely that. If the provided - * query is not an INSERT, or if the table is a reference table, or if the - * INSERT does not have a VALUES RTE (i.e. it is not a multi-row INSERT), this - * function returns NULL. If all those conditions are met, an RTE representing - * the multiple values of a multi-row INSERT is returned. - */ -static RangeTblEntry * -ExtractDistributedInsertValuesRTE(Query *query, Oid distributedTableId) -{ - RangeTblEntry *valuesRTE = NULL; - uint32 rangeTableId = 1; - Var *partitionColumn = NULL; - TargetEntry *targetEntry = NULL; - - if (query->commandType != CMD_INSERT) - { - return NULL; - } - - partitionColumn = PartitionColumn(distributedTableId, rangeTableId); - if (partitionColumn == NULL) - { - return NULL; - } - - targetEntry = get_tle_by_resno(query->targetList, partitionColumn->varattno); - Assert(targetEntry != NULL); - - if (IsA(targetEntry->expr, Var)) - { - Var *partitionVar = (Var *) targetEntry->expr; - - valuesRTE = rt_fetch(partitionVar->varno, query->rtable); - if (valuesRTE->rtekind != RTE_VALUES) - { - return NULL; - } - } - - return valuesRTE; -} - - /* * UpdateTaskQueryString updates the query string stored within the provided * Task. If the Task has row values from a multi-row INSERT, those are injected @@ -169,6 +123,7 @@ UpdateTaskQueryString(Query *query, Oid distributedTableId, RangeTblEntry *value if (valuesRTE != NULL) { Assert(valuesRTE->rtekind == RTE_VALUES); + Assert(task->rowValuesLists != NULL); oldValuesLists = valuesRTE->values_lists; valuesRTE->values_lists = task->rowValuesLists; diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 284496427..60c6d043c 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -1854,6 +1854,7 @@ BuildRoutesForInsert(Query *query, DeferredErrorMessage **planningError) { int shardCount = 0; List *shardIntervalList = LoadShardIntervalList(distributedTableId); + RangeTblEntry *valuesRTE = NULL; ShardInterval *shardInterval = NULL; ModifyRoute *modifyRoute = NULL; @@ -1867,7 +1868,17 @@ BuildRoutesForInsert(Query *query, DeferredErrorMessage **planningError) modifyRoute = palloc(sizeof(ModifyRoute)); modifyRoute->shardId = shardInterval->shardId; - modifyRoute->rowValuesLists = NIL; + + valuesRTE = ExtractDistributedInsertValuesRTE(query); + if (valuesRTE != NULL) + { + /* add the values list for a multi-row INSERT */ + modifyRoute->rowValuesLists = valuesRTE->values_lists; + } + else + { + modifyRoute->rowValuesLists = NIL; + } modifyRouteList = lappend(modifyRouteList, modifyRoute); @@ -1986,6 +1997,39 @@ BuildRoutesForInsert(Query *query, DeferredErrorMessage **planningError) } +/* + * ExtractDistributedInsertValuesRTE does precisely that. If the provided + * query is not an INSERT, or if the INSERT does not have a VALUES RTE + * (i.e. it is not a multi-row INSERT), this function returns NULL. + * If all those conditions are met, an RTE representing the multiple values + * of a multi-row INSERT is returned. + */ +RangeTblEntry * +ExtractDistributedInsertValuesRTE(Query *query) +{ + ListCell *rteCell = NULL; + RangeTblEntry *valuesRTE = NULL; + + if (query->commandType != CMD_INSERT) + { + return NULL; + } + + foreach(rteCell, query->rtable) + { + RangeTblEntry *rte = (RangeTblEntry *) lfirst(rteCell); + + if (rte->rtekind == RTE_VALUES) + { + valuesRTE = rte; + break; + } + } + + return valuesRTE; +} + + /* * IntersectPlacementList performs placement pruning based on matching on * nodeName:nodePort fields of shard placement data. We start pruning from all diff --git a/src/include/distributed/multi_router_planner.h b/src/include/distributed/multi_router_planner.h index 169d271e7..3ff9f048a 100644 --- a/src/include/distributed/multi_router_planner.h +++ b/src/include/distributed/multi_router_planner.h @@ -48,6 +48,7 @@ extern RelationRestrictionContext * CopyRelationRestrictionContext( extern Oid ExtractFirstDistributedTableId(Query *query); extern RangeTblEntry * ExtractSelectRangeTableEntry(Query *query); extern RangeTblEntry * ExtractInsertRangeTableEntry(Query *query); +extern RangeTblEntry * ExtractDistributedInsertValuesRTE(Query *query); extern void AddShardIntervalRestrictionToSelect(Query *subqery, ShardInterval *shardInterval); diff --git a/src/test/regress/expected/multi_modifications.out b/src/test/regress/expected/multi_modifications.out index 874b54a42..0d235119a 100644 --- a/src/test/regress/expected/multi_modifications.out +++ b/src/test/regress/expected/multi_modifications.out @@ -888,9 +888,14 @@ SELECT create_reference_table('reference_summary_table'); INSERT INTO reference_raw_table VALUES (1, 100); INSERT INTO reference_raw_table VALUES (1, 200); INSERT INTO reference_raw_table VALUES (1, 200); -INSERT INTO reference_raw_table VALUES (1, 300); -INSERT INTO reference_raw_table VALUES (2, 400); -INSERT INTO reference_raw_table VALUES (2, 500); +INSERT INTO reference_raw_table VALUES (1,300), (2, 400), (2,500) RETURNING *; + id | value +----+------- + 1 | 300 + 2 | 400 + 2 | 500 +(3 rows) + INSERT INTO reference_summary_table VALUES (1); INSERT INTO reference_summary_table VALUES (2); SELECT * FROM reference_summary_table ORDER BY id; diff --git a/src/test/regress/sql/multi_modifications.sql b/src/test/regress/sql/multi_modifications.sql index b105fb189..9828c765a 100644 --- a/src/test/regress/sql/multi_modifications.sql +++ b/src/test/regress/sql/multi_modifications.sql @@ -588,9 +588,7 @@ SELECT create_reference_table('reference_summary_table'); INSERT INTO reference_raw_table VALUES (1, 100); INSERT INTO reference_raw_table VALUES (1, 200); INSERT INTO reference_raw_table VALUES (1, 200); -INSERT INTO reference_raw_table VALUES (1, 300); -INSERT INTO reference_raw_table VALUES (2, 400); -INSERT INTO reference_raw_table VALUES (2, 500); +INSERT INTO reference_raw_table VALUES (1,300), (2, 400), (2,500) RETURNING *; INSERT INTO reference_summary_table VALUES (1); INSERT INTO reference_summary_table VALUES (2);