Fix upsert TLE translation when mapping variable numbers

Previously, there was a bug when chunks contained dropped columns
in the mapping of variable numbers. This PR fixes the bug and adds a
test for this case.

src/chunk_insert_state.c

@@ -20,6 +20,8 @@
 #include <miscadmin.h>
 #include <parser/parsetree.h>
 #include <rewrite/rewriteManip.h>
+#include <nodes/makefuncs.h>
+#include <catalog/pg_type.h>
 
 #include "errors.h"
 #include "chunk_insert_state.h"
@@ -252,67 +254,92 @@ get_adjusted_onconflictupdate_where(ExprState *hyper_where_state, List *where_qu
 #endif
 }
 
+/*
+ * adjust_hypertable_tlist - from Postgres source code `adjust_partition_tlist`
+ *		Adjust the targetlist entries for a given chunk to account for
+ *		attribute differences between hypertable and the chunk
+ *
+ * The expressions have already been fixed, but here we fix the list to make
+ * target resnos match the chunk's attribute numbers.  This results in a
+ * copy of the original target list in which the entries appear in resno
+ * order, including both the existing entries (that may have their resno
+ * changed in-place) and the newly added entries for columns that don't exist
+ * in the parent.
+ *
+ * Scribbles on the input tlist's entries resno so be aware.
+ */
+static List *
+adjust_hypertable_tlist(List *tlist, TupleConversionMap *map)
+{
+	List	   *new_tlist = NIL;
+	TupleDesc	chunk_tupdesc = map->outdesc;
+	AttrNumber *attrMap = map->attrMap;
+	AttrNumber	chunk_attrno;
+
+	for (chunk_attrno = 1; chunk_attrno <= chunk_tupdesc->natts; chunk_attrno++)
+	{
+		Form_pg_attribute att_tup = TupleDescAttr(chunk_tupdesc, chunk_attrno - 1);
+		TargetEntry *tle;
+
+		if (attrMap[chunk_attrno - 1] != InvalidAttrNumber)
+		{
+			Assert(!att_tup->attisdropped);
+
+			/*
+			 * Use the corresponding entry from the parent's tlist, adjusting
+			 * the resno the match the partition's attno.
+			 */
+			tle = (TargetEntry *) list_nth(tlist, attrMap[chunk_attrno - 1] - 1);
+			if (namestrcmp(&att_tup->attname, tle->resname) != 0)
+				elog(ERROR, "invalid translation of ON CONFLICT update statements");
+			tle->resno = chunk_attrno;
+		}
+		else
+		{
+			Const	   *expr;
+
+			/*
+			 * For a dropped attribute in the partition, generate a dummy
+			 * entry with resno matching the partition's attno.
+			 */
+			Assert(att_tup->attisdropped);
+			expr = makeConst(INT4OID,
+							 -1,
+							 InvalidOid,
+							 sizeof(int32),
+							 (Datum) 0,
+							 true,	/* isnull */
+							 true /* byval */ );
+			tle = makeTargetEntry((Expr *) expr,
+								  chunk_attrno,
+								  pstrdup(NameStr(att_tup->attname)),
+								  false);
+		}
+		new_tlist = lappend(new_tlist, tle);
+	}
+	return new_tlist;
+}
+
 static ProjectionInfo *
