Fix incorrect NEW references to generated columns in rule rewriting

When a rule action or rule qualification references NEW.col where col
is a generated column (stored or virtual), the rewriter produces
incorrect results.

rewriteTargetListIU removes generated columns from the query's target
list, since stored generated columns are recomputed by the executor
and virtual ones store nothing.  However, ReplaceVarsFromTargetList
then cannot find these columns when resolving NEW references during
rule rewriting.  For UPDATE, the REPLACEVARS_CHANGE_VARNO fallback
redirects NEW.col to the original target relation, making it read the
pre-update value (same as OLD.col).  For INSERT,
REPLACEVARS_SUBSTITUTE_NULL replaces it with NULL.  Both are wrong
when the generated column depends on columns being modified.

Fix by building target list entries for generated columns from their
generation expressions, pre-resolving the NEW.attribute references
within those expressions against the query's targetlist, and passing
them together with the query's targetlist to ReplaceVarsFromTargetList.

Back-patch to all supported branches.  Virtual generated columns were
added in v18, so the back-patches in pre-v18 branches only handle
stored generated columns.

Reported-by: SATYANARAYANA NARLAPURAM <satyanarlapuram@gmail.com>
Author: Richard Guo <guofenglinux@gmail.com>
Author: Dean Rasheed <dean.a.rasheed@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/CAHg+QDexGTmCZzx=73gXkY2ZADS6LRhpnU+-8Y_QmrdTS6yUhA@mail.gmail.com
Backpatch-through: 14

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 4f7f8a65ce256065ebe0f477b738d2015aed2d5f..62f54786be475808e6b1ea6b1e899bf4d2390e41 100644 (file)
--- 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
- *
- * If the table contains virtual generated columns, build a target list
- * containing the expanded expressions and use ReplaceVarsFromTargetList() to
- * do the replacements.
+ * Get a table's generated columns
  *
- * Vars matching rt_index at the current query level are replaced by the
- * virtual generated column expressions from rel, if there are any.
+ * If include_stored is true, both stored and virtual generated columns are
+ * returned.  Otherwise, only virtual generated columns are returned.
  *
- * 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 8b7a71d8f0c482db54797244bfeaf3d741814276..86a04548a02eb3fd19eedb16c50d976395d974c8 100644 (file)
--- 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 b7da03ce7ea211929bdb876a0e9d8641839eb52e..15760f4eae60748289cad2e0aca8810a326d4c84 100644 (file)
--- 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 2001a47bcc6a3b6e0e5652f5ca18652a06f09211..6abae289ffed7d60e92820b65467866a29a35b5b 100644 (file)
--- 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 81a98995d8949bea8b86dad052feb8dcba3bc37a..fdfc764e24acd7af57f2574d98c90ab7dcbd4a31 100644 (file)
--- 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');