From b43fcbbfabdead3faab40f27117a6feeaaf23857 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Fri, 12 Mar 2021 17:24:04 +0900
Subject: [PATCH v4] Fix RETURNING system columns when inserting into
 partitioned tables

Currently, the use of a virtual TupleTableSlots to hold the tuples
being inserted into partitioned tables leads to the following fixed
error when the RETURNING list contains references to system columns:

  ERROR:  virtual tuple table slot does not have system attributes

This commit arranges the inserted tuple to be placed in a
TupleTableSlot specified by a leaf partition's table AM after the
tuple is routed to it.  That allows any references to system columns
present in the RETURNING list to be resolved appropriately by the
leaf partition table AM's TupleTableSlotOps implementation.
---
 src/backend/executor/execPartition.c   | 37 +++++++++++++++++++++++---
 src/backend/executor/nodeModifyTable.c | 13 ++++-----
 src/test/regress/expected/insert.out   | 10 +++----
 src/test/regress/sql/insert.sql        |  4 +--
 4 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index b8da4c5967..077ded4d3b 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -24,6 +24,7 @@
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
+#include "optimizer/optimizer.h"
 #include "partitioning/partbounds.h"
 #include "partitioning/partdesc.h"
 #include "partitioning/partprune.h"
@@ -168,6 +169,7 @@ static ResultRelInfo *ExecInitPartitionInfo(ModifyTableState *mtstate,
 											PartitionDispatch dispatch,
 											ResultRelInfo *rootResultRelInfo,
 											int partidx);
+static bool tlist_has_system_column(List *returningList, int firstVarno);
 static void ExecInitRoutingInfo(ModifyTableState *mtstate,
 								EState *estate,
 								PartitionTupleRouting *proute,
@@ -936,6 +938,31 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 	return leaf_part_rri;
 }
 
+/*
+ * Returns if any of the Vars in returningList refers to a system column.
+ */
+static bool
+tlist_has_system_column(List *returningList, int firstVarno)
+{
+	Bitmapset  *returning_attrs = NULL;
+	int			col;
+
+	if (returningList == NIL)
+		return false;
+
+	pull_varattnos((Node *) returningList, firstVarno, &returning_attrs);
+	col = -1;
+	while ((col = bms_next_member(returning_attrs, col)) >= 0)
+	{
+		AttrNumber  attno = col + FirstLowInvalidHeapAttributeNumber;
+
+		if (attno < 0)
+			return true;
+	}
+
+	return false;
+}
+
 /*
  * ExecInitRoutingInfo
  *		Set up information needed for translating tuples between root
@@ -953,6 +980,7 @@ ExecInitRoutingInfo(ModifyTableState *mtstate,
 	ResultRelInfo *rootRelInfo = partRelInfo->ri_RootResultRelInfo;
 	MemoryContext oldcxt;
 	int			rri_index;
+	int			firstVarno = mtstate->resultRelInfo[0].ri_RangeTableIndex;
 
 	oldcxt = MemoryContextSwitchTo(proute->memcxt);
 
@@ -968,9 +996,12 @@ ExecInitRoutingInfo(ModifyTableState *mtstate,
 	 * If a partition has a different rowtype than the root parent, initialize
 	 * a slot dedicated to storing this partition's tuples.  The slot is used
 	 * for various operations that are applied to tuples after routing, such
-	 * as checking constraints.
+	 * as checking constraints.  We would also need this slot when there are
+	 * system columns present in the query's RETURNING list, because the root
+	 * parent's virtual slot would fail on trying to access them.
 	 */
-	if (partRelInfo->ri_RootToPartitionMap != NULL)
+	if (partRelInfo->ri_RootToPartitionMap != NULL ||
+		tlist_has_system_column(partRelInfo->ri_returningList, firstVarno))
 	{
 		Relation	partrel = partRelInfo->ri_RelationDesc;
 
@@ -982,8 +1013,6 @@ ExecInitRoutingInfo(ModifyTableState *mtstate,
 		partRelInfo->ri_PartitionTupleSlot =
 			table_slot_create(partrel, &estate->es_tupleTable);
 	}
-	else
-		partRelInfo->ri_PartitionTupleSlot = NULL;
 
 	/*
 	 * If the partition is a foreign table, let the FDW init itself for
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 2993ba43e3..143eda16a7 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -1993,6 +1993,7 @@ ExecPrepareTupleRouting(ModifyTableState *mtstate,
 {
 	ResultRelInfo *partrel;
 	TupleConversionMap *map;
+	TupleTableSlot *partition_slot;
 
 	/*
 	 * Lookup the target partition's ResultRelInfo.  If ExecFindPartition does
@@ -2022,15 +2023,15 @@ ExecPrepareTupleRouting(ModifyTableState *mtstate,
 	}
 
 	/*
-	 * Convert the tuple, if necessary.
+	 * Convert the tuple, if necessary, and store in the partition-specific
+	 * slot, if any.
 	 */
 	map = partrel->ri_RootToPartitionMap;
+	partition_slot = partrel->ri_PartitionTupleSlot;
 	if (map != NULL)
-	{
-		TupleTableSlot *new_slot = partrel->ri_PartitionTupleSlot;
-
-		slot = execute_attr_map_slot(map->attrMap, slot, new_slot);
-	}
+		slot = execute_attr_map_slot(map->attrMap, slot, partition_slot);
+	else if (partition_slot)
+		slot = ExecCopySlot(partition_slot, slot);
 
 	*partRelInfo = partrel;
 	return slot;
diff --git a/src/test/regress/expected/insert.out b/src/test/regress/expected/insert.out
index da50ee3b67..02403a8505 100644
--- a/src/test/regress/expected/insert.out
+++ b/src/test/regress/expected/insert.out
@@ -960,13 +960,13 @@ select tableoid::regclass, * from mcrparted order by a, b;
 (11 rows)
 
 drop table mcrparted;
--- check that wholerow vars in the RETURNING list work with partitioned tables
+-- check that wholerow and system columns in the RETURNING list work with partitioned tables
 create table returningwrtest (a int) partition by list (a);
 create table returningwrtest1 partition of returningwrtest for values in (1);
-insert into returningwrtest values (1) returning returningwrtest;
- returningwrtest 
------------------
- (1)
+insert into returningwrtest values (1) returning returningwrtest, ctid, xmin = pg_current_xact_id()::xid AS xmin_correct;
+ returningwrtest | ctid  | xmin_correct 
+-----------------+-------+--------------
+ (1)             | (0,1) | t
 (1 row)
 
 -- check also that the wholerow vars in RETURNING list are converted as needed
diff --git a/src/test/regress/sql/insert.sql b/src/test/regress/sql/insert.sql
index 963faa1614..c40a4b9c31 100644
--- a/src/test/regress/sql/insert.sql
+++ b/src/test/regress/sql/insert.sql
@@ -590,10 +590,10 @@ insert into mcrparted values ('aaa', 0), ('b', 0), ('bz', 10), ('c', -10),
 select tableoid::regclass, * from mcrparted order by a, b;
 drop table mcrparted;
 
--- check that wholerow vars in the RETURNING list work with partitioned tables
+-- check that wholerow and system columns in the RETURNING list work with partitioned tables
 create table returningwrtest (a int) partition by list (a);
 create table returningwrtest1 partition of returningwrtest for values in (1);
-insert into returningwrtest values (1) returning returningwrtest;
+insert into returningwrtest values (1) returning returningwrtest, ctid, xmin = pg_current_xact_id()::xid AS xmin_correct;
 
 -- check also that the wholerow vars in RETURNING list are converted as needed
 alter table returningwrtest add b text;
-- 
2.24.1

