From a2517bd315034de7f6c5a4728f66729136918e88 Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Tue, 27 Feb 2018 20:52:56 -0300
Subject: [PATCH v1] fix ON CONFLICT DO UPDATE for partitioned tables

---
 src/backend/catalog/pg_inherits.c             |  72 ++++++++++++++++
 src/backend/executor/execPartition.c          |  29 +++++++
 src/backend/executor/nodeModifyTable.c        |  33 +++++++-
 src/backend/optimizer/util/plancat.c          |   3 +-
 src/backend/parser/analyze.c                  |   7 --
 src/include/catalog/pg_inherits_fn.h          |   3 +
 src/test/regress/expected/insert_conflict.out | 113 ++++++++++++++++++++++++--
 src/test/regress/sql/insert_conflict.sql      |  75 +++++++++++++++--
 8 files changed, 311 insertions(+), 24 deletions(-)

diff --git a/src/backend/catalog/pg_inherits.c b/src/backend/catalog/pg_inherits.c
index 5a5beb9273..e1a46bcd2b 100644
--- a/src/backend/catalog/pg_inherits.c
+++ b/src/backend/catalog/pg_inherits.c
@@ -407,6 +407,78 @@ typeInheritsFrom(Oid subclassTypeId, Oid superclassTypeId)
 }
 
 /*
+ * Given a list of index OIDs in the rootrel, return a list of OIDs of the
+ * corresponding indexes in the partrel.  If any index in the rootrel does not
+ * correspond to any index in the child, an error is raised.
+ *
+ * This processes the index list for INSERT ON CONFLICT DO UPDATE at execution
+ * time.  This fact is hardcoded in the error messages.
+ *
+ * XXX this implementation fails if the partition is not a direct child of
+ * rootrel.
+ */
+List *
+MapPartitionIndexList(Relation rootrel, Relation partrel, List *indexlist)
+{
+	List	   *result = NIL;
+	List	   *partIdxs;
+	Relation	inhRel;
+	ScanKeyData	key;
+	ListCell   *cell;
+
+	partIdxs = RelationGetIndexList(partrel);
+	/* quick exit if partition has no indexes */
+	if (partIdxs == NIL)
+		return NIL;
+
+	inhRel = heap_open(InheritsRelationId, AccessShareLock);
+
+	foreach(cell, indexlist)
+	{
+		Oid			parentIdx = lfirst_oid(cell);
+		SysScanDesc	scan;
+		HeapTuple	tuple;
+		bool		found = false;
+
+		ScanKeyInit(&key,
+					Anum_pg_inherits_inhparent,
+					BTEqualStrategyNumber, F_OIDEQ,
+					ObjectIdGetDatum(parentIdx));
+
+		scan = systable_beginscan(inhRel, InheritsParentIndexId, true,
+								  NULL, 1, &key);
+		while (HeapTupleIsValid(tuple = systable_getnext(scan)))
+		{
+			Oid indexoid = ((Form_pg_inherits) GETSTRUCT(tuple))->inhrelid;
+
+			if (list_member_oid(partIdxs, indexoid))
+			{
+				result = lappend_oid(result, indexoid);
+				found = true;
+				break;
+			}
+		}
+		systable_endscan(scan);
+
+		/*
+		 * Indexes can only be used as inference targets if they exist in the
+		 * partition that receives the tuple; bail out if we cannot find it.
+		 */
+		if (!found)
+			ereport(ERROR,
+					(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+					 errmsg("invalid ON CONFLICT DO UPDATE specification"),
+					 errdetail("An inferred index was not found in partition \"%s\".",
+							   RelationGetRelationName(partrel))));
+	}
+
+	relation_close(inhRel, AccessShareLock);
+	list_free(partIdxs);
+
+	return result;
+}
+
+/*
  * Create a single pg_inherits row with the given data
  */
 void
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 54efc9e545..95a814e975 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -475,6 +475,35 @@ ExecInitPartitionInfo(ModifyTableState *mtstate,
 									&mtstate->ps, RelationGetDescr(partrel));
 	}
 
+	/*
+	 * If needed, initialize projection and qual for ON CONFLICT DO UPDATE for
+	 * this partition.
+	 */
+	if (node && node->onConflictAction == ONCONFLICT_UPDATE)
+	{
+		ExprContext *econtext = mtstate->ps.ps_ExprContext;
+		List	   *leaf_oc_set;
+
+		leaf_oc_set = map_partition_varattnos(node->onConflictSet,
+											  node->nominalRelation,
+											  partrel, rootrel, NULL);
+		leaf_part_rri->ri_onConflictSetProj =
+			ExecBuildProjectionInfo(leaf_oc_set, econtext,
+									mtstate->mt_conflproj, &mtstate->ps,
+									RelationGetDescr(partrel));
+		if (node->onConflictWhere)
+		{
+			List	   *leaf_oc_where;
+
+			leaf_oc_where =
+				map_partition_varattnos((List *) node->onConflictWhere,
+										node->nominalRelation,
+										partrel, rootrel, NULL);
+			leaf_part_rri->ri_onConflictSetWhere =
+				ExecInitQual(leaf_oc_where, &mtstate->ps);
+		}
+	}
+
 	Assert(proute->partitions[partidx] == NULL);
 	proute->partitions[partidx] = leaf_part_rri;
 
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index c32928d9bd..9748f80ddc 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -39,6 +39,7 @@
 
 #include "access/htup_details.h"
 #include "access/xact.h"
