diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c index 0ba9931b2fd..e7cf63cb47a 100644 --- a/src/backend/rewrite/rewriteHandler.c +++ b/src/backend/rewrite/rewriteHandler.c @@ -332,6 +332,57 @@ acquireLocksOnSubLinks(Node *node, acquireLocksOnSubLinks_context *context) return expression_tree_walker(node, acquireLocksOnSubLinks, context); } +static List * +add_generated_columns_to_targetlist(List *targetList, + CmdType commandType, + Relation target_relation, + RangeTblEntry *target_rte, + int target_rt_index, + int result_relation, + int nomatch_varno) +{ + TupleDesc tupdesc = RelationGetDescr(target_relation); + List *result; + + if (!(tupdesc->constr && + (tupdesc->constr->has_generated_stored || + tupdesc->constr->has_generated_virtual))) + return targetList; + + result = list_copy(targetList); + + for (int i = 0; i < tupdesc->natts; i++) + { + Form_pg_attribute attr = TupleDescAttr(tupdesc, i); + Node *expr; + TargetEntry *te; + + if (!attr->attgenerated) + continue; + + expr = build_generation_expression(target_relation, i + 1); + ChangeVarNodes(expr, 1, target_rt_index, 0); + expr = ReplaceVarsFromTargetList(expr, + target_rt_index, + 0, + target_rte, + targetList, + result_relation, + (commandType == CMD_UPDATE) ? + REPLACEVARS_CHANGE_VARNO : + REPLACEVARS_SUBSTITUTE_NULL, + nomatch_varno, + NULL); + + te = makeTargetEntry((Expr *) expr, + i + 1, + pstrdup(NameStr(attr->attname)), + false); + result = lappend(result, te); + } + + return result; +} /* * rewriteRuleAction - @@ -645,24 +696,16 @@ rewriteRuleAction(Query *parsetree, { RangeTblEntry *new_rte = rt_fetch(new_varno, sub_action->rtable); Relation new_rel; + List *new_targetlist; - /* - * 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); + new_targetlist = add_generated_columns_to_targetlist(parsetree->targetList, + event, + new_rel, + new_rte, + new_varno, + sub_action->resultRelation, + current_varno); relation_close(new_rel, NoLock); sub_action = (Query *) @@ -670,7 +713,7 @@ rewriteRuleAction(Query *parsetree, new_varno, 0, new_rte, - parsetree->targetList, + new_targetlist, sub_action->resultRelation, (event == CMD_UPDATE) ? REPLACEVARS_CHANGE_VARNO : diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out index 363bd7ccad5..be421b738bc 100644 --- a/src/test/regress/expected/generated_stored.out +++ b/src/test/regress/expected/generated_stored.out @@ -1600,20 +1600,22 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); -- 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 (a int, c int, b int GENERATED ALWAYS AS (a + c) 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); +INSERT INTO gtest_rule (a, c) VALUES (1, 5); UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET c = 7; SELECT * FROM gtest_rule_log; op | old_b | new_b -----+-------+------- - INS | | 2 - UPD | 2 | 20 -(2 rows) + INS | | 6 + UPD | 6 | 15 + UPD | 15 | 17 +(3 rows) DROP TABLE gtest_rule, gtest_rule_log; -- sanity check of system catalog diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out index 5bbd2f684ce..35b54f98ca1 100644 --- a/src/test/regress/expected/generated_virtual.out +++ b/src/test/regress/expected/generated_virtual.out @@ -1522,20 +1522,22 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); -- 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 (a int, c int, b int GENERATED ALWAYS AS (a + c) VIRTUAL); 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); +INSERT INTO gtest_rule (a, c) VALUES (1, 5); UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET c = 7; SELECT * FROM gtest_rule_log; op | old_b | new_b -----+-------+------- - INS | | 2 - UPD | 2 | 20 -(2 rows) + INS | | 6 + UPD | 6 | 15 + UPD | 15 | 17 +(3 rows) DROP TABLE gtest_rule, gtest_rule_log; -- sanity check of system catalog diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql index 3265a1e88c6..bda1fe760e8 100644 --- a/src/test/regress/sql/generated_stored.sql +++ b/src/test/regress/sql/generated_stored.sql @@ -800,14 +800,15 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); -- 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 (a int, c int, b int GENERATED ALWAYS AS (a + c) 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); +INSERT INTO gtest_rule (a, c) VALUES (1, 5); UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET c = 7; SELECT * FROM gtest_rule_log; DROP TABLE gtest_rule, gtest_rule_log; diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql index cd1cd938975..b7c4c711e61 100644 --- a/src/test/regress/sql/generated_virtual.sql +++ b/src/test/regress/sql/generated_virtual.sql @@ -813,14 +813,15 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); -- 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 (a int, c int, b int GENERATED ALWAYS AS (a + c) VIRTUAL); 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); +INSERT INTO gtest_rule (a, c) VALUES (1, 5); UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET c = 7; SELECT * FROM gtest_rule_log; DROP TABLE gtest_rule, gtest_rule_log;