From 490a8302b71cbe78e7028879b0dffa2c43d03f33 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Mon, 25 Dec 2017 18:45:33 +0900
Subject: [PATCH 3/3] Teach executor to handle ON CONFLICT (key) on partitioned
 tables

---
 src/backend/executor/execIndexing.c           | 14 +++++++++---
 src/backend/executor/nodeModifyTable.c        | 32 +++++++++++++++++++++++++++
 src/test/regress/expected/insert_conflict.out | 27 +++++++++++++++++-----
 src/test/regress/sql/insert_conflict.sql      | 18 ++++++++++-----
 4 files changed, 77 insertions(+), 14 deletions(-)

diff --git a/src/backend/executor/execIndexing.c b/src/backend/executor/execIndexing.c
index 89e189fa71..f76a2ede76 100644
--- a/src/backend/executor/execIndexing.c
+++ b/src/backend/executor/execIndexing.c
@@ -531,10 +531,18 @@ ExecCheckIndexConstraints(TupleTableSlot *slot,
 		if (!indexInfo->ii_ReadyForInserts)
 			continue;
 
-		/* When specific arbiter indexes requested, only examine them */
+		/*
+		 * When specific arbiter indexes requested, only examine them.  If
+		 * this is a partition (after a tuple is routed to it from the
+		 * parent into which the original tuple has been inserted), we must
+		 * check the parent index id, instead of our own id, because that's
+		 * the one that appears in the arbiterIndexes list.
+		 */
 		if (arbiterIndexes != NIL &&
-			!list_member_oid(arbiterIndexes,
-							 indexRelation->rd_index->indexrelid))
+			!(list_member_oid(arbiterIndexes,
+							  indexRelation->rd_index->indexrelid) ||
+			  list_member_oid(arbiterIndexes,
+							  indexRelation->rd_index->indparentidx)))
 			continue;
 
 		if (!indexRelation->rd_index->indimmediate)
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index afb83ed3ae..94f819006a 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -2186,6 +2186,38 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 
 			resultRelInfo->ri_onConflictSetWhere = qualexpr;
 		}
+
+		/* Build the above information for each leaf partition rel */
+		for (i = 0; i < mtstate->mt_num_partitions; i++)
+		{
+			Relation	partrel;
+			List	   *leaf_oc_set;
+
+			resultRelInfo = mtstate->mt_partitions[i];
+			partrel = resultRelInfo->ri_RelationDesc;
+
+			/* varno = node->nominalRelation */
+			leaf_oc_set = map_partition_varattnos(node->onConflictSet,
+											node->nominalRelation,
+											partrel, rel, NULL);
+			resultRelInfo->ri_onConflictSetProj =
+					ExecBuildProjectionInfo(leaf_oc_set, econtext,
+									mtstate->mt_conflproj, &mtstate->ps,
+									resultRelInfo->ri_RelationDesc->rd_att);
+
+			if (node->onConflictWhere)
+			{
+				List   *leaf_oc_where;
+
+				/* varno = node->nominalRelation */
+				leaf_oc_where =
+					map_partition_varattnos((List *) node->onConflictWhere,
+											node->nominalRelation,
+											partrel, rel, NULL);
+				resultRelInfo->ri_onConflictSetWhere =
+								ExecInitQual(leaf_oc_where, &mtstate->ps);
+			}
+		}
 	}
 
 	/*
diff --git a/src/test/regress/expected/insert_conflict.out b/src/test/regress/expected/insert_conflict.out
index 8fd2027d6a..dcb07fc09e 100644
--- a/src/test/regress/expected/insert_conflict.out
+++ b/src/test/regress/expected/insert_conflict.out
@@ -792,10 +792,25 @@ create table parted_conflict_test (a int, b char) partition by list (a);
 create table parted_conflict_test_1 partition of parted_conflict_test (b unique) for values in (1);
 insert into parted_conflict_test values (1, 'a') on conflict do nothing;
 insert into parted_conflict_test values (1, 'a') on conflict do nothing;
--- however, on conflict do update is not supported yet
-insert into parted_conflict_test values (1) on conflict (b) do update set a = excluded.a;
-ERROR:  there is no unique or exclusion constraint matching the ON CONFLICT specification
--- but it works OK if we target the partition directly
-insert into parted_conflict_test_1 values (1) on conflict (b) do
-update set a = excluded.a;
+-- create one more partition and a partitioned unique index
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (2);
+create unique index on parted_conflict_test (a);
+insert into parted_conflict_test values (1, 'a') on conflict (a) do nothing;
+insert into parted_conflict_test values (1, 'b') on conflict (a) do update set b = excluded.b;
+insert into parted_conflict_test values (1, 'b') on conflict (a) do update set b = excluded.b where parted_conflict_test.b = 'a';
+select * from parted_conflict_test;
+ a | b 
+---+---
+ 1 | b
+(1 row)
+
+-- also works OK if we target the partition directly
+insert into parted_conflict_test_1 values (1, 'c') on conflict (a) do
+update set b = excluded.b;
+select * from parted_conflict_test;
+ a | b 
+---+---
+ 1 | c
+(1 row)
+
 drop table parted_conflict_test;
diff --git a/src/test/regress/sql/insert_conflict.sql b/src/test/regress/sql/insert_conflict.sql
index 32c647e3f8..264f67ce89 100644
--- a/src/test/regress/sql/insert_conflict.sql
+++ b/src/test/regress/sql/insert_conflict.sql
@@ -478,9 +478,17 @@ create table parted_conflict_test (a int, b char) partition by list (a);
 create table parted_conflict_test_1 partition of parted_conflict_test (b unique) for values in (1);
 insert into parted_conflict_test values (1, 'a') on conflict do nothing;
 insert into parted_conflict_test values (1, 'a') on conflict do nothing;
--- however, on conflict do update is not supported yet
-insert into parted_conflict_test values (1) on conflict (b) do update set a = excluded.a;
--- but it works OK if we target the partition directly
-insert into parted_conflict_test_1 values (1) on conflict (b) do
-update set a = excluded.a;
+
+-- create one more partition and a partitioned unique index
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (2);
+create unique index on parted_conflict_test (a);
+insert into parted_conflict_test values (1, 'a') on conflict (a) do nothing;
+insert into parted_conflict_test values (1, 'b') on conflict (a) do update set b = excluded.b;
+insert into parted_conflict_test values (1, 'b') on conflict (a) do update set b = excluded.b where parted_conflict_test.b = 'a';
+select * from parted_conflict_test;
+
+-- also works OK if we target the partition directly
+insert into parted_conflict_test_1 values (1, 'c') on conflict (a) do
+update set b = excluded.b;
+select * from parted_conflict_test;
 drop table parted_conflict_test;
-- 
2.11.0

