From 5467147f6b37898263c78ce64b086072c76caaad Mon Sep 17 00:00:00 2001 From: Ashutosh Bapat Date: Mon, 21 Apr 2025 18:06:58 +0530 Subject: [PATCH] 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. "()" 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. Discussion: https://postgr.es/m/CACJufxF=WDGthXSAQr9thYUsfx_1_t9E6N8tE3B8EqXcVoVfQw@mail.gmail.com --- src/backend/commands/tablecmds.c | 90 ++++++++++--------- .../regress/expected/generated_stored.out | 15 ++++ .../regress/expected/generated_virtual.out | 15 ++++ src/test/regress/sql/generated_stored.sql | 3 + src/test/regress/sql/generated_virtual.sql | 3 + 5 files changed, 86 insertions(+), 40 deletions(-) diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c index 265b1c397fb..f7b7162a99c 100644 --- a/src/backend/commands/tablecmds.c +++ b/src/backend/commands/tablecmds.c @@ -19768,6 +19768,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); @@ -19791,43 +19793,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 @@ -19837,20 +19832,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 16de30ab191..eb981b0689b 100644 --- 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 6300e7c1d96..75a1730e4d2 100644 --- a/src/test/regress/expected/generated_virtual.out +++ b/src/test/regress/expected/generated_virtual.out @@ -1027,11 +1027,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 4ec155f2da9..296429fbb23 100644 --- 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 b4eedeee2fb..033a251fb20 100644 --- a/src/test/regress/sql/generated_virtual.sql +++ b/src/test/regress/sql/generated_virtual.sql @@ -534,7 +534,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); base-commit: e29df428a1dca4112aad640c889a9a54642759c9 -- 2.34.1