Add OLD/NEW support to RETURNING in DML queries.

This allows the RETURNING list of INSERT/UPDATE/DELETE/MERGE queries
to explicitly return old and new values by using the special aliases
"old" and "new", which are automatically added to the query (if not
already defined) while parsing its RETURNING list, allowing things
like:

  RETURNING old.colname, new.colname, ...

  RETURNING old.*, new.*

Additionally, a new syntax is supported, allowing the names "old" and
"new" to be changed to user-supplied alias names, e.g.:

  RETURNING WITH (OLD AS o, NEW AS n) o.colname, n.colname, ...

This is useful when the names "old" and "new" are already defined,
such as inside trigger functions, allowing backwards compatibility to
be maintained -- the interpretation of any existing queries that
happen to already refer to relations called "old" or "new", or use
those as aliases for other relations, is not changed.

For an INSERT, old values will generally be NULL, and for a DELETE,
new values will generally be NULL, but that may change for an INSERT
with an ON CONFLICT ... DO UPDATE clause, or if a query rewrite rule
changes the command type. Therefore, we put no restrictions on the use
of old and new in any DML queries.

Dean Rasheed, reviewed by Jian He and Jeff Davis.

Discussion: https://postgr.es/m/CAEZATCWx0J0-v=Qjc6gXzR=KtsdvAE7Ow=D=mu50AgOe+pvisQ@mail.gmail.com

src/backend/optimizer/path/allpaths.c

@@ -3985,6 +3985,7 @@ subquery_push_qual(Query *subquery, RangeTblEntry *rte, Index rti, Node *qual)
 		 */
 		qual = ReplaceVarsFromTargetList(qual, rti, 0, rte,
 										 subquery->targetList,
+										 subquery->resultRelation,
 										 REPLACEVARS_REPORT_ERROR, 0,
 										 &subquery->hasSubLinks);
 

src/backend/optimizer/plan/createplan.c

@@ -7121,6 +7121,8 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 				 int epqParam)
 {
 	ModifyTable *node = makeNode(ModifyTable);
+	bool		returning_old_or_new = false;
+	bool		returning_old_or_new_valid = false;
 	List	   *fdw_private_list;
 	Bitmapset  *direct_modify_plans;
 	ListCell   *lc;
@@ -7185,6 +7187,8 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 	}
 	node->updateColnosLists = updateColnosLists;
 	node->withCheckOptionLists = withCheckOptionLists;
+	node->returningOldAlias = root->parse->returningOldAlias;
+	node->returningNewAlias = root->parse->returningNewAlias;
 	node->returningLists = returningLists;
 	node->rowMarks = rowMarks;
 	node->mergeActionLists = mergeActionLists;
@@ -7265,7 +7269,8 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 		 * callback functions needed for that and (2) there are no local
 		 * structures that need to be run for each modified row: row-level
 		 * triggers on the foreign table, stored generated columns, WITH CHECK
-		 * OPTIONs from parent views.
+		 * OPTIONs from parent views, or Vars returning OLD/NEW in the
+		 * RETURNING list.
 		 */
 		direct_modify = false;
 		if (fdwroutine != NULL &&
@@ -7276,7 +7281,18 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 			withCheckOptionLists == NIL &&
 			!has_row_triggers(root, rti, operation) &&
 			!has_stored_generated_columns(root, rti))
-			direct_modify = fdwroutine->PlanDirectModify(root, node, rti, i);
+		{
+			/* returning_old_or_new is the same for all result relations */
+			if (!returning_old_or_new_valid)
+			{
+				returning_old_or_new =
+					contain_vars_returning_old_or_new((Node *)
+													  root->parse->returningList);
+				returning_old_or_new_valid = true;
+			}
+			if (!returning_old_or_new)
+				direct_modify = fdwroutine->PlanDirectModify(root, node, rti, i);
+		}
 		if (direct_modify)
 			direct_modify_plans = bms_add_member(direct_modify_plans, i);
 

src/backend/optimizer/plan/setrefs.c

