virt-gen-cols-with-returning-old-new.patch.no-cfbot

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
new file mode 100644
index b6464d4..2c9d8ae
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -96,7 +96,8 @@ static List *matchLocks(CmdType event, R
 						int varno, Query *parsetree, bool *hasUpdate);
 static Query *fireRIRrules(Query *parsetree, List *activeRIRs);
 static Bitmapset *adjust_view_column_set(Bitmapset *cols, List *targetlist);
-static Node *expand_generated_columns_internal(Node *node, Relation rel, int rt_index, RangeTblEntry *rte);
+static Node *expand_generated_columns_internal(Node *node, Relation rel, int rt_index,
+								  RangeTblEntry *rte, int result_relation);
 
 
 /*
@@ -2326,8 +2327,10 @@ fireRIRrules(Query *parsetree, List *act
 		 * Note that subqueries in virtual generated column expressions are
 		 * not currently supported, so this cannot add any more sublinks.
 		 */
-		parsetree = (Query *) expand_generated_columns_internal((Node *) parsetree,
-																rel, rt_index, rte);
+		parsetree = (Query *)
+			expand_generated_columns_internal((Node *) parsetree,
+											  rel, rt_index, rte,
+											  parsetree->resultRelation);
 
 		table_close(rel, NoLock);
 	}
@@ -4449,9 +4452,18 @@ RewriteQuery(Query *parsetree, List *rew
  * If the table contains virtual generated columns, build a target list
  * containing the expanded expressions and use ReplaceVarsFromTargetList() to
  * do the replacements.
+ *
+ * Vars matching rt_index at the current query level are replaced by the
+ * virtual generated column expressions from rel, if there are any.
+ *
+ * The caller must also provide rte, the RTE describing the target relation,
+ * in order to handle any whole-row Vars referencing the target, and
+ * result_relation, the index of the result relation, if this is part of an
+ * INSERT/UPDATE/DELETE/MERGE query.
  */
 static Node *
-expand_generated_columns_internal(Node *node, Relation rel, int rt_index, RangeTblEntry *rte)
+expand_generated_columns_internal(Node *node, Relation rel, int rt_index,
+								  RangeTblEntry *rte, int result_relation)
 {
 	TupleDesc	tupdesc;
 
@@ -4502,7 +4514,10 @@ expand_generated_columns_internal(Node *
 
 		Assert(list_length(tlist) > 0);
 
-		node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist, REPLACEVARS_CHANGE_VARNO, rt_index, NULL);
+		node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist,
+										 result_relation,
+										 REPLACEVARS_CHANGE_VARNO, rt_index,
+										 NULL);
 	}
 
 	return node;
@@ -4529,7 +4544,7 @@ expand_generated_columns_in_expr(Node *n
 		rte->rtekind = RTE_RELATION;
 		rte->relid = RelationGetRelid(rel);
 
-		node = expand_generated_columns_internal(node, rel, rt_index, rte);
+		node = expand_generated_columns_internal(node, rel, rt_index, rte, 0);
 	}
 
 	return node;
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
new file mode 100644
index 5a74c0a..367605d
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -190,7 +190,12 @@ SELECT * FROM gtest1 ORDER BY a;
  2 | 4
 (2 rows)
 
-UPDATE gtest1 SET a = 3 WHERE b = 4;
+UPDATE gtest1 SET a = 3 WHERE b = 4 RETURNING old.*, new.*;
+ a | b | a | b 
+---+---+---+---
+ 2 | 4 | 3 | 6
+(1 row)
+
 SELECT * FROM gtest1 ORDER BY a;
  a | b 
 ---+---
@@ -216,7 +221,14 @@ CREATE TABLE gtestm (
 INSERT INTO gtestm VALUES (1, 5, 100);
 MERGE INTO gtestm t USING (VALUES (1, 10), (2, 20)) v(id, f1) ON t.id = v.id
   WHEN MATCHED THEN UPDATE SET f1 = v.f1
-  WHEN NOT MATCHED THEN INSERT VALUES (v.id, v.f1, 200);
+  WHEN NOT MATCHED THEN INSERT VALUES (v.id, v.f1, 200)
+  RETURNING merge_action(), old.*, new.*;
+ merge_action | id | f1 | f2  | f3 | f4  | id | f1 | f2  | f3 | f4  
+--------------+----+----+-----+----+-----+----+----+-----+----+-----
+ UPDATE       |  1 |  5 | 100 | 10 | 200 |  1 | 10 | 100 | 20 | 200
+ INSERT       |    |    |     |    |     |  2 | 20 | 200 | 40 | 400
+(2 rows)
+
 SELECT * FROM gtestm ORDER BY id;
  id | f1 | f2  | f3 | f4  
 ----+----+-----+----+-----
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
new file mode 100644
index f855898..91aec4e
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -84,7 +84,7 @@ DROP TABLE gtestx;
 
 -- test UPDATE/DELETE quals
 SELECT * FROM gtest1 ORDER BY a;
-UPDATE gtest1 SET a = 3 WHERE b = 4;
+UPDATE gtest1 SET a = 3 WHERE b = 4 RETURNING old.*, new.*;
 SELECT * FROM gtest1 ORDER BY a;
 DELETE FROM gtest1 WHERE b = 2;
 SELECT * FROM gtest1 ORDER BY a;
@@ -100,7 +100,8 @@ CREATE TABLE gtestm (
 INSERT INTO gtestm VALUES (1, 5, 100);
 MERGE INTO gtestm t USING (VALUES (1, 10), (2, 20)) v(id, f1) ON t.id = v.id
   WHEN MATCHED THEN UPDATE SET f1 = v.f1
-  WHEN NOT MATCHED THEN INSERT VALUES (v.id, v.f1, 200);
+  WHEN NOT MATCHED THEN INSERT VALUES (v.id, v.f1, 200)
+  RETURNING merge_action(), old.*, new.*;
 SELECT * FROM gtestm ORDER BY id;
 DROP TABLE gtestm;