diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c index 4f7f8a65ce2..62f54786be4 100644 --- a/src/backend/rewrite/rewriteHandler.c +++ b/src/backend/rewrite/rewriteHandler.c @@ -96,8 +96,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); -static Node *expand_generated_columns_internal(Node *node, Relation rel, int rt_index, - RangeTblEntry *rte, int result_relation); +static List *get_generated_columns(Relation rel, int rt_index, bool include_stored); /* @@ -640,12 +639,46 @@ 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; + List *gen_cols; + + /* + * The target list does not contain entries for generated columns + * (they are removed by rewriteTargetListIU), so we must build entries + * for them here, so that new.gen_col can be rewritten correctly. + */ + new_rel = relation_open(new_rte->relid, NoLock); + gen_cols = get_generated_columns(new_rel, new_varno, true); + relation_close(new_rel, NoLock); + + /* + * The generated column expressions refer to new.attribute, so they + * must be rewritten before they can be used as replacements. + */ + gen_cols = (List *) + ReplaceVarsFromTargetList((Node *) gen_cols, + new_varno, + 0, + new_rte, + parsetree->targetList, + sub_action->resultRelation, + (event == CMD_UPDATE) ? + REPLACEVARS_CHANGE_VARNO : + REPLACEVARS_SUBSTITUTE_NULL, + current_varno, + &sub_action->hasSubLinks); + + /* + * Now rewrite new.attribute in sub_action, using both the target list + * and the rewritten generated column expressions. + */ sub_action = (Query *) ReplaceVarsFromTargetList((Node *) sub_action, new_varno, 0, - rt_fetch(new_varno, sub_action->rtable), - parsetree->targetList, + new_rte, + list_concat(gen_cols, parsetree->targetList), sub_action->resultRelation, (event == CMD_UPDATE) ? REPLACEVARS_CHANGE_VARNO : @@ -2342,18 +2375,50 @@ CopyAndAddInvertedQual(Query *parsetree, ChangeVarNodes(new_qual, PRS2_OLD_VARNO, rt_index, 0); /* Fix references to NEW */ if (event == CMD_INSERT || event == CMD_UPDATE) + { + RangeTblEntry *rte = rt_fetch(rt_index, parsetree->rtable); + Relation rel; + List *gen_cols; + + /* + * As in rewriteRuleAction, build entries for generated columns so + * that new.gen_col in the rule qualification can be rewritten + * correctly. + */ + rel = relation_open(rte->relid, NoLock); + gen_cols = get_generated_columns(rel, PRS2_NEW_VARNO, true); + relation_close(rel, NoLock); + + /* + * The generated column expressions refer to new.attribute, so they + * must be rewritten before they can be used as replacements. + */ + gen_cols = (List *) + ReplaceVarsFromTargetList((Node *) gen_cols, + PRS2_NEW_VARNO, + 0, + rte, + parsetree->targetList, + parsetree->resultRelation, + (event == CMD_UPDATE) ? + REPLACEVARS_CHANGE_VARNO : + REPLACEVARS_SUBSTITUTE_NULL, + rt_index, + &parsetree->hasSubLinks); + new_qual = ReplaceVarsFromTargetList(new_qual, PRS2_NEW_VARNO, 0, - rt_fetch(rt_index, - parsetree->rtable), - parsetree->targetList, + rte, + list_concat(gen_cols, + parsetree->targetList), parsetree->resultRelation, (event == CMD_UPDATE) ? REPLACEVARS_CHANGE_VARNO : REPLACEVARS_SUBSTITUTE_NULL, rt_index, &parsetree->hasSubLinks); + } /* And attach the fixed qual */ AddInvertedQual(parsetree, new_qual); @@ -4431,36 +4496,31 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length, /* - * Expand virtual generated columns + * Get a table's 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 include_stored is true, both stored and virtual generated columns are + * returned. Otherwise, only virtual generated columns are returned. * - * Vars matching rt_index at the current query level are replaced by the - * virtual 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. + * Returns a list of TargetEntry, one for each generated column, containing + * the attribute numbers and generation expressions. */ -static Node * -expand_generated_columns_internal(Node *node, Relation rel, int rt_index, - RangeTblEntry *rte, int result_relation) +static List * +get_generated_columns(Relation rel, int rt_index, bool include_stored) { + List *gen_cols = NIL; 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; - for (int i = 0; i < tupdesc->natts; i++) { 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; @@ -4469,19 +4529,12 @@ expand_generated_columns_internal(Node *node, Relation rel, int rt_index, ChangeVarNodes(defexpr, 1, rt_index, 0); te = makeTargetEntry((Expr *) defexpr, i + 1, 0, false); - tlist = lappend(tlist, te); + gen_cols = lappend(gen_cols, te); } } - - Assert(list_length(tlist) > 0); - - node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, tlist, - result_relation, - REPLACEVARS_CHANGE_VARNO, rt_index, - NULL); } - return node; + return gen_cols; } /* @@ -4498,6 +4551,7 @@ expand_generated_columns_in_expr(Node *node, Relation rel, int rt_index) if (tupdesc->constr && tupdesc->constr->has_generated_virtual) { RangeTblEntry *rte; + List *vcols; rte = makeNode(RangeTblEntry); /* eref needs to be set, but the actual name doesn't matter */ @@ -4505,14 +4559,26 @@ 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); + vcols = get_generated_columns(rel, rt_index, false); + + if (vcols) + { + /* + * Passing NULL for outer_hasSubLinks is safe because generation + * expressions cannot contain SubLinks, so the replacement cannot + * introduce any. + */ + node = ReplaceVarsFromTargetList(node, rt_index, 0, rte, vcols, 0, + REPLACEVARS_CHANGE_VARNO, rt_index, + NULL); + } } return node; } /* - * Build the generation expression for the virtual generated column. + * Build the generation expression for a generated column. * * Error out if there is no generation expression found for the given column. */ @@ -4524,8 +4590,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 8b7a71d8f0c..86a04548a02 100644 --- a/src/test/regress/expected/generated_stored.out +++ b/src/test/regress/expected/generated_stored.out @@ -1527,6 +1527,40 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); c | integer | | | x | integer | | | generated always as (b * 2) stored +-- rule actions referring to generated columns: +-- 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; +UPDATE gtest_rule SET a = (SELECT max(b) FROM gtest_rule); +SELECT * FROM gtest_rule_log; + op | old_b | new_b +-----+-------+------- + INS | | 2 + UPD | 2 | 20 + UPD | 20 | 40 +(3 rows) + +DROP RULE gtest_rule_upd ON gtest_rule; +DROP RULE gtest_rule_ins ON gtest_rule; +DROP TABLE gtest_rule_log; +-- rule quals referring to generated columns: +-- NEW.b in the rule qual should reflect the generated column's new value +CREATE RULE gtest_rule_qual AS ON UPDATE TO gtest_rule WHERE NEW.b > 100 + DO INSTEAD NOTHING; +UPDATE gtest_rule SET a = 100; +SELECT * FROM gtest_rule; + a | b +----+---- + 20 | 40 +(1 row) + +DROP TABLE gtest_rule; -- 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 b7da03ce7ea..15760f4eae6 100644 --- a/src/test/regress/expected/generated_virtual.out +++ b/src/test/regress/expected/generated_virtual.out @@ -1494,6 +1494,40 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); c | integer | | | x | integer | | | generated always as (b * 2) +-- rule actions referring to generated columns: +-- 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) 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); +UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET a = (SELECT max(b) FROM gtest_rule); +SELECT * FROM gtest_rule_log; + op | old_b | new_b +-----+-------+------- + INS | | 2 + UPD | 2 | 20 + UPD | 20 | 40 +(3 rows) + +DROP RULE gtest_rule_upd ON gtest_rule; +DROP RULE gtest_rule_ins ON gtest_rule; +DROP TABLE gtest_rule_log; +-- rule quals referring to generated columns: +-- NEW.b in the rule qual should reflect the generated column's new value +CREATE RULE gtest_rule_qual AS ON UPDATE TO gtest_rule WHERE NEW.b > 100 + DO INSTEAD NOTHING; +UPDATE gtest_rule SET a = 100; +SELECT * FROM gtest_rule; + a | b +----+---- + 20 | 40 +(1 row) + +DROP TABLE gtest_rule; -- 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 2001a47bcc6..6abae289ffe 100644 --- a/src/test/regress/sql/generated_stored.sql +++ b/src/test/regress/sql/generated_stored.sql @@ -751,6 +751,29 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); \d gtest28* +-- rule actions referring to generated columns: +-- 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; +UPDATE gtest_rule SET a = (SELECT max(b) FROM gtest_rule); +SELECT * FROM gtest_rule_log; +DROP RULE gtest_rule_upd ON gtest_rule; +DROP RULE gtest_rule_ins ON gtest_rule; +DROP TABLE gtest_rule_log; + +-- rule quals referring to generated columns: +-- NEW.b in the rule qual should reflect the generated column's new value +CREATE RULE gtest_rule_qual AS ON UPDATE TO gtest_rule WHERE NEW.b > 100 + DO INSTEAD NOTHING; +UPDATE gtest_rule SET a = 100; +SELECT * FROM gtest_rule; +DROP TABLE gtest_rule; -- 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 81a98995d89..fdfc764e24a 100644 --- a/src/test/regress/sql/generated_virtual.sql +++ b/src/test/regress/sql/generated_virtual.sql @@ -802,6 +802,29 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED); \d gtest28* +-- rule actions referring to generated columns: +-- 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) 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); +UPDATE gtest_rule SET a = 10; +UPDATE gtest_rule SET a = (SELECT max(b) FROM gtest_rule); +SELECT * FROM gtest_rule_log; +DROP RULE gtest_rule_upd ON gtest_rule; +DROP RULE gtest_rule_ins ON gtest_rule; +DROP TABLE gtest_rule_log; + +-- rule quals referring to generated columns: +-- NEW.b in the rule qual should reflect the generated column's new value +CREATE RULE gtest_rule_qual AS ON UPDATE TO gtest_rule WHERE NEW.b > 100 + DO INSTEAD NOTHING; +UPDATE gtest_rule SET a = 100; +SELECT * FROM gtest_rule; +DROP TABLE gtest_rule; -- sanity check of system catalog SELECT attrelid, attname, attgenerated FROM pg_attribute WHERE attgenerated NOT IN ('', 's', 'v');