@@ -3070,6 +3070,21 @@ fix_join_expr_mutator(Node *node, fix_join_expr_context *context)
 	{
 		Var		   *var = (Var *) node;
 
+		/*
+		 * Verify that Vars with non-default varreturningtype only appear in
+		 * the RETURNING list, and refer to the target relation.
+		 */
+		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
+		{
+			if (context->inner_itlist != NULL ||
+				context->outer_itlist == NULL ||
+				context->acceptable_rel == 0)
+				elog(ERROR, "variable returning old/new found outside RETURNING list");
+			if (var->varno != context->acceptable_rel)
+				elog(ERROR, "wrong varno %d (expected %d) for variable returning old/new",
+					 var->varno, context->acceptable_rel);
+		}
+
 		/* Look for the var in the input tlists, first in the outer */
 		if (context->outer_itlist)
 		{

src/backend/optimizer/plan/subselect.c

@@ -354,17 +354,19 @@ build_subplan(PlannerInfo *root, Plan *plan, Path *path,
 		Node	   *arg = pitem->item;
 
 		/*
-		 * The Var, PlaceHolderVar, Aggref or GroupingFunc has already been
-		 * adjusted to have the correct varlevelsup, phlevelsup, or
-		 * agglevelsup.
+		 * The Var, PlaceHolderVar, Aggref, GroupingFunc, or ReturningExpr has
+		 * already been adjusted to have the correct varlevelsup, phlevelsup,
+		 * agglevelsup, or retlevelsup.
 		 *
-		 * If it's a PlaceHolderVar, Aggref or GroupingFunc, its arguments
-		 * might contain SubLinks, which have not yet been processed (see the
-		 * comments for SS_replace_correlation_vars).  Do that now.
+		 * If it's a PlaceHolderVar, Aggref, GroupingFunc, or ReturningExpr,
+		 * its arguments might contain SubLinks, which have not yet been
+		 * processed (see the comments for SS_replace_correlation_vars).  Do
+		 * that now.
 		 */
 		if (IsA(arg, PlaceHolderVar) ||
 			IsA(arg, Aggref) ||
-			IsA(arg, GroupingFunc))
+			IsA(arg, GroupingFunc) ||
+			IsA(arg, ReturningExpr))
 			arg = SS_process_sublinks(root, arg, false);
 
 		splan->parParam = lappend_int(splan->parParam, pitem->paramId);
@@ -1863,8 +1865,8 @@ convert_EXISTS_to_ANY(PlannerInfo *root, Query *subselect,
 /*
  * Replace correlation vars (uplevel vars) with Params.
  *
- * Uplevel PlaceHolderVars, aggregates, GROUPING() expressions, and
- * MergeSupportFuncs are replaced, too.
+ * Uplevel PlaceHolderVars, aggregates, GROUPING() expressions,
+ * MergeSupportFuncs, and ReturningExprs are replaced, too.
  *
  * Note: it is critical that this runs immediately after SS_process_sublinks.
  * Since we do not recurse into the arguments of uplevel PHVs and aggregates,
@@ -1924,6 +1926,12 @@ replace_correlation_vars_mutator(Node *node, PlannerInfo *root)
 			return (Node *) replace_outer_merge_support(root,
 														(MergeSupportFunc *) node);
 	}
+	if (IsA(node, ReturningExpr))
+	{
+		if (((ReturningExpr *) node)->retlevelsup > 0)
+			return (Node *) replace_outer_returning(root,
+													(ReturningExpr *) node);
+	}
 	return expression_tree_mutator(node, replace_correlation_vars_mutator, root);
 }
 
@@ -1977,11 +1985,11 @@ process_sublinks_mutator(Node *node, process_sublinks_context *context)
 	}
 
 	/*
-	 * Don't recurse into the arguments of an outer PHV, Aggref or
-	 * GroupingFunc here.  Any SubLinks in the arguments have to be dealt with
-	 * at the outer query level; they'll be handled when build_subplan
-	 * collects the PHV, Aggref or GroupingFunc into the arguments to be
-	 * passed down to the current subplan.
+	 * Don't recurse into the arguments of an outer PHV, Aggref, GroupingFunc,
+	 * or ReturningExpr here.  Any SubLinks in the arguments have to be dealt
+	 * with at the outer query level; they'll be handled when build_subplan
+	 * collects the PHV, Aggref, GroupingFunc, or ReturningExpr into the
+	 * arguments to be passed down to the current subplan.
 	 */
 	if (IsA(node, PlaceHolderVar))
 	{
@@ -1998,6 +2006,11 @@ process_sublinks_mutator(Node *node, process_sublinks_context *context)
 		if (((GroupingFunc *) node)->agglevelsup > 0)
 			return node;
 	}
+	else if (IsA(node, ReturningExpr))
+	{
+		if (((ReturningExpr *) node)->retlevelsup > 0)
+			return node;
+	}
 
 	/*
 	 * We should never see a SubPlan expression in the input (since this is
@@ -2110,7 +2123,9 @@ SS_identify_outer_params(PlannerInfo *root)
 	outer_params = NULL;
 	for (proot = root->parent_root; proot != NULL; proot = proot->parent_root)
 	{
-		/* Include ordinary Var/PHV/Aggref/GroupingFunc params */
+		/*
+		 * Include ordinary Var/PHV/Aggref/GroupingFunc/ReturningExpr params.
+		 */
 		foreach(l, proot->plan_params)
 		{
 			PlannerParamItem *pitem = (PlannerParamItem *) lfirst(l);

src/backend/optimizer/prep/prepjointree.c

@@ -2539,7 +2539,8 @@ pullup_replace_vars_callback(Var *var,
 		 * expansion with varlevelsup = 0, and then adjust below if needed.
 		 */
 		expandRTE(rcon->target_rte,
-				  var->varno, 0 /* not varlevelsup */ , var->location,
+				  var->varno, 0 /* not varlevelsup */ ,
+				  var->varreturningtype, var->location,
 				  (var->vartype != RECORDOID),
 				  &colnames, &fields);
 		/* Expand the generated per-field Vars, but don't insert PHVs there */

src/backend/optimizer/util/appendinfo.c

@@ -253,6 +253,13 @@ adjust_appendrel_attrs_mutator(Node *node,
 		 * all non-Var outputs of such subqueries, and then we could look up
 		 * the pre-existing PHV here.  Or perhaps just wrap the translations
 		 * that way to begin with?
+		 *
+		 * If var->varreturningtype is not VAR_RETURNING_DEFAULT, then that
+		 * also needs to be copied to the translated Var.  That too would fail
+		 * if the translation wasn't a Var, but that should never happen since
+		 * a non-default var->varreturningtype is only used for Vars referring
+		 * to the result relation, which should never be a flattened UNION ALL
+		 * subquery.
 		 */
 
 		for (cnt = 0; cnt < nappinfos; cnt++)
@@ -283,9 +290,17 @@ adjust_appendrel_attrs_mutator(Node *node,
 					elog(ERROR, "attribute %d of relation \"%s\" does not exist",
 						 var->varattno, get_rel_name(appinfo->parent_reloid));
 				if (IsA(newnode, Var))
+				{
+					((Var *) newnode)->varreturningtype = var->varreturningtype;
 					((Var *) newnode)->varnullingrels = var->varnullingrels;
-				else if (var->varnullingrels != NULL)
-					elog(ERROR, "failed to apply nullingrels to a non-Var");
+				}
+				else
+				{
+					if (var->varreturningtype != VAR_RETURNING_DEFAULT)
+						elog(ERROR, "failed to apply returningtype to a non-Var");
+					if (var->varnullingrels != NULL)
+						elog(ERROR, "failed to apply nullingrels to a non-Var");
+				}
 				return newnode;
 			}
 			else if (var->varattno == 0)
@@ -339,6 +354,8 @@ adjust_appendrel_attrs_mutator(Node *node,
 					rowexpr->colnames = copyObject(rte->eref->colnames);
 					rowexpr->location = -1;
 
+					if (var->varreturningtype != VAR_RETURNING_DEFAULT)
+						elog(ERROR, "failed to apply returningtype to a non-Var");
 					if (var->varnullingrels != NULL)
 						elog(ERROR, "failed to apply nullingrels to a non-Var");
 

src/backend/optimizer/util/clauses.c

@@ -1295,6 +1295,7 @@ contain_leaked_vars_walker(Node *node, void *context)
 		case T_NullTest:
 		case T_BooleanTest:
 		case T_NextValueExpr:
+		case T_ReturningExpr:
 		case T_List:
 
 			/*
@@ -3404,6 +3405,8 @@ eval_const_expressions_mutator(Node *node,
 										 fselect->resulttypmod,
 										 fselect->resultcollid,
 										 ((Var *) arg)->varlevelsup);
+						/* New Var has same OLD/NEW returning as old one */
+						newvar->varreturningtype = ((Var *) arg)->varreturningtype;
 						/* New Var is nullable by same rels as the old one */
 						newvar->varnullingrels = ((Var *) arg)->varnullingrels;
 						return (Node *) newvar;

src/backend/optimizer/util/paramassign.c

@@ -91,6 +91,7 @@ assign_param_for_var(PlannerInfo *root, Var *var)
 				pvar->vartype == var->vartype &&
 				pvar->vartypmod == var->vartypmod &&
 				pvar->varcollid == var->varcollid &&
+				pvar->varreturningtype == var->varreturningtype &&
 				bms_equal(pvar->varnullingrels, var->varnullingrels))
 				return pitem->paramId;
 		}
@@ -358,6 +359,52 @@ replace_outer_merge_support(PlannerInfo *root, MergeSupportFunc *msf)
 	return retval;
 }
 
+/*
+ * Generate a Param node to replace the given ReturningExpr expression which
+ * is expected to have retlevelsup > 0 (ie, it is not local).  Record the need
+ * for the ReturningExpr in the proper upper-level root->plan_params.
+ */
+Param *
+replace_outer_returning(PlannerInfo *root, ReturningExpr *rexpr)
+{
+	Param	   *retval;
+	PlannerParamItem *pitem;
+	Index		levelsup;
+	Oid			ptype = exprType((Node *) rexpr->retexpr);
+
+	Assert(rexpr->retlevelsup > 0 && rexpr->retlevelsup < root->query_level);
+
+	/* Find the query level the ReturningExpr belongs to */
+	for (levelsup = rexpr->retlevelsup; levelsup > 0; levelsup--)
+		root = root->parent_root;
+
+	/*
+	 * It does not seem worthwhile to try to de-duplicate references to outer
+	 * ReturningExprs.  Just make a new slot every time.
+	 */
+	rexpr = copyObject(rexpr);
+	IncrementVarSublevelsUp((Node *) rexpr, -((int) rexpr->retlevelsup), 0);
+	Assert(rexpr->retlevelsup == 0);
+
+	pitem = makeNode(PlannerParamItem);
+	pitem->item = (Node *) rexpr;
+	pitem->paramId = list_length(root->glob->paramExecTypes);
+	root->glob->paramExecTypes = lappend_oid(root->glob->paramExecTypes,
+											 ptype);
+
+	root->plan_params = lappend(root->plan_params, pitem);
+
+	retval = makeNode(Param);
+	retval->paramkind = PARAM_EXEC;
+	retval->paramid = pitem->paramId;
+	retval->paramtype = ptype;
+	retval->paramtypmod = exprTypmod((Node *) rexpr->retexpr);
+	retval->paramcollid = exprCollation((Node *) rexpr->retexpr);
+	retval->location = exprLocation((Node *) rexpr->retexpr);
+
+	return retval;
+}
+
 /*
  * Generate a Param node to replace the given Var,
  * which is expected to come from some upper NestLoop plan node.

src/backend/optimizer/util/plancat.c

@@ -1857,8 +1857,8 @@ build_physical_tlist(PlannerInfo *root, RelOptInfo *rel)
 		case RTE_NAMEDTUPLESTORE:
 		case RTE_RESULT:
 			/* Not all of these can have dropped cols, but share code anyway */
-			expandRTE(rte, varno, 0, -1, true /* include dropped */ ,
-					  NULL, &colvars);
+			expandRTE(rte, varno, 0, VAR_RETURNING_DEFAULT, -1,
+					  true /* include dropped */ , NULL, &colvars);
 			foreach(l, colvars)
 			{
 				var = (Var *) lfirst(l);

src/backend/optimizer/util/var.c

@@ -76,6 +76,7 @@ static bool pull_varattnos_walker(Node *node, pull_varattnos_context *context);
 static bool pull_vars_walker(Node *node, pull_vars_context *context);
 static bool contain_var_clause_walker(Node *node, void *context);
 static bool contain_vars_of_level_walker(Node *node, int *sublevels_up);
+static bool contain_vars_returning_old_or_new_walker(Node *node, void *context);
 static bool locate_var_of_level_walker(Node *node,
 									   locate_var_of_level_context *context);
 static bool pull_var_clause_walker(Node *node,
@@ -492,6 +493,49 @@ contain_vars_of_level_walker(Node *node, int *sublevels_up)
 }
 
 
+/*
+ * contain_vars_returning_old_or_new
+ *	  Recursively scan a clause to discover whether it contains any Var nodes
+ *	  (of the current query level) whose varreturningtype is VAR_RETURNING_OLD
+ *	  or VAR_RETURNING_NEW.
+ *
+ *	  Returns true if any found.
+ *
+ * Any ReturningExprs are also detected --- if an OLD/NEW Var was rewritten,
+ * we still regard this as a clause that returns OLD/NEW values.
+ *
+ * Does not examine subqueries, therefore must only be used after reduction
+ * of sublinks to subplans!
+ */
+bool
+contain_vars_returning_old_or_new(Node *node)
+{
+	return contain_vars_returning_old_or_new_walker(node, NULL);
+}
+
+static bool
+contain_vars_returning_old_or_new_walker(Node *node, void *context)
+{
+	if (node == NULL)
+		return false;
+	if (IsA(node, Var))
+	{
+		if (((Var *) node)->varlevelsup == 0 &&
+			((Var *) node)->varreturningtype != VAR_RETURNING_DEFAULT)
+			return true;		/* abort the tree traversal and return true */
+		return false;
+	}
+	if (IsA(node, ReturningExpr))
+	{
+		if (((ReturningExpr *) node)->retlevelsup == 0)
+			return true;		/* abort the tree traversal and return true */
+		return false;
+	}
+	return expression_tree_walker(node, contain_vars_returning_old_or_new_walker,
+								  context);
+}
+
+
 /*
  * locate_var_of_level
  *	  Find the parse location of any Var of the specified query level.