v6-0002-Use-ReplaceVarsFromTargetList-for-rewriting-virtu.patch

From 7bfe4fc9ce81124d9e7fdc24a1f46099bcaec726 Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter@eisentraut.org>
Date: Wed, 4 Sep 2024 09:38:42 +0200
Subject: [PATCH v6 2/2] Use ReplaceVarsFromTargetList() for rewriting virtual
 columns

---
 src/backend/parser/analyze.c         |   8 --
 src/backend/parser/parse_clause.c    |   4 -
 src/backend/parser/parse_merge.c     |   1 -
 src/backend/parser/parse_relation.c  |   3 -
 src/backend/rewrite/rewriteHandler.c | 185 ++++++++-------------------
 src/include/nodes/parsenodes.h       |   2 -
 src/include/parser/parse_node.h      |   1 -
 7 files changed, 54 insertions(+), 150 deletions(-)

diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index e3936a0b8ea..e901203424d 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -568,7 +568,6 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	qry->hasWindowFuncs = pstate->p_hasWindowFuncs;
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasAggs = pstate->p_hasAggs;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
@@ -995,7 +994,6 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasSubLinks = pstate->p_hasSubLinks;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
@@ -1461,7 +1459,6 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	qry->hasWindowFuncs = pstate->p_hasWindowFuncs;
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasAggs = pstate->p_hasAggs;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	foreach(l, stmt->lockingClause)
 	{
@@ -1688,7 +1685,6 @@ transformValuesClause(ParseState *pstate, SelectStmt *stmt)
 	qry->jointree = makeFromExpr(pstate->p_joinlist, NULL);
 
 	qry->hasSubLinks = pstate->p_hasSubLinks;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
@@ -1940,7 +1936,6 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	qry->hasWindowFuncs = pstate->p_hasWindowFuncs;
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasAggs = pstate->p_hasAggs;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	foreach(l, lockingClause)
 	{
@@ -2415,7 +2410,6 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	qry->hasWindowFuncs = pstate->p_hasWindowFuncs;
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasAggs = pstate->p_hasAggs;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
@@ -2483,7 +2477,6 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasSubLinks = pstate->p_hasSubLinks;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
@@ -2850,7 +2843,6 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	qry->hasWindowFuncs = pstate->p_hasWindowFuncs;
 	qry->hasTargetSRFs = pstate->p_hasTargetSRFs;
 	qry->hasAggs = pstate->p_hasAggs;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	foreach(l, sstmt->lockingClause)
 	{
diff --git a/src/backend/parser/parse_clause.c b/src/backend/parser/parse_clause.c
index 48524ac3fc2..8118036495b 100644
--- a/src/backend/parser/parse_clause.c
+++ b/src/backend/parser/parse_clause.c
@@ -207,10 +207,6 @@ setTargetTable(ParseState *pstate, RangeVar *relation,
 	pstate->p_target_relation = parserOpenTable(pstate, relation,
 												RowExclusiveLock);
 
-	if (pstate->p_target_relation->rd_att->constr &&
-		pstate->p_target_relation->rd_att->constr->has_generated_virtual)
-		pstate->p_hasGeneratedVirtual = true;
-
 	/*
 	 * Now build an RTE and a ParseNamespaceItem.
 	 */
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 350e9e18885..87df79027d7 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -405,7 +405,6 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 
 	qry->hasTargetSRFs = false;
 	qry->hasSubLinks = pstate->p_hasSubLinks;
-	qry->hasGeneratedVirtual = pstate->p_hasGeneratedVirtual;
 
 	assign_query_collations(pstate, qry);
 
diff --git a/src/backend/parser/parse_relation.c b/src/backend/parser/parse_relation.c
index b864749f09c..7a3199ae1d3 100644
--- a/src/backend/parser/parse_relation.c
+++ b/src/backend/parser/parse_relation.c
@@ -1515,9 +1515,6 @@ addRangeTableEntry(ParseState *pstate,
 	rte->eref = makeAlias(refname, NIL);
 	buildRelationAliases(rel->rd_att, alias, rte->eref);
 
-	if (rel->rd_att->constr && rel->rd_att->constr->has_generated_virtual)
-		pstate->p_hasGeneratedVirtual = true;
-
 	/*
 	 * Set flags and initialize access permissions.
 	 *
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index f2713eaef23..0241d471571 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -44,7 +44,6 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
-#include "utils/syscache.h"
 
 
 /* We use a list of these to detect recursion in RewriteQuery */
@@ -91,8 +90,7 @@ static List *matchLocks(CmdType event, Relation relation,
 						int varno, Query *parsetree, bool *hasUpdate);
 static Query *fireRIRrules(Query *parsetree, List *activeRIRs);
 static Bitmapset *adjust_view_column_set(Bitmapset *cols, List *targetlist);
-struct expand_generated_context;
-static Query *expand_generated_columns_in_query(Query *query, struct expand_generated_context *context);
+static Node *expand_generated_columns_internal(Node *node, Relation rel, int rt_index, RangeTblEntry *rte);
 
 
 /*
@@ -2139,6 +2137,9 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 			}
 		}
 
+		/* Expand virtual generated columns of this table */
+		parsetree = (Query *) expand_generated_columns_internal((Node *) parsetree, rel, rt_index, rte);
+
 		table_close(rel, NoLock);
 	}
 
@@ -4370,84 +4371,66 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 
 
 /*
- * Virtual generated columns support
+ * Expand virtual generated columns
+ *
+ * If the table contains virtual generated columns, build a target list
+ * containing the expanded expressions and use ReplaceVarsFromTargetList() to
+ * do the replacements.
  */
-
-struct expand_generated_context
-{
-	/* list of range tables, innermost last */
-	List	   *rtables;
-
-	/* incremented for every level where it's true */
-	int			ancestor_has_virtual;
-};
-
 static Node *
-expand_generated_columns_mutator(Node *node, struct expand_generated_context *context)
+expand_generated_columns_internal(Node *node, Relation rel, int rt_index, RangeTblEntry *rte)
 {
-	if (node == NULL)
-		return NULL;
+	TupleDesc	tupdesc;
 
-	if (IsA(node, Var))
+	tupdesc = RelationGetDescr(rel);
+	if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
 	{
-		Var		   *v = (Var *) node;
-		Oid			relid;
-		AttrNumber	attnum;
-		List	   *rtable = list_nth_node(List,
-										   context->rtables,
-										   list_length(context->rtables) - v->varlevelsup - 1);
-
-		relid = rt_fetch(v->varno, rtable)->relid;
-		attnum = v->varattno;
+		List	   *tlist = NIL;
 
-		if (!relid || !attnum)
-			return node;
-
-		if (get_attgenerated(relid, attnum) == ATTRIBUTE_GENERATED_VIRTUAL)
+		for (int i = 0; i < tupdesc->natts; i++)
 		{
-			Relation	rt_entry_relation = table_open(relid, NoLock);
-			Oid			attcollid;
-
-			node = build_column_default(rt_entry_relation, attnum);
-			if (node == NULL)
-				elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
-					 attnum, RelationGetRelationName(rt_entry_relation));
+			Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
 
-			/*
-			 * If the column definition has a collation and it is different
-			 * from the collation of the generation expression, put a COLLATE
-			 * clause around the expression.
-			 */
-			attcollid = GetSysCacheOid(ATTNUM, Anum_pg_attribute_attcollation, relid, attnum, 0, 0);
-			if (attcollid && attcollid != exprCollation(node))
+			if (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
 			{
-				CollateExpr *ce = makeNode(CollateExpr);
+				Node	   *defexpr;
+				int			attnum = i + 1;
+				Oid			attcollid;
+				TargetEntry *te;
 
-				ce->arg = (Expr *) node;
-				ce->collOid = attcollid;
-				ce->location = -1;
+				defexpr = build_column_default(rel, attnum);
+				if (defexpr == NULL)
+					elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
+						 attnum, RelationGetRelationName(rel));
 
-				node = (Node *) ce;
-			}
+				/*
+				 * If the column definition has a collation and it is
+				 * different from the collation of the generation expression,
+				 * put a COLLATE clause around the expression.
+				 */
+				attcollid = attr->attcollation;
+				if (attcollid && attcollid != exprCollation(defexpr))
+				{
+					CollateExpr *ce = makeNode(CollateExpr);
 
-			IncrementVarSublevelsUp(node, v->varlevelsup, 0);
-			ChangeVarNodes(node, 1, v->varno, v->varlevelsup);
+					ce->arg = (Expr *) defexpr;
+					ce->collOid = attcollid;
+					ce->location = -1;
 
-			table_close(rt_entry_relation, NoLock);
-		}
+					defexpr = (Node *) ce;
+				}
 
-		return node;
-	}
-	else if (IsA(node, Query))
-	{
-		Query	   *query = (Query *) node;
+				ChangeVarNodes(defexpr, 1, rt_index, 0);
 
-		query = expand_generated_columns_in_query(query, context);
+				te = makeTargetEntry((Expr *) defexpr, attnum, 0, false);
+				tlist = lappend(tlist, te);
+			}
+		}
 
-		return (Node *) query;
+		node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist, REPLACEVARS_CHANGE_VARNO, rt_index, NULL);
 	}
-	else
-		return expression_tree_mutator(node, expand_generated_columns_mutator, context);
+
+	return node;
 }
 
 Node *
