From 6a15e1c9fcd94698216c691330f420c9973c0f3c Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Thu, 19 Jul 2018 15:51:40 +0200
Subject: [PATCH] extract ErrorIfOnConflictNotSupported function for reuse

---
 .../planner/multi_router_planner.c            | 52 +++++++++++++++----
 .../distributed/multi_router_planner.h        |  1 +
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c
index f39527052..277b5deed 100644
--- a/src/backend/distributed/planner/multi_router_planner.c
+++ b/src/backend/distributed/planner/multi_router_planner.c
@@ -530,12 +530,8 @@ ModifyQuerySupported(Query *queryTree, Query *originalQuery, bool multiShardQuer
 	List *rangeTableList = NIL;
 	ListCell *rangeTableCell = NULL;
 	uint32 queryTableCount = 0;
-	bool specifiesPartitionValue = false;
-	ListCell *setTargetCell = NULL;
-	List *onConflictSet = NIL;
-	Node *arbiterWhere = NULL;
-	Node *onConflictWhere = NULL;
 	CmdType commandType = queryTree->commandType;
+	DeferredErrorMessage *deferredError = NULL;
 
 	/*
 	 * Here, we check if a recursively planned query tries to modify
@@ -769,7 +765,10 @@ ModifyQuerySupported(Query *queryTree, Query *originalQuery, bool multiShardQuer
 				TargetEntryChangesValue(targetEntry, partitionColumn,
 										queryTree->jointree))
 			{
-				specifiesPartitionValue = true;
+				return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
+									 "modifying the partition value of rows is not "
+									 "allowed",
+									 NULL, NULL);
 			}
 
 			if (commandType == CMD_UPDATE &&
@@ -829,13 +828,46 @@ ModifyQuerySupported(Query *queryTree, Query *originalQuery, bool multiShardQuer
 		}
 	}
 
-	if (commandType == CMD_INSERT && queryTree->onConflict != NULL)
+	deferredError = ErrorIfOnConflictNotSupported(queryTree);
+	if (deferredError != NULL)
 	{
-		onConflictSet = queryTree->onConflict->onConflictSet;
-		arbiterWhere = queryTree->onConflict->arbiterWhere;
-		onConflictWhere = queryTree->onConflict->onConflictWhere;
+		return deferredError;
 	}
 
+	return NULL;
+}
+
+
+/*
+ * ErrorIfOnConflictNotSupprted returns an error if an INSERT query has an
+ * unsupported ON CONFLICT clause. In particular, changing the partition
+ * column value or using volatile functions is not allowed.
+ */
+DeferredErrorMessage *
+ErrorIfOnConflictNotSupported(Query *queryTree)
+{
+	Oid distributedTableId = InvalidOid;
+	uint32 rangeTableId = 1;
+	Var *partitionColumn = NULL;
+	List *onConflictSet = NIL;
+	Node *arbiterWhere = NULL;
+	Node *onConflictWhere = NULL;
+	ListCell *setTargetCell = NULL;
+	bool specifiesPartitionValue = false;
+
+	CmdType commandType = queryTree->commandType;
+	if (commandType != CMD_INSERT || queryTree->onConflict == NULL)
+	{
+		return NULL;
+	}
+
+	distributedTableId = ExtractFirstDistributedTableId(queryTree);
+	partitionColumn = PartitionColumn(distributedTableId, rangeTableId);
+
+	onConflictSet = queryTree->onConflict->onConflictSet;
+	arbiterWhere = queryTree->onConflict->arbiterWhere;
+	onConflictWhere = queryTree->onConflict->onConflictWhere;
+
 	/*
 	 * onConflictSet is expanded via expand_targetlist() on the standard planner.
 	 * This ends up adding all the columns to the onConflictSet even if the user
diff --git a/src/include/distributed/multi_router_planner.h b/src/include/distributed/multi_router_planner.h
index 9206d1337..ff22d1335 100644
--- a/src/include/distributed/multi_router_planner.h
+++ b/src/include/distributed/multi_router_planner.h
@@ -52,6 +52,7 @@ extern DeferredErrorMessage * ModifyQuerySupported(Query *queryTree, Query *orig
 												   bool multiShardQuery,
 												   PlannerRestrictionContext *
 												   plannerRestrictionContext);
+extern DeferredErrorMessage * ErrorIfOnConflictNotSupported(Query *queryTree);
 extern List * ShardIntervalOpExpressions(ShardInterval *shardInterval, Index rteIndex);
 extern RelationRestrictionContext * CopyRelationRestrictionContext(
 	RelationRestrictionContext *oldContext);