-get_adjusted_projection_info_onconflicupdate(ProjectionInfo *orig, List *update_tles, AttrNumber *map,
+get_adjusted_projection_info_onconflicupdate(ProjectionInfo *orig, List *update_tles, AttrNumber *variable_attnos_map,
-											 int map_size, Index varno, Oid rowtype,
+											 int variable_attnos_map_size, Index varno, Oid rowtype,
-											 TupleDesc chunk_desc)
+											 TupleDesc chunk_desc, TupleConversionMap *hypertable_to_chunk_map)
 {
 	bool		found_whole_row;
-	int			i;
-	ListCell   *lc,
-			   *prev,
-			   *next;
 
 	Assert(update_tles != NIL);
+	/* copy the tles so as not to scribble on the hypertable tles */
+	update_tles = copyObject(update_tles);
+
 	/* map hypertable attnos -> chunk attnos for the hypertable */
-	update_tles = (List *) map_variable_attnos_compat((Node *) update_tles, varno, 0, map, map_size,
+	update_tles = (List *) map_variable_attnos_compat((Node *) update_tles, varno, 0, variable_attnos_map, variable_attnos_map_size,
 													  rowtype, &found_whole_row);
 	/* map hypertable attnos -> chunk attnos for the "excluded" table */
-	update_tles = (List *) map_variable_attnos_compat((Node *) update_tles, INNER_VAR, 0, map, map_size,
+	update_tles = (List *) map_variable_attnos_compat((Node *) update_tles, INNER_VAR, 0, variable_attnos_map, variable_attnos_map_size,
 													  rowtype, &found_whole_row);
 
 
-	Assert(list_length(update_tles) == map_size);
+	update_tles = adjust_hypertable_tlist(update_tles, hypertable_to_chunk_map);
-
-	/*
-	 * delete any hypertable tle entries for which there are no chunk
-	 * attributes. This can happen for dropped columns that exist on the
-	 * hyperable but not on the chunks.
-	 */
-	lc = list_head(update_tles);
-	prev = NULL;
-	for (i = 0; i < map_size; i++)
-	{
-		AttrNumber	att = map[i];
-
-		next = lnext(lc);
-
-		if (att == InvalidAttrNumber)
-			update_tles = list_delete_cell(update_tles, lc, prev);
-		else
-			prev = lc;
-		lc = next;
-	}
-
-	/*
-	 * Double-check that the TLE is in the same order as the chunk attributes
-	 * and also fix up the TLE resno after the deletes above
-	 */
-	i = 0;
-	foreach(lc, update_tles)
-	{
-		TargetEntry *t = lfirst(lc);
-		Form_pg_attribute attribute;
-
-		Assert(i < chunk_desc->natts);
-		attribute = chunk_desc->attrs[i];
-
-		if (namestrcmp(&attribute->attname, t->resname) != 0)
-			elog(ERROR, "invalid translation of ON CONFLICT update statements");
-
-		t->resno = i + 1;
-		i++;
-	}
 
 	/* set the projection result slot to the chunk schema */
 	ExecSetSlotDescriptor(get_projection_info_slot_compat(orig), chunk_desc);
@@ -363,7 +390,7 @@ adjust_projections(ChunkInsertState *cis, ChunkDispatch *dispatch, Oid rowtype)
 														 variable_attnos_map_size,
 														 dispatch->hypertable_result_rel_info->ri_RangeTableIndex,
 														 rowtype,
-														 chunk_desc);
+														 chunk_desc, cis->tup_conv_map);
 
 		if (rri->ri_onConflictSetWhere != NULL)
 		{

test/expected/upsert.out

@@ -218,16 +218,39 @@ DO UPDATE SET color = excluded.color WHERE upsert_test_space.temp < 20 RETURNING
  Fri Jan 20 09:00:01 2017 |  3.5 | orange8 | dev5                 | 
 (1 row)
 
-SELECT * FROM upsert_test_space; 
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_1) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-1-new') ON CONFLICT (time, device_id)
-           time           | temp |             color              |      device_id       | device_id_1 
+DO UPDATE SET device_id_1 = excluded.device_id_1 RETURNING *;
---------------------------+------+--------------------------------+----------------------+-------------
+           time           | temp |  color  |      device_id       |     device_id_1      
+--------------------------+------+---------+----------------------+----------------------
+ Fri Jan 20 09:00:01 2017 |  3.5 | orange8 | dev5                 | device-id-1-new     
+(1 row)
+
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_1) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-1-new') ON CONFLICT (time, device_id)
+DO UPDATE SET device_id_1 = 'device-id-1-new-2', color = 'orange9'  RETURNING *;
+           time           | temp |  color  |      device_id       |     device_id_1      
+--------------------------+------+---------+----------------------+----------------------
+ Fri Jan 20 09:00:01 2017 |  3.5 | orange9 | dev5                 | device-id-1-new-2   
+(1 row)
+
+SELECT * FROM upsert_test_space;
+           time           | temp |             color              |      device_id       |     device_id_1      
+--------------------------+------+--------------------------------+----------------------+----------------------
  Fri Jan 20 09:00:01 2017 | 25.9 | orange                         | dev1                 | 
  Fri Jan 20 09:00:01 2017 | 25.9 | orange3 (originally yellow)    | dev2                 | 
  Fri Jan 20 09:00:01 2017 | 23.5 | orange3.1                      | dev3                 | 
  Fri Jan 20 09:00:01 2017 | 23.5 | orange5.1 (originally orange5) | dev4                 | 
- Fri Jan 20 09:00:01 2017 |  3.5 | orange8                        | dev5                 | 
+ Fri Jan 20 09:00:01 2017 |  3.5 | orange9                        | dev5                 | device-id-1-new-2   
 (5 rows)
 
+ALTER TABLE upsert_test_space DROP device_id_1, ADD device_id_2 char(20);
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_2) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-2')
+ON CONFLICT (time, device_id)
+DO UPDATE SET device_id_2 = 'device-id-2-new', color = 'orange10' RETURNING *;
+           time           | temp |  color   |      device_id       |     device_id_2      
+--------------------------+------+----------+----------------------+----------------------
+ Fri Jan 20 09:00:01 2017 |  3.5 | orange10 | dev5                 | device-id-2-new     
+(1 row)
+
 WITH CTE AS (
     INSERT INTO upsert_test_multi_unique
     VALUES ('2017-01-20T09:00:01', 25.9, 'purple')

test/sql/upsert.sql

@@ -94,9 +94,18 @@ INSERT INTO upsert_test_space (time, device_id, temp, color) VALUES ('2017-01-20
 DO UPDATE SET color = excluded.color, temp=excluded.temp WHERE excluded.temp < 20 RETURNING *;
 INSERT INTO upsert_test_space (time, device_id, temp, color) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8') ON CONFLICT (time, device_id)
 DO UPDATE SET color = excluded.color WHERE upsert_test_space.temp < 20 RETURNING *;
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_1) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-1-new') ON CONFLICT (time, device_id)
+DO UPDATE SET device_id_1 = excluded.device_id_1 RETURNING *;
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_1) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-1-new') ON CONFLICT (time, device_id)
+DO UPDATE SET device_id_1 = 'device-id-1-new-2', color = 'orange9'  RETURNING *;
 
+SELECT * FROM upsert_test_space;
+
+ALTER TABLE upsert_test_space DROP device_id_1, ADD device_id_2 char(20);
+INSERT INTO upsert_test_space (time, device_id, temp, color, device_id_2) VALUES ('2017-01-20T09:00:01', 'dev5', 43.5, 'orange8', 'device-id-2')
+ON CONFLICT (time, device_id)
+DO UPDATE SET device_id_2 = 'device-id-2-new', color = 'orange10' RETURNING *;
 
-SELECT * FROM upsert_test_space; 
 
 WITH CTE AS (
     INSERT INTO upsert_test_multi_unique