citus/src/backend/distributed/planner/merge_planner.c

520 lines
14 KiB
C

/*-------------------------------------------------------------------------
*
* merge_planner.c
*
* This file contains functions to help plan MERGE queries.
*
* Copyright (c) Citus Data, Inc.
*
*-------------------------------------------------------------------------
*/
#include <stddef.h>
#include "postgres.h"
#include "parser/parsetree.h"
#include "distributed/pg_version_constants.h"
#include "distributed/merge_planner.h"
#include "distributed/multi_logical_optimizer.h"
#include "distributed/multi_router_planner.h"
#include "distributed/listutils.h"
static bool QueryHasMergeCommand(Query *queryTree);
static DeferredErrorMessage * ErrorIfMergeHasUnsupportedTables(Query *parse,
List *rangeTableList,
PlannerRestrictionContext *
restrictionContext);
static DeferredErrorMessage * ErrorIfDistTablesNotColocated(Query *parse,
List *distTablesList,
PlannerRestrictionContext *
plannerRestrictionContext);
static bool IsPartitionColumnInMergeSource(Expr *columnExpression, Query *query, bool
skipOuterVars);
#if PG_VERSION_NUM >= PG_VERSION_15
static DeferredErrorMessage * InsertPartitionColumnMatchesSource(Query *query, RangeTblEntry *resultRte);
#endif
/*
* MergeQuerySupported does check for a MERGE command in the query, if it finds
* one, it will verify the below criteria
* - Supported tables and combinations in ErrorIfMergeHasUnsupportedTables
* - Distributed tables requirements in ErrorIfDistTablesNotColocated
* - Checks target-lists and functions-in-quals in TargetlistAndFunctionsSupported
*/
DeferredErrorMessage *
MergeQuerySupported(Query *originalQuery,
PlannerRestrictionContext *plannerRestrictionContext)
{
/* For non-MERGE commands it's a no-op */
if (!QueryHasMergeCommand(originalQuery))
{
return NULL;
}
List *rangeTableList = ExtractRangeTableEntryList(originalQuery);
RangeTblEntry *resultRte = ExtractResultRelationRTE(originalQuery);
/*
* Fast path queries cannot have merge command, and we prevent the remaining here.
* In Citus we have limited support for MERGE, it's allowed only if all
* the tables(target, source or any CTE) tables are are local i.e. a
* combination of Citus local and Non-Citus tables (regular Postgres tables)
* or distributed tables with some restrictions, please see header of routine
* ErrorIfDistTablesNotColocated for details.
*/
DeferredErrorMessage *deferredError =
ErrorIfMergeHasUnsupportedTables(originalQuery,
rangeTableList,
plannerRestrictionContext);
if (deferredError)
{
return deferredError;
}
Oid resultRelationId = resultRte->relid;
deferredError =
TargetlistAndFunctionsSupported(resultRelationId,
originalQuery->jointree,
originalQuery->jointree->quals,
originalQuery->targetList,
originalQuery->commandType,
originalQuery->returningList);
if (deferredError)
{
return deferredError;
}
#if PG_VERSION_NUM >= PG_VERSION_15
/*
* MERGE is a special case where we have multiple modify statements
* within itself. Check each INSERT/UPDATE/DELETE individually.
*/
MergeAction *action = NULL;
foreach_ptr(action, originalQuery->mergeActionList)
{
Assert(originalQuery->returningList == NULL);
deferredError =
TargetlistAndFunctionsSupported(resultRelationId,
originalQuery->jointree,
action->qual,
action->targetList,
action->commandType,
originalQuery->returningList);
if (deferredError)
{
return deferredError;
}
}
deferredError =
InsertPartitionColumnMatchesSource(originalQuery, resultRte);
if (deferredError)
{
return deferredError;
}
#endif
return NULL;
}
/*
* IsMergeAllowedOnRelation takes a relation entry and checks if MERGE command is
* permitted on special relations, such as materialized view, returns true only if
* it's a "source" relation.
*/
bool
IsMergeAllowedOnRelation(Query *parse, RangeTblEntry *rte)
{
if (!IsMergeQuery(parse))
{
return false;
}
RangeTblEntry *targetRte = rt_fetch(parse->resultRelation, parse->rtable);
/* Is it a target relation? */
if (targetRte->relid == rte->relid)
{
return false;
}
return true;
}
/*
* ErrorIfDistTablesNotColocated Checks to see if
*
* - There are a minimum of two distributed tables (source and a target).
* - All the distributed tables are indeed colocated.
* - MERGE relations are joined on the distribution column
* MERGE .. USING .. ON target.dist_key = source.dist_key
*
* If any of the conditions are not met, it raises an exception.
*/
static DeferredErrorMessage *
ErrorIfDistTablesNotColocated(Query *parse, List *distTablesList,
PlannerRestrictionContext *plannerRestrictionContext)
{
/* All MERGE tables must be distributed */
if (list_length(distTablesList) < 2)
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"For MERGE command, both the source and target "
"must be distributed", NULL, NULL);
}
/* All distributed tables must be colocated */
if (!AllRelationsInListColocated(distTablesList, RANGETABLE_ENTRY))
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"For MERGE command, all the distributed tables "
"must be colocated", NULL, NULL);
}
/* Are source and target tables joined on distribution column? */
if (!RestrictionEquivalenceForPartitionKeys(plannerRestrictionContext))
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE command is only supported when distributed "
"tables are joined on their distribution column",
NULL, NULL);
}
return NULL;
}
/*
* ErrorIfMergeHasUnsupportedTables checks if all the tables(target, source or any CTE
* present) in the MERGE command are local i.e. a combination of Citus local and Non-Citus
* tables (regular Postgres tables), or distributed tables with some restrictions, please
* see header of routine ErrorIfDistTablesNotColocated for details, raises an exception
* for all other combinations.
*/
static DeferredErrorMessage *
ErrorIfMergeHasUnsupportedTables(Query *parse, List *rangeTableList,
PlannerRestrictionContext *restrictionContext)
{
List *distTablesList = NIL;
bool foundLocalTables = false;
RangeTblEntry *rangeTableEntry = NULL;
foreach_ptr(rangeTableEntry, rangeTableList)
{
Oid relationId = rangeTableEntry->relid;
switch (rangeTableEntry->rtekind)
{
case RTE_RELATION:
{
/* Check the relation type */
break;
}
case RTE_SUBQUERY:
case RTE_FUNCTION:
case RTE_TABLEFUNC:
case RTE_VALUES:
case RTE_JOIN:
case RTE_CTE:
{
/* Skip them as base table(s) will be checked */
continue;
}
/*
* RTE_NAMEDTUPLESTORE is typically used in ephmeral named relations,
* such as, trigger data; until we find a genuine use case, raise an
* exception.
* RTE_RESULT is a node added by the planner and we shouldn't
* encounter it in the parse tree.
*/
case RTE_NAMEDTUPLESTORE:
case RTE_RESULT:
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE command is not supported with "
"Tuplestores and results",
NULL, NULL);
}
default:
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE command: Unrecognized range table entry.",
NULL, NULL);
}
}
/* RTE Relation can be of various types, check them now */
/* skip the regular views as they are replaced with subqueries */
if (rangeTableEntry->relkind == RELKIND_VIEW)
{
continue;
}
if (rangeTableEntry->relkind == RELKIND_MATVIEW ||
rangeTableEntry->relkind == RELKIND_FOREIGN_TABLE)
{
/* Materialized view or Foreign table as target is not allowed */
if (IsMergeAllowedOnRelation(parse, rangeTableEntry))
{
/* Non target relation is ok */
continue;
}
else
{
/* Usually we don't reach this exception as the Postgres parser catches it */
StringInfo errorMessage = makeStringInfo();
appendStringInfo(errorMessage,
"MERGE command is not allowed on "
"relation type(relkind:%c)", rangeTableEntry->relkind);
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, errorMessage->data,
NULL, NULL);
}
}
if (rangeTableEntry->relkind != RELKIND_RELATION &&
rangeTableEntry->relkind != RELKIND_PARTITIONED_TABLE)
{
StringInfo errorMessage = makeStringInfo();
appendStringInfo(errorMessage, "Unexpected table type(relkind:%c) "
"in MERGE command", rangeTableEntry->relkind);
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED, errorMessage->data,
NULL, NULL);
}
Assert(rangeTableEntry->relid != 0);
/* Reference tables are not supported yet */
if (IsCitusTableType(relationId, REFERENCE_TABLE))
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE command is not supported on reference "
"tables yet", NULL, NULL);
}
/* Append/Range tables are not supported */
if (IsCitusTableType(relationId, APPEND_DISTRIBUTED) ||
IsCitusTableType(relationId, RANGE_DISTRIBUTED))
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"For MERGE command, all the distributed tables "
"must be colocated, for append/range distribution, "
"colocation is not supported", NULL,
"Consider using hash distribution instead");
}
/*
* For now, save all distributed tables, later (below) we will
* check for supported combination(s).
*/
if (IsCitusTableType(relationId, DISTRIBUTED_TABLE))
{
distTablesList = lappend(distTablesList, rangeTableEntry);
continue;
}
/* Regular Postgres tables and Citus local tables are allowed */
if (!IsCitusTable(relationId) ||
IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
{
foundLocalTables = true;
continue;
}
/* Any other Citus table type missing ? */
}
/* Ensure all tables are indeed local */
if (foundLocalTables && list_length(distTablesList) == 0)
{
/* All the tables are local, supported */
return NULL;
}
else if (foundLocalTables && list_length(distTablesList) > 0)
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE command is not supported with "
"combination of distributed/local tables yet",
NULL, NULL);
}
/* Ensure all distributed tables are indeed co-located */
return ErrorIfDistTablesNotColocated(parse, distTablesList, restrictionContext);
}
/*
* QueryHasMergeCommand walks over the query tree and returns false if there
* is no Merge command (e.g., CMD_MERGE), true otherwise.
*/
static bool
QueryHasMergeCommand(Query *queryTree)
{
/* function is void for pre-15 versions of Postgres */
#if PG_VERSION_NUM < PG_VERSION_15
return false;
#else
/*
* Postgres currently doesn't support Merge queries inside subqueries and
* ctes, but lets be defensive and do query tree walk anyway.
*
* We do not call this path for fast-path queries to avoid this additional
* overhead.
*/
if (!ContainsMergeCommandWalker((Node *) queryTree))
{
/* No MERGE found */
return false;
}
return true;
#endif
}
/*
* IsPartitionColumnInMerge returns true if the given column is a partition column.
* The function uses FindReferencedTableColumn to find the original relation
* id and column that the column expression refers to. It then checks whether
* that column is a partition column of the relation.
*
* Also, the function returns always false for reference tables given that
* reference tables do not have partition column.
*
* If skipOuterVars is true, then it doesn't process the outervars.
*/
bool
IsPartitionColumnInMergeSource(Expr *columnExpression, Query *query, bool skipOuterVars)
{
bool isPartitionColumn = false;
Var *column = NULL;
RangeTblEntry *relationRTE = NULL;
/* ParentQueryList is same as the original query for MERGE */
FindReferencedTableColumn(columnExpression, list_make1(query), query, &column,
&relationRTE,
skipOuterVars);
Oid relationId = relationRTE ? relationRTE->relid : InvalidOid;
if (relationId != InvalidOid && column != NULL)
{
Var *partitionColumn = DistPartitionKey(relationId);
/* not all distributed tables have partition column */
if (partitionColumn != NULL && column->varattno == partitionColumn->varattno)
{
isPartitionColumn = true;
}
}
return isPartitionColumn;
}
#if PG_VERSION_NUM >= PG_VERSION_15
/*
* InsertPartitionColumnMatchesSource check to see if MERGE is inserting a
* value into the target which is not from the source table, if so, it
* raises an exception.
* Note: Inserting random values other than the joined column values will
* result in unexpected behaviour of rows ending up in incorrect shards.
*/
static DeferredErrorMessage *
InsertPartitionColumnMatchesSource(Query *query, RangeTblEntry *resultRte)
{
if (!IsCitusTableType(resultRte->relid, DISTRIBUTED_TABLE))
{
return NULL;
}
bool foundDistributionColumn = false;
MergeAction *action = NULL;
foreach_ptr(action, query->mergeActionList)
{
/* Skip MATCHED clauses */
if (action->matched)
{
continue;
}
/* NOT MATCHED can have either INSERT or DO NOTHING */
if (action->commandType == CMD_NOTHING)
{
return NULL;
}
if (action->targetList == NIL)
{
/* INSERT DEFAULT VALUES is not allowed */
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"cannot perform MERGE INSERT with DEFAULTS",
NULL, NULL);
}
Assert(action->commandType == CMD_INSERT);
Var *targetKey = PartitionColumn(resultRte->relid, 1);
TargetEntry *targetEntry = NULL;
foreach_ptr(targetEntry, action->targetList)
{
if (targetEntry->resjunk)
{
continue;
}
AttrNumber originalAttrNo = targetEntry->resno;
/* skip processing of target table non-partition columns */
if (originalAttrNo != targetKey->varattno)
{
continue;
}
foundDistributionColumn = true;
if (targetEntry->expr->type == T_Var)
{
if (IsPartitionColumnInMergeSource(targetEntry->expr, query, true))
{
return NULL;
}
else
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE INSERT must use the source table "
"distribution column value",
NULL, NULL);
}
}
else
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE INSERT must refer a source column "
"for distribution column ",
NULL, NULL);
}
}
if (!foundDistributionColumn)
{
return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
"MERGE INSERT must have distribution column as value",
NULL, NULL);
}
}
return NULL;
}
#endif