1 files changed, 51 insertions, 56 deletions

diff --git a/src/backend/optimizer/util/inherit.c b/src/backend/optimizer/util/inherit.c
index 6c93c224774..bf3ea26cf4b 100644
--- a/src/backend/optimizer/util/inherit.c
+++ b/src/backend/optimizer/util/inherit.c
@@ -25,6 +25,7 @@
 #include "optimizer/inherit.h"
 #include "optimizer/optimizer.h"
 #include "optimizer/pathnode.h"
+#include "optimizer/plancat.h"
 #include "optimizer/planmain.h"
 #include "optimizer/planner.h"
 #include "optimizer/prep.h"
@@ -51,7 +52,7 @@ static Bitmapset *translate_col_privs(const Bitmapset *parent_privs,
 									  List *translated_vars);
 static Bitmapset *translate_col_privs_multilevel(PlannerInfo *root,
 												 RelOptInfo *rel,
-												 RelOptInfo *top_parent_rel,
+												 RelOptInfo *parent_rel,
 												 Bitmapset *parent_cols);
 static void expand_appendrel_subquery(PlannerInfo *root, RelOptInfo *rel,
 									  RangeTblEntry *rte, Index rti);
@@ -345,11 +346,6 @@ expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 		root->partColsUpdated =
 			has_partition_attrs(parentrel, parent_updatedCols, NULL);
 
-	/*
-	 * There shouldn't be any generated columns in the partition key.
-	 */
-	Assert(!has_partition_attrs(parentrel, parentrte->extraUpdatedCols, NULL));
-
 	/* Nothing further to do here if there are no partitions. */
 	if (partdesc->nparts == 0)
 		return;
@@ -566,13 +562,6 @@ expand_single_inheritance_child(PlannerInfo *root, RangeTblEntry *parentrte,
 	childrte->alias = childrte->eref = makeAlias(parentrte->eref->aliasname,
 												 child_colnames);
 
-	/* Translate the bitmapset of generated columns being updated. */
-	if (childOID != parentOID)
-		childrte->extraUpdatedCols = translate_col_privs(parentrte->extraUpdatedCols,
-														 appinfo->translated_vars);
-	else
-		childrte->extraUpdatedCols = bms_copy(parentrte->extraUpdatedCols);
-
 	/*
 	 * Store the RTE and appinfo in the respective PlannerInfo arrays, which
 	 * the caller must already have allocated space for.
@@ -672,21 +661,16 @@ get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel)
 	Assert(IS_SIMPLE_REL(rel));
 
 	/*
-	 * We obtain updatedCols and extraUpdatedCols for the query's result
-	 * relation.  Then, if necessary, we map it to the column numbers of the
-	 * relation for which they were requested.
+	 * We obtain updatedCols for the query's result relation.  Then, if
+	 * necessary, we map it to the column numbers of the relation for which
+	 * they were requested.
 	 */
 	relid = root->parse->resultRelation;
 	rte = planner_rt_fetch(relid, root);
 	perminfo = getRTEPermissionInfo(root->parse->rteperminfos, rte);
 
 	updatedCols = perminfo->updatedCols;
-	extraUpdatedCols = rte->extraUpdatedCols;
 
-	/*
-	 * For "other" rels, we must look up the root parent relation mentioned in
-	 * the query, and translate the column numbers.
-	 */
 	if (rel->relid != relid)
 	{
 		RelOptInfo *top_parent_rel = find_base_rel(root, relid);
@@ -695,10 +679,15 @@ get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel)
 
 		updatedCols = translate_col_privs_multilevel(root, rel, top_parent_rel,
 													 updatedCols);
-		extraUpdatedCols = translate_col_privs_multilevel(root, rel, top_parent_rel,
-														  extraUpdatedCols);
 	}
 
+	/*
+	 * Now we must check to see if there are any generated columns that depend
+	 * on the updatedCols, and add them to the result.
+	 */
+	extraUpdatedCols = get_dependent_generated_columns(root, rel->relid,
+													   updatedCols);
+
 	return bms_union(updatedCols, extraUpdatedCols);
 }
 
@@ -755,6 +744,45 @@ translate_col_privs(const Bitmapset *parent_privs,
 }
 
 /*
+ * translate_col_privs_multilevel
+ *		Recursively translates the column numbers contained in 'parent_cols'
+ *		to the column numbers of a descendant relation given by 'rel'
+ *
+ * Note that because this is based on translate_col_privs, it will expand
+ * a whole-row reference into all inherited columns.  This is not an issue
+ * for current usages, but beware.
+ */
+static Bitmapset *
+translate_col_privs_multilevel(PlannerInfo *root, RelOptInfo *rel,
+							   RelOptInfo *parent_rel,
+							   Bitmapset *parent_cols)
+{
+	AppendRelInfo *appinfo;
+
+	/* Fast path for easy case. */
+	if (parent_cols == NULL)
+		return NULL;
+
+	/* Recurse if immediate parent is not the top parent. */
+	if (rel->parent != parent_rel)
+	{
+		if (rel->parent)
+			parent_cols = translate_col_privs_multilevel(root, rel->parent,
+														 parent_rel,
+														 parent_cols);
+		else
+			elog(ERROR, "rel with relid %u is not a child rel", rel->relid);
+	}
+
+	/* Now translate for this child. */
+	Assert(root->append_rel_array != NULL);
+	appinfo = root->append_rel_array[rel->relid];
+	Assert(appinfo != NULL);
+
+	return translate_col_privs(parent_cols, appinfo->translated_vars);
+}
+
+/*
  * expand_appendrel_subquery
  *		Add "other rel" RelOptInfos for the children of an appendrel baserel
  *
@@ -920,36 +948,3 @@ apply_child_basequals(PlannerInfo *root, RelOptInfo *parentrel,
 
 	return true;
 }
-
-/*
- * translate_col_privs_multilevel
- *		Recursively translates the column numbers contained in 'parent_cols'
- *		to the columns numbers of a descendent relation given by 'rel'
- */
-static Bitmapset *
-translate_col_privs_multilevel(PlannerInfo *root, RelOptInfo *rel,
-							   RelOptInfo *top_parent_rel,
-							   Bitmapset *parent_cols)
-{
-	AppendRelInfo *appinfo;
-
-	if (parent_cols == NULL)
-		return NULL;
-
-	/* Recurse if immediate parent is not the top parent. */
-	if (rel->parent != top_parent_rel)
-	{
-		if (rel->parent)
-			parent_cols = translate_col_privs_multilevel(root, rel->parent,
-														 top_parent_rel,
-														 parent_cols);
-		else
-			elog(ERROR, "rel with relid %u is not a child rel", rel->relid);
-	}
-
-	Assert(root->append_rel_array != NULL);
-	appinfo = root->append_rel_array[rel->relid];
-	Assert(appinfo != NULL);
-
-	return translate_col_privs(parent_cols, appinfo->translated_vars);
-}