Tighten check for generated column in partition key expression

A generated column may end up being part of the partition key
expression, if it's specified as an expression e.g. "(<generated
column name>)" or if the partition key expression contains a whole-row
reference, even though we do not allow a generated column to be part
of partition key expression.  Fix this hole.

Co-authored-by: jian he <jian.universality@gmail.com>
Co-authored-by: Ashutosh Bapat <ashutosh.bapat.oss@gmail.com>
Reviewed-by: Fujii Masao <masao.fujii@oss.nttdata.com>
Discussion: https://www.postgresql.org/message-id/flat/CACJufxF%3DWDGthXSAQr9thYUsfx_1_t9E6N8tE3B8EqXcVoVfQw%40mail.gmail.com

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 34f9e342362cea67cf3fbfb16b0050199cef5825..3aac459e483db9240dcbdf4a9d6b4720c14c0d66 100644 (file)
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -19835,6 +19835,8 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
                        /* Expression */
                        Node       *expr = pelem->expr;
                        char            partattname[16];
+                       Bitmapset  *expr_attrs = NULL;
+                       int                     i;
 
                        Assert(expr != NULL);
                        atttype = exprType(expr);
@@ -19858,43 +19860,36 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
                        while (IsA(expr, CollateExpr))
                                expr = (Node *) ((CollateExpr *) expr)->arg;
 
-                       if (IsA(expr, Var) &&
-                               ((Var *) expr)->varattno > 0)
+                       /*
+                        * Examine all the columns in the partition key expression. When
+                        * the whole-row reference is present, examine all the columns of
+                        * the partitioned table.
+                        */
+                       pull_varattnos(expr, 1, &expr_attrs);
+                       if (bms_is_member(0 - FirstLowInvalidHeapAttributeNumber, expr_attrs))
                        {
-                               /*
-                                * User wrote "(column)" or "(column COLLATE something)".
-                                * Treat it like simple attribute anyway.
-                                */
-                               partattrs[attn] = ((Var *) expr)->varattno;
+                               expr_attrs = bms_add_range(expr_attrs,
+                                                                                  1 - FirstLowInvalidHeapAttributeNumber,
+                                                                                  RelationGetNumberOfAttributes(rel) - FirstLowInvalidHeapAttributeNumber);
+                               expr_attrs = bms_del_member(expr_attrs, 0 - FirstLowInvalidHeapAttributeNumber);
                        }