+#include "catalog/pg_inherits_fn.h"
 #include "commands/trigger.h"
 #include "executor/execPartition.h"
 #include "executor/executor.h"
@@ -510,6 +511,20 @@ ExecInsert(ModifyTableState *mtstate,
 			uint32		specToken;
 			ItemPointerData conflictTid;
 			bool		specConflict;
+			List	   *mappedArbiterIndexes;
+
+			/*
+			 * Map the arbiter index list to the OIDs in the corresponding
+			 * partition.
+			 */
+			if (saved_resultRelInfo &&
+				resultRelInfo->ri_RelationDesc->rd_rel->relispartition)
+				mappedArbiterIndexes =
+					MapPartitionIndexList(saved_resultRelInfo->ri_RelationDesc,
+										  resultRelInfo->ri_RelationDesc,
+										  arbiterIndexes);
+			else
+				mappedArbiterIndexes = arbiterIndexes;
 
 			/*
 			 * Do a non-conclusive check for conflicts first.
@@ -526,7 +541,7 @@ ExecInsert(ModifyTableState *mtstate,
 	vlock:
 			specConflict = false;
 			if (!ExecCheckIndexConstraints(slot, estate, &conflictTid,
-										   arbiterIndexes))
+										   mappedArbiterIndexes))
 			{
 				/* committed conflict tuple found */
 				if (onconflict == ONCONFLICT_UPDATE)
@@ -581,7 +596,7 @@ ExecInsert(ModifyTableState *mtstate,
 			/* insert index entries for tuple */
 			recheckIndexes = ExecInsertIndexTuples(slot, &(tuple->t_self),
 												   estate, true, &specConflict,
-												   arbiterIndexes);
+												   mappedArbiterIndexes);
 
 			/* adjust the tuple's state accordingly */
 			if (!specConflict)
@@ -1146,6 +1161,18 @@ lreplace:;
 			TupleConversionMap *tupconv_map;
 
 			/*
+			 * Disallow an INSERT ON CONFLICT DO UPDATE that causes the
+			 * original row to migrate to a different partition.  Maybe this
+			 * can be implemented some day, but it seems a fringe feature with
+			 * little redeeming value.
+			 */
+			if (((ModifyTable *) mtstate->ps.plan)->onConflictAction == ONCONFLICT_UPDATE)
+				ereport(ERROR,
+						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						 errmsg("invalid ON UPDATE specification"),
+						 errdetail("The result tuple would appear in a different partition than the original tuple.")));
+
+			/*
 			 * When an UPDATE is run on a leaf partition, we will not have
 			 * partition tuple routing set up. In that case, fail with
 			 * partition constraint violation error.
@@ -2329,7 +2356,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	}
 
 	/*
-	 * If needed, Initialize target list, projection and qual for ON CONFLICT
+	 * If needed, initialize target list, projection and qual for ON CONFLICT
 	 * DO UPDATE.
 	 */
 	resultRelInfo = mtstate->resultRelInfo;
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
index 60f21711f4..db7c0030ca 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -558,7 +558,8 @@ get_relation_foreign_keys(PlannerInfo *root, RelOptInfo *rel,
 
 /*
  * infer_arbiter_indexes -
- *	  Determine the unique indexes used to arbitrate speculative insertion.
+ *	  Determine the unique indexes used to arbitrate speculative insertion,
+ *	  and return them as a list of OIDs.
  *
  * Uses user-supplied inference clause expressions and predicate to match a
  * unique index from those defined and ready on the heap relation (target).
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index c3a9617f67..92696f0607 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -1025,13 +1025,6 @@ transformOnConflictClause(ParseState *pstate,
 		TargetEntry *te;
 		int			attno;
 
-		if (targetrel->rd_partdesc)
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("%s cannot be applied to partitioned table \"%s\"",
-							"ON CONFLICT DO UPDATE",
-							RelationGetRelationName(targetrel))));
-
 		/*
 		 * All INSERT expressions have been parsed, get ready for potentially
 		 * existing SET statements that need to be processed like an UPDATE.
diff --git a/src/include/catalog/pg_inherits_fn.h b/src/include/catalog/pg_inherits_fn.h
index eebee977a5..20fb96db51 100644
--- a/src/include/catalog/pg_inherits_fn.h
+++ b/src/include/catalog/pg_inherits_fn.h
@@ -16,6 +16,7 @@
 
 #include "nodes/pg_list.h"
 #include "storage/lock.h"
+#include "utils/relcache.h"
 
 extern List *find_inheritance_children(Oid parentrelId, LOCKMODE lockmode);
 extern List *find_all_inheritors(Oid parentrelId, LOCKMODE lockmode,
@@ -23,6 +24,8 @@ extern List *find_all_inheritors(Oid parentrelId, LOCKMODE lockmode,
 extern bool has_subclass(Oid relationId);
 extern bool has_superclass(Oid relationId);
 extern bool typeInheritsFrom(Oid subclassTypeId, Oid superclassTypeId);
+extern List *MapPartitionIndexList(Relation rootrel, Relation partrel,
+					  List *indexlist);
 extern void StoreSingleInheritance(Oid relationId, Oid parentOid,
 					   int32 seqNumber);
 extern bool DeleteInheritsTuple(Oid inhrelid, Oid inhparent);
diff --git a/src/test/regress/expected/insert_conflict.out b/src/test/regress/expected/insert_conflict.out
index 2650faedee..da8fe11120 100644
--- a/src/test/regress/expected/insert_conflict.out
+++ b/src/test/regress/expected/insert_conflict.out
@@ -786,16 +786,115 @@ select * from selfconflict;
 (3 rows)
 
 drop table selfconflict;
--- check that the following works:
+--
+-- INSERT ON CONFLICT and partitioned tables
+--
+-- DO NOTHING works
 -- insert into partitioned_table on conflict do nothing
 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:  ON CONFLICT DO UPDATE cannot be applied to partitioned table "parted_conflict_test"
--- 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;
+drop table parted_conflict_test;
+-- simple DO UPDATE works, as long as the tuple remains in the same partition
+create table parted_conflict_test (a int primary key, b text) partition by list (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values in (1, 2);
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (3, 4);
+insert into parted_conflict_test values (1, 'first');
+insert into parted_conflict_test values (1, 'second')
+  on conflict (a) do nothing;
+insert into parted_conflict_test values (1, 'third')
+  on conflict (a) do update set b = format('%s (was %s)', excluded.b, parted_conflict_test.b);
+select * from parted_conflict_test;
+ a |         b         
+---+-------------------
+ 1 | third (was first)
+(1 row)
+
+insert into parted_conflict_test values (1, 'b')
+  on conflict (a) do update set b = 'fourth'
+  where parted_conflict_test.b = 'third (was first)';
+select * from parted_conflict_test;
+ a |   b    
+---+--------
+ 1 | fourth
+(1 row)
+
+insert into parted_conflict_test values (1, 'c')
+  on conflict (a) do update set b = 'fourth'
+  where parted_conflict_test.b = 'b';
+select * from parted_conflict_test;
+ a |   b    
+---+--------
+ 1 | fourth
+(1 row)
+
+insert into parted_conflict_test values (1, 'fifth')
+  on conflict (a) do update set a = parted_conflict_test.a * 2,
+  b = format('%s (was %s)', excluded.b, parted_conflict_test.b);
+select * from parted_conflict_test;
+ a |         b          
+---+--------------------
+ 2 | fifth (was fourth)
+(1 row)
+
+-- targetting the partition directly also works
+insert into parted_conflict_test_1 values (2, 'sixth') on conflict (a) do
+  update set b = format('%s (was %s)', excluded.b, parted_conflict_test_1.b);
+select * from parted_conflict_test;
+ a |               b                
+---+--------------------------------
+ 2 | sixth (was fifth (was fourth))
+(1 row)
+
+drop table parted_conflict_test;
+-- moving tuple to another partition in the UPDATE clause is not supported
+create table parted_conflict_test (a int, b text) partition by list (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values in (1);
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (2);
+insert into parted_conflict_test values (1, 'one');
+insert into parted_conflict_test values (1, 'one two')
+  on conflict (a) do update set a = excluded.a * 2;
+ERROR:  there is no unique or exclusion constraint matching the ON CONFLICT specification
+drop table parted_conflict_test;
+-- multiple-layered partitioning
+create table parted_conflict_test (a int primary key, b text) partition by range (a);
+create table parted_conflict_test_1 partition of parted_conflict_test
+  for values from (0) to (10000) partition by range (a);
+create table parted_conflict_test_1_1 partition of parted_conflict_test_1
+  for values from (0) to (100);
+insert into parted_conflict_test values ('10', 'ten');
+insert into parted_conflict_test values ('10', 'ten two')
+  on conflict (a) do update set b = excluded.b;
+ERROR:  invalid ON CONFLICT DO UPDATE specification
+DETAIL:  An inferred index was not found in partition "parted_conflict_test_1_1".
+select * from parted_conflict_test;
+ a  |  b  
+----+-----
+ 10 | ten
+(1 row)
+
+insert into parted_conflict_test_1 values ('10', 'ten three')
+  on conflict (a) do update set b = excluded.b;
+select * from parted_conflict_test;
+ a  |     b     
+----+-----------
+ 10 | ten three
+(1 row)
+
+drop table parted_conflict_test;
+-- a partitioned table with an index and no corresponding index on the
+-- partition; should raise an error
+create table parted_conflict_test (a int, b text) partition by range (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values from (0) to (10000);
+alter table only parted_conflict_test add primary key (a);
+insert into parted_conflict_test values (100, 'hundred');
+insert into parted_conflict_test values (100, 'hundred (two)') on conflict (a) do update set b = excluded.b;
+ERROR:  there is no unique or exclusion constraint matching the ON CONFLICT specification
+select * from parted_conflict_test;
+  a  |    b    
+-----+---------
+ 100 | hundred
+(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..61758d2ea9 100644
--- a/src/test/regress/sql/insert_conflict.sql
+++ b/src/test/regress/sql/insert_conflict.sql
@@ -472,15 +472,78 @@ select * from selfconflict;
 
 drop table selfconflict;
 
--- check that the following works:
+--
+-- INSERT ON CONFLICT and partitioned tables
+--
+
+-- DO NOTHING works
 -- insert into partitioned_table on conflict do nothing
 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;
+drop table parted_conflict_test;
+
+-- simple DO UPDATE works, as long as the tuple remains in the same partition
+create table parted_conflict_test (a int primary key, b text) partition by list (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values in (1, 2);
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (3, 4);
+insert into parted_conflict_test values (1, 'first');
+insert into parted_conflict_test values (1, 'second')
+  on conflict (a) do nothing;
+insert into parted_conflict_test values (1, 'third')
+  on conflict (a) do update set b = format('%s (was %s)', excluded.b, parted_conflict_test.b);
+select * from parted_conflict_test;
+insert into parted_conflict_test values (1, 'b')
+  on conflict (a) do update set b = 'fourth'
+  where parted_conflict_test.b = 'third (was first)';
+select * from parted_conflict_test;
+insert into parted_conflict_test values (1, 'c')
+  on conflict (a) do update set b = 'fourth'
+  where parted_conflict_test.b = 'b';
+select * from parted_conflict_test;
+insert into parted_conflict_test values (1, 'fifth')
+  on conflict (a) do update set a = parted_conflict_test.a * 2,
+  b = format('%s (was %s)', excluded.b, parted_conflict_test.b);
+select * from parted_conflict_test;
+
+-- targetting the partition directly also works
+insert into parted_conflict_test_1 values (2, 'sixth') on conflict (a) do
+  update set b = format('%s (was %s)', excluded.b, parted_conflict_test_1.b);
+select * from parted_conflict_test;
+drop table parted_conflict_test;
+
+-- moving tuple to another partition in the UPDATE clause is not supported
+create table parted_conflict_test (a int, b text) partition by list (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values in (1);
+create table parted_conflict_test_2 partition of parted_conflict_test for values in (2);
+insert into parted_conflict_test values (1, 'one');
+insert into parted_conflict_test values (1, 'one two')
+  on conflict (a) do update set a = excluded.a * 2;
+drop table parted_conflict_test;
+
+-- multiple-layered partitioning
+create table parted_conflict_test (a int primary key, b text) partition by range (a);
+create table parted_conflict_test_1 partition of parted_conflict_test
+  for values from (0) to (10000) partition by range (a);
+create table parted_conflict_test_1_1 partition of parted_conflict_test_1
+  for values from (0) to (100);
+insert into parted_conflict_test values ('10', 'ten');
+insert into parted_conflict_test values ('10', 'ten two')
+  on conflict (a) do update set b = excluded.b;
+select * from parted_conflict_test;
+
+insert into parted_conflict_test_1 values ('10', 'ten three')
+  on conflict (a) do update set b = excluded.b;
+select * from parted_conflict_test;
+drop table parted_conflict_test;
+
+-- a partitioned table with an index and no corresponding index on the
+-- partition; should raise an error
+create table parted_conflict_test (a int, b text) partition by range (a);
+create table parted_conflict_test_1 partition of parted_conflict_test for values from (0) to (10000);
+alter table only parted_conflict_test add primary key (a);
+insert into parted_conflict_test values (100, 'hundred');
+insert into parted_conflict_test values (100, 'hundred (two)') on conflict (a) do update set b = excluded.b;
+select * from parted_conflict_test;
 drop table parted_conflict_test;
-- 
2.11.0

