From 7cc29faf775b2368d4381856d6f853df01ab01c2 Mon Sep 17 00:00:00 2001
From: Richard Guo <guofenglinux@gmail.com>
Date: Mon, 13 Apr 2026 16:54:23 +0900
Subject: [PATCH v1] Fix incorrect NEW references to generated columns in rule
 actions

---
 src/backend/rewrite/rewriteHandler.c          | 75 ++++++++++++++-----
 .../regress/expected/generated_stored.out     | 18 +++++
 .../regress/expected/generated_virtual.out    | 18 +++++
 src/test/regress/sql/generated_stored.sql     | 13 ++++
 src/test/regress/sql/generated_virtual.sql    | 13 ++++
 5 files changed, 118 insertions(+), 19 deletions(-)

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 021c73f1b67..0ba9931b2fd 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -98,7 +98,8 @@ static List *matchLocks(CmdType event, Relation relation,
 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, int result_relation);
+											   RangeTblEntry *rte, int result_relation,
+											   bool include_stored);
 
 
 /*
@@ -642,11 +643,33 @@ rewriteRuleAction(Query *parsetree,
 	if ((event == CMD_INSERT || event == CMD_UPDATE) &&
 		sub_action->commandType != CMD_UTILITY)
 	{
+		RangeTblEntry *new_rte = rt_fetch(new_varno, sub_action->rtable);
+		Relation	new_rel;
+
+		/*
+		 * First, expand any generated column references in the NEW relation.
+		 * This must happen before ReplaceVarsFromTargetList, because the
+		 * query's target list won't contain entries for generated columns
+		 * (they are removed by rewriteTargetListIU).  Without this step,
+		 * ReplaceVarsFromTargetList would fall through to
+		 * REPLACEVARS_CHANGE_VARNO (for UPDATE) or
+		 * REPLACEVARS_SUBSTITUTE_NULL (for INSERT), producing wrong results
+		 * when the generated column depends on columns being modified.
+		 */
+		new_rel = relation_open(new_rte->relid, NoLock);
+		sub_action = (Query *)
+			expand_generated_columns_internal((Node *) sub_action,
+											  new_rel, new_varno,
+											  new_rte,
+											  sub_action->resultRelation,
+											  true);
+		relation_close(new_rel, NoLock);
+
 		sub_action = (Query *)
 			ReplaceVarsFromTargetList((Node *) sub_action,
 									  new_varno,
 									  0,
-									  rt_fetch(new_varno, sub_action->rtable),
+									  new_rte,
 									  parsetree->targetList,
 									  sub_action->resultRelation,
 									  (event == CMD_UPDATE) ?
@@ -4530,28 +4553,38 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length,
 
 
 /*
- * Expand virtual generated columns
+ * Expand generated columns
  *
- * If the table contains virtual generated columns, build a target list
- * containing the expanded expressions and use ReplaceVarsFromTargetList() to
- * do the replacements.
+ * If the table contains 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.
+ * 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.
+ *
+ * If include_stored is false, only virtual generated columns are expanded.
+ * If include_stored is true, both stored and virtual generated columns are
+ * expanded.  The latter is needed when expanding NEW references in rule
+ * actions, because rewriteTargetListIU removes both types from the target
+ * list, and both need to be expanded before ReplaceVarsFromTargetList can
+ * correctly resolve them.
  */
 static Node *
 expand_generated_columns_internal(Node *node, Relation rel, int rt_index,
-								  RangeTblEntry *rte, int result_relation)
+								  RangeTblEntry *rte, int result_relation,
+								  bool include_stored)
 {
 	TupleDesc	tupdesc;
 
 	tupdesc = RelationGetDescr(rel);
-	if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
+	if (tupdesc->constr &&
+		(tupdesc->constr->has_generated_virtual ||
+		 (include_stored && tupdesc->constr->has_generated_stored)))
 	{
 		List	   *tlist = NIL;
 
@@ -4559,7 +4592,8 @@ expand_generated_columns_internal(Node *node, Relation rel, int rt_index,
 		{
 			Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
 
-			if (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+			if (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL ||
+				(include_stored && attr->attgenerated == ATTRIBUTE_GENERATED_STORED))
 			{
 				Node	   *defexpr;
 				TargetEntry *te;
@@ -4572,12 +4606,11 @@ expand_generated_columns_internal(Node *node, Relation rel, int rt_index,
 			}
 		}
 
-		Assert(list_length(tlist) > 0);
-
-		node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist,
-										 result_relation,
-										 REPLACEVARS_CHANGE_VARNO, rt_index,
-										 NULL);
+		if (tlist != NIL)
+			node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist,
+											 result_relation,
+											 REPLACEVARS_CHANGE_VARNO, rt_index,
+											 NULL);
 	}
 
 	return node;
@@ -4604,7 +4637,8 @@ expand_generated_columns_in_expr(Node *node, Relation rel, int rt_index)
 		rte->rtekind = RTE_RELATION;
 		rte->relid = RelationGetRelid(rel);
 
-		node = expand_generated_columns_internal(node, rel, rt_index, rte, 0);
+		node = expand_generated_columns_internal(node, rel, rt_index, rte, 0,
+												 false);
 	}
 
 	return node;
@@ -4623,8 +4657,11 @@ build_generation_expression(Relation rel, int attrno)
 	Node	   *defexpr;
 	Oid			attcollid;
 
-	Assert(rd_att->constr && rd_att->constr->has_generated_virtual);
-	Assert(att_tup->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL);
+	Assert(rd_att->constr &&
+		   (rd_att->constr->has_generated_virtual ||
+			rd_att->constr->has_generated_stored));
+	Assert(att_tup->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL ||
+		   att_tup->attgenerated == ATTRIBUTE_GENERATED_STORED);
 
 	defexpr = build_column_default(rel, attrno);
 	if (defexpr == NULL)
diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out
index 43cddeac373..363bd7ccad5 100644
--- a/src/test/regress/expected/generated_stored.out
+++ b/src/test/regress/expected/generated_stored.out
@@ -1598,6 +1598,24 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
  c      | integer |           |          | 
  x      | integer |           |          | generated always as (b * 2) stored
 
+-- rules
+-- NEW.b in a rule action should reflect the generated column's new value
+CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
+CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
+CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
+INSERT INTO gtest_rule (a) VALUES (1);
+UPDATE gtest_rule SET a = 10;
+SELECT * FROM gtest_rule_log;
+ op  | old_b | new_b 
+-----+-------+-------
+ INS |       |     2
+ UPD |     2 |    20
+(2 rows)
+
+DROP TABLE gtest_rule, gtest_rule_log;
 -- sanity check of system catalog
 SELECT attrelid, attname, attgenerated FROM pg_attribute WHERE attgenerated NOT IN ('', 's', 'v');
  attrelid | attname | attgenerated 
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index 234061fa1f7..5bbd2f684ce 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1520,6 +1520,24 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
  c      | integer |           |          | 
  x      | integer |           |          | generated always as (b * 2)
 
+-- rules
+-- NEW.b in a rule action should reflect the generated column's new value
+CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
+CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
+CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
+INSERT INTO gtest_rule (a) VALUES (1);
+UPDATE gtest_rule SET a = 10;
+SELECT * FROM gtest_rule_log;
+ op  | old_b | new_b 
+-----+-------+-------
+ INS |       |     2
+ UPD |     2 |    20
+(2 rows)
+
+DROP TABLE gtest_rule, gtest_rule_log;
 -- sanity check of system catalog
 SELECT attrelid, attname, attgenerated FROM pg_attribute WHERE attgenerated NOT IN ('', 's', 'v');
  attrelid | attname | attgenerated 
diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql
index 280021d79b7..3265a1e88c6 100644
--- a/src/test/regress/sql/generated_stored.sql
+++ b/src/test/regress/sql/generated_stored.sql
@@ -798,6 +798,19 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 \d gtest28*
 
+-- rules
+-- NEW.b in a rule action should reflect the generated column's new value
+CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
+CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
+CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
+INSERT INTO gtest_rule (a) VALUES (1);
+UPDATE gtest_rule SET a = 10;
+SELECT * FROM gtest_rule_log;
+DROP TABLE gtest_rule, gtest_rule_log;
+
 
 -- sanity check of system catalog
 SELECT attrelid, attname, attgenerated FROM pg_attribute WHERE attgenerated NOT IN ('', 's', 'v');
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index 4d9ad3c5dca..cd1cd938975 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -811,6 +811,19 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 \d gtest28*
 
+-- rules
+-- NEW.b in a rule action should reflect the generated column's new value
+CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
+CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
+CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
+  DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
+INSERT INTO gtest_rule (a) VALUES (1);
+UPDATE gtest_rule SET a = 10;
+SELECT * FROM gtest_rule_log;
+DROP TABLE gtest_rule, gtest_rule_log;
+
 
 -- sanity check of system catalog
 SELECT attrelid, attname, attgenerated FROM pg_attribute WHERE attgenerated NOT IN ('', 's', 'v');
-- 
2.39.5 (Apple Git-154)