-                       else
-                       {
-                               Bitmapset  *expr_attrs = NULL;
-                               int                     i;
 
-                               partattrs[attn] = 0;    /* marks the column as expression */
-                               *partexprs = lappend(*partexprs, expr);
+                       i = -1;
+                       while ((i = bms_next_member(expr_attrs, i)) >= 0)
+                       {
+                               AttrNumber      attno = i + FirstLowInvalidHeapAttributeNumber;
 
-                               /*
-                                * transformPartitionSpec() should have already rejected
-                                * subqueries, aggregates, window functions, and SRFs, based
-                                * on the EXPR_KIND_ for partition expressions.
-                                */
+                               Assert(attno != 0);
 
                                /*
                                 * Cannot allow system column references, since that would
                                 * make partition routing impossible: their values won't be
                                 * known yet when we need to do that.
                                 */
-                               pull_varattnos(expr, 1, &expr_attrs);
-                               for (i = FirstLowInvalidHeapAttributeNumber; i < 0; i++)
-                               {
-                                       if (bms_is_member(i - FirstLowInvalidHeapAttributeNumber,
-                                                                         expr_attrs))
-                                               ereport(ERROR,
-                                                               (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-                                                                errmsg("partition key expressions cannot contain system column references")));
-                               }
+                               if (attno < 0)
+                                       ereport(ERROR,
+                                                       (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+                                                        errmsg("partition key expressions cannot contain system column references")));
 
                                /*
                                 * Stored generated columns cannot work: They are computed
@@ -19904,20 +19899,35 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
                                 * SET EXPRESSION would need to check whether the column is
                                 * used in partition keys).  Seems safer to prohibit for now.
                                 */
-                               i = -1;
-                               while ((i = bms_next_member(expr_attrs, i)) >= 0)
-                               {
-                                       AttrNumber      attno = i + FirstLowInvalidHeapAttributeNumber;
+                               if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
+                                       ereport(ERROR,
+                                                       (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+                                                        errmsg("cannot use generated column in partition key"),
+                                                        errdetail("Column \"%s\" is a generated column.",
+                                                                          get_attname(RelationGetRelid(rel), attno, false)),
+                                                        parser_errposition(pstate, pelem->location)));
+                       }
 
-                                       if (attno > 0 &&
-                                               TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
-                                               ereport(ERROR,
-                                                               (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-                                                                errmsg("cannot use generated column in partition key"),
-                                                                errdetail("Column \"%s\" is a generated column.",
-                                                                                  get_attname(RelationGetRelid(rel), attno, false)),
-                                                                parser_errposition(pstate, pelem->location)));
-                               }
+                       if (IsA(expr, Var) &&
+                               ((Var *) expr)->varattno > 0)
+                       {
+
+                               /*
+                                * User wrote "(column)" or "(column COLLATE something)".
+                                * Treat it like simple attribute anyway.
+                                */
+                               partattrs[attn] = ((Var *) expr)->varattno;
+                       }
+                       else
+                       {
+                               partattrs[attn] = 0;    /* marks the column as expression */
+                               *partexprs = lappend(*partexprs, expr);
+
+                               /*
+                                * transformPartitionSpec() should have already rejected
+                                * subqueries, aggregates, window functions, and SRFs, based
+                                * on the EXPR_KIND_ for partition expressions.
+                                */
 
                                /*
                                 * Preprocess the expression before checking for mutability.

diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out
index adac2cedfb2a3383be828bd56b81101d98c6ccd5..b3710a49de628f0a22caa2bcc6906a174b3acdc8 100644 (file)
--- a/src/test/regress/expected/generated_stored.out
+++ b/src/test/regress/expected/generated_stored.out
@@ -1074,11 +1074,26 @@ ERROR:  cannot use generated column in partition key
 LINE 1: ...ENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
                                                                    ^
 DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...ERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+                                                                 ^
+DETAIL:  Column "f3" is a generated column.
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
 ERROR:  cannot use generated column in partition key
 LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
                                                              ^
 DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_part_key));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_par...
+                                                             ^
+DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_part_key is not null));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_par...
+                                                             ^
+DETAIL:  Column "f3" is a generated column.
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);

diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index c861bd36c5a293f4c981a2238360ba1f6f3007d9..4ec3d330017503a0a99a469210a6738750e3f392 100644 (file)
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1036,11 +1036,26 @@ ERROR:  cannot use generated column in partition key
 LINE 1: ...NERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
                                                                    ^
 DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...RATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
+                                                                 ^
+DETAIL:  Column "f3" is a generated column.
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
 ERROR:  cannot use generated column in partition key
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
                                                              ^
 DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
+                                                             ^
+DETAIL:  Column "f3" is a generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key is not null));
+ERROR:  cannot use generated column in partition key
+LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
+                                                             ^
+DETAIL:  Column "f3" is a generated column.
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);

diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql
index f56fde8d4e5d0a718729f4a024db7d9019333d3c..99ea0105685cff2f011249c0fc04bdbc9a9a17a8 100644 (file)
--- a/src/test/regress/sql/generated_stored.sql
+++ b/src/test/regress/sql/generated_stored.sql
@@ -500,7 +500,10 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_part_key));
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_part_key is not null));
 
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);

diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index adfe88d74ae98382794593b0b47203a4c9341375..992c0cdae65ec0ba8af45ba963e1df4415753e24 100644 (file)
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -543,7 +543,10 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key));
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key is not null));
 
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);