Teach map_partition_varattnos to handle whole-row expressions.

Otherwise, partitioned tables with RETURNING expressions or subject
to a WITH CHECK OPTION do not work properly.

Amit Langote, reviewed by Amit Khandekar and Etsuro Fujita.  A few
comment changes by me.

Discussion: http://postgr.es/m/9a39df80-871e-6212-0684-f93c83be4097@lab.ntt.co.jp

src/backend/catalog/partition.c

@@ -898,16 +898,20 @@ get_qual_from_partbound(Relation rel, Relation parent,
  * We must allow for cases where physical attnos of a partition can be
  * different from the parent's.
  *
+ * If found_whole_row is not NULL, *found_whole_row returns whether a
+ * whole-row variable was found in the input expression.
+ *
  * Note: this will work on any node tree, so really the argument and result
  * should be declared "Node *".  But a substantial majority of the callers
  * are working on Lists, so it's less messy to do the casts internally.
  */
 List *
 map_partition_varattnos(List *expr, int target_varno,
-						Relation partrel, Relation parent)
+						Relation partrel, Relation parent,
+						bool *found_whole_row)
 {
 	AttrNumber *part_attnos;
-	bool		found_whole_row;
+	bool		my_found_whole_row;
 
 	if (expr == NIL)
 		return NIL;
@@ -919,10 +923,10 @@ map_partition_varattnos(List *expr, int target_varno,
 										target_varno, 0,
 										part_attnos,
 										RelationGetDescr(parent)->natts,
-										&found_whole_row);
-	/* There can never be a whole-row reference here */
+										RelationGetForm(partrel)->reltype,
+										&my_found_whole_row);
 	if (found_whole_row)
-		elog(ERROR, "unexpected whole-row reference found in partition key");
+		*found_whole_row = my_found_whole_row;
 
 	return expr;
 }
@@ -1783,6 +1787,7 @@ generate_partition_qual(Relation rel)
 	List	   *my_qual = NIL,
 			   *result = NIL;
 	Relation	parent;
+	bool		found_whole_row;
 
 	/* Guard against stack overflow due to overly deep partition tree */
 	check_stack_depth();
@@ -1825,7 +1830,11 @@ generate_partition_qual(Relation rel)
 	 * in it to bear this relation's attnos. It's safe to assume varno = 1
 	 * here.
 	 */
-	result = map_partition_varattnos(result, 1, rel, parent);
+	result = map_partition_varattnos(result, 1, rel, parent,
+									 &found_whole_row);
+	/* There can never be a whole-row reference here */
+	if (found_whole_row)
+		elog(ERROR, "unexpected whole-row reference found in partition key");
 
 	/* Save a copy in the relcache */
 	oldcxt = MemoryContextSwitchTo(CacheMemoryContext);

src/backend/commands/tablecmds.c

@@ -1989,7 +1989,7 @@ MergeAttributes(List *schema, List *supers, char relpersistence,
 				expr = map_variable_attnos(stringToNode(check[i].ccbin),
 										   1, 0,
 										   newattno, tupleDesc->natts,
-										   &found_whole_row);
+										   InvalidOid, &found_whole_row);
 
 				/*
 				 * For the moment we have to reject whole-row variables. We
@@ -8874,7 +8874,7 @@ ATPrepAlterColumnType(List **wqueue,
 					map_variable_attnos(def->cooked_default,
 										1, 0,
 										attmap, RelationGetDescr(rel)->natts,
-										&found_whole_row);
+										InvalidOid, &found_whole_row);
 				if (found_whole_row)
 					ereport(ERROR,
 							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
@@ -13713,6 +13713,7 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 			Oid			part_relid = lfirst_oid(lc);
 			Relation	part_rel;
 			Expr	   *constr;
+			bool		found_whole_row;
 
 			/* Lock already taken */
 			if (part_relid != RelationGetRelid(attachRel))
@@ -13738,7 +13739,12 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 			constr = linitial(partConstraint);
 			tab->partition_constraint = (Expr *)
 				map_partition_varattnos((List *) constr, 1,
-										part_rel, rel);
+										part_rel, rel,
+										&found_whole_row);
+			/* There can never be a whole-row reference here */
+			if (found_whole_row)
+				elog(ERROR, "unexpected whole-row reference found in partition key");
+
 			/* keep our lock until commit */
 			if (part_rel != attachRel)
 				heap_close(part_rel, NoLock);

src/backend/executor/nodeModifyTable.c

@@ -1996,7 +1996,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 			/* varno = node->nominalRelation */
 			mapped_wcoList = map_partition_varattnos(wcoList,
 													 node->nominalRelation,
-													 partrel, rel);
+													 partrel, rel, NULL);
 			foreach(ll, mapped_wcoList)
 			{
 				WithCheckOption *wco = castNode(WithCheckOption, lfirst(ll));
@@ -2069,7 +2069,7 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 			/* varno = node->nominalRelation */
 			rlist = map_partition_varattnos(returningList,
 											node->nominalRelation,
-											partrel, rel);
+											partrel, rel, NULL);
 			resultRelInfo->ri_projectReturning =
 				ExecBuildProjectionInfo(rlist, econtext, slot, &mtstate->ps,
 										resultRelInfo->ri_RelationDesc->rd_att);