@@ -4458,76 +4441,19 @@ expand_generated_columns_in_expr(Node *node, Relation rel)
 	if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
 	{
 		RangeTblEntry *rte;
-		List	   *rtable;
-		struct expand_generated_context context;
 
-		/*
-		 * Make a dummy range table for a single relation.  For the benefit of
-		 * triggers, add the same entry twice, so it covers PRS2_OLD_VARNO and
-		 * PRS2_NEW_VARNO.
-		 */
 		rte = makeNode(RangeTblEntry);
 		rte->relid = RelationGetRelid(rel);
-		rtable = list_make2(rte, rte);
-		context.rtables = list_make1(rtable);
-
-		return expression_tree_mutator(node, expand_generated_columns_mutator, &context);
-	}
-	else
-		return node;
-}
-
-/*
- * Expand virtual generated columns in a Query.  We do some optimizations here
- * to avoid digging through the whole Query unless necessary.
- */
-static Query *
-expand_generated_columns_in_query(Query *query, struct expand_generated_context *context)
-{
-	context->rtables = lappend(context->rtables, query->rtable);
-	if (query->hasGeneratedVirtual)
-		context->ancestor_has_virtual++;
-
-	/*
-	 * If any table in the query has a virtual column or there is a sublink,
-	 * then we need to do the whole walk.
-	 */
-	if (query->hasGeneratedVirtual || query->hasSubLinks || context->ancestor_has_virtual)
-	{
-		query = query_tree_mutator(query,
-								   expand_generated_columns_mutator,
-								   context,
-								   QTW_DONT_COPY_QUERY);
-	}
-
-	/*
-	 * Else we only need to process subqueries.
-	 */
-	else
-	{
-		ListCell   *lc;
-
-		foreach(lc, query->rtable)
-		{
-			RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
-
-			if (rte->rtekind == RTE_SUBQUERY)
-				rte->subquery = expand_generated_columns_in_query(rte->subquery, context);
-		}
 
-		foreach(lc, query->cteList)
-		{
-			CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc);
-
-			cte->ctequery = (Node *) expand_generated_columns_in_query(castNode(Query, cte->ctequery), context);
-		}
+		/*
+		 * XXX For the benefit of triggers, make two passes, so it covers
+		 * PRS2_OLD_VARNO and PRS2_NEW_VARNO.
+		 */
+		node = expand_generated_columns_internal(node, rel, 1, rte);
+		node = expand_generated_columns_internal(node, rel, 2, rte);
 	}
 
-	if (query->hasGeneratedVirtual)
-		context->ancestor_has_virtual--;
-	context->rtables = list_truncate(context->rtables, list_length(context->rtables) - 1);
-
-	return query;
+	return node;
 }
 
 
@@ -4575,12 +4501,9 @@ QueryRewrite(Query *parsetree)
 	foreach(l, querylist)
 	{
 		Query	   *query = (Query *) lfirst(l);
-		struct expand_generated_context context = {0};
 
 		query = fireRIRrules(query, NIL);
 
-		query = expand_generated_columns_in_query(query, &context);
-
 		query->queryId = input_query_id;
 
 		results = lappend(results, query);
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 355f52199fa..0e91c084a01 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -160,8 +160,6 @@ typedef struct Query
 	bool		hasForUpdate pg_node_attr(query_jumble_ignore);
 	/* rewriter has applied some RLS policy */
 	bool		hasRowSecurity pg_node_attr(query_jumble_ignore);
-	/* some table has a virtual generated column */
-	bool		hasGeneratedVirtual pg_node_attr(query_jumble_ignore);
 	/* is a RETURN statement */
 	bool		isReturn pg_node_attr(query_jumble_ignore);
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 1b02128548b..5b781d87a9d 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -225,7 +225,6 @@ struct ParseState
 	bool		p_hasTargetSRFs;
 	bool		p_hasSubLinks;
 	bool		p_hasModifyingCTE;
-	bool		p_hasGeneratedVirtual;
 
 	Node	   *p_last_srf;		/* most recent set-returning func/op found */
 
-- 
2.46.0