src/backend/parser/parse_utilcmd.c

@@ -1107,7 +1107,7 @@ transformTableLikeClause(CreateStmtContext *cxt, TableLikeClause *table_like_cla
 			ccbin_node = map_variable_attnos(stringToNode(ccbin),
 											 1, 0,
 											 attmap, tupleDesc->natts,
-											 &found_whole_row);
+											 InvalidOid, &found_whole_row);
 
 			/*
 			 * We reject whole-row variables because the whole point of LIKE
@@ -1463,7 +1463,7 @@ generateClonedIndexStmt(CreateStmtContext *cxt, Relation source_idx,
 			indexkey = map_variable_attnos(indexkey,
 										   1, 0,
 										   attmap, attmap_length,
-										   &found_whole_row);
+										   InvalidOid, &found_whole_row);
 
 			/* As in transformTableLikeClause, reject whole-row variables */
 			if (found_whole_row)
@@ -1539,7 +1539,7 @@ generateClonedIndexStmt(CreateStmtContext *cxt, Relation source_idx,
 		pred_tree = map_variable_attnos(pred_tree,
 										1, 0,
 										attmap, attmap_length,
-										&found_whole_row);
+										InvalidOid, &found_whole_row);
 
 		/* As in transformTableLikeClause, reject whole-row variables */
 		if (found_whole_row)

src/backend/rewrite/rewriteManip.c

@@ -1203,14 +1203,12 @@ replace_rte_variables_mutator(Node *node,
  * appear in the expression.
  *
  * If the expression tree contains a whole-row Var for the target RTE,
- * the Var is not changed but *found_whole_row is returned as TRUE.
- * For most callers this is an error condition, but we leave it to the caller
- * to report the error so that useful context can be provided.  (In some
- * usages it would be appropriate to modify the Var's vartype and insert a
- * ConvertRowtypeExpr node to map back to the original vartype.  We might
- * someday extend this function's API to support that.  For now, the only
- * concession to that future need is that this function is a tree mutator
- * not just a walker.)
+ * *found_whole_row is returned as TRUE.  In addition, if to_rowtype is
+ * not InvalidOid, we modify the Var's vartype and insert a ConvertRowTypeExpr
+ * to map back to the orignal rowtype.  Callers that don't provide to_rowtype
+ * should report an error if *found_row_type is true; we don't do that here
+ * because we don't know exactly what wording for the error message would
+ * be most appropriate.  The caller will be aware of the context.
  *
  * This could be built using replace_rte_variables and a callback function,
  * but since we don't ever need to insert sublinks, replace_rte_variables is
@@ -1223,6 +1221,8 @@ typedef struct
 	int			sublevels_up;	/* (current) nesting depth */
 	const AttrNumber *attno_map;	/* map array for user attnos */
 	int			map_length;		/* number of entries in attno_map[] */
+	/* Target type when converting whole-row vars */
+	Oid			to_rowtype;
 	bool	   *found_whole_row;	/* output flag */
 } map_variable_attnos_context;
 
@@ -1257,6 +1257,34 @@ map_variable_attnos_mutator(Node *node,
 			{
 				/* whole-row variable, warn caller */
 				*(context->found_whole_row) = true;
+
+				/* If the callers expects us to convert the same, do so. */
+				if (OidIsValid(context->to_rowtype))
+				{
+					/* No support for RECORDOID. */
+					Assert(var->vartype != RECORDOID);
+
+					/* Don't convert unless necessary. */
+					if (context->to_rowtype != var->vartype)
+					{
+						ConvertRowtypeExpr *r;
+
+						/* Var itself is converted to the requested type. */
+						newvar->vartype = context->to_rowtype;
+
+						/*
+						 * And a conversion node on top to convert back to the
+						 * original type.
+						 */
+						r = makeNode(ConvertRowtypeExpr);
+						r->arg = (Expr *) newvar;
+						r->resulttype = var->vartype;
+						r->convertformat = COERCE_IMPLICIT_CAST;
+						r->location = -1;
+
+						return (Node *) r;
+					}
+				}
 			}
 			return (Node *) newvar;
 		}
@@ -1283,7 +1311,7 @@ Node *
 map_variable_attnos(Node *node,
 					int target_varno, int sublevels_up,
 					const AttrNumber *attno_map, int map_length,
-					bool *found_whole_row)
+					Oid to_rowtype, bool *found_whole_row)
 {
 	map_variable_attnos_context context;
 
@@ -1291,6 +1319,7 @@ map_variable_attnos(Node *node,
 	context.sublevels_up = sublevels_up;
 	context.attno_map = attno_map;
 	context.map_length = map_length;
+	context.to_rowtype = to_rowtype;
 	context.found_whole_row = found_whole_row;
 
 	*found_whole_row = false;