From 7f239c72a897f7e0f884f92ea2f92efde92c60f8 Mon Sep 17 00:00:00 2001 From: Peter Eisentraut Date: Tue, 4 Nov 2025 14:31:57 +0100 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. Co-authored-by: jian he Co-authored-by: Ashutosh Bapat Reviewed-by: Fujii Masao Discussion: https://www.postgresql.org/message-id/flat/CACJufxF%3DWDGthXSAQr9thYUsfx_1_t9E6N8tE3B8EqXcVoVfQw%40mail.gmail.com --- src/backend/commands/tablecmds.c | 86 ++++++++++++++----------- src/test/regress/expected/generated.out | 15 +++++ src/test/regress/sql/generated.sql | 3 + 3 files changed, 66 insertions(+), 38 deletions(-) diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c index 6b33b46d518..232e91c90fd 100644 --- a/src/backend/commands/tablecmds.c +++ b/src/backend/commands/tablecmds.c @@ -16548,6 +16548,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); @@ -16571,9 +16573,55 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu while (IsA(expr, CollateExpr)) expr = (Node *) ((CollateExpr *) expr)->arg; + /* + * 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)) + { + expr_attrs = bms_add_range(expr_attrs, + 1 - FirstLowInvalidHeapAttributeNumber, + RelationGetNumberOfAttributes(rel) - FirstLowInvalidHeapAttributeNumber); + expr_attrs = bms_del_member(expr_attrs, 0 - FirstLowInvalidHeapAttributeNumber); + } + + i = -1; + while ((i = bms_next_member(expr_attrs, i)) >= 0) + { + AttrNumber attno = i + FirstLowInvalidHeapAttributeNumber; + + 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. + */ + if (attno < 0) + ereport(ERROR, + (errcode(ERRCODE_INVALID_OBJECT_DEFINITION), + errmsg("partition key expressions cannot contain system column references"))); + + /* + * Generated columns cannot work: They are computed after + * BEFORE triggers, but partition routing is done before all + * triggers. + */ + 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 (IsA(expr, Var) && ((Var *) expr)->varattno > 0) { + /* * User wrote "(column)" or "(column COLLATE something)". * Treat it like simple attribute anyway. @@ -16582,9 +16630,6 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu } else { - Bitmapset *expr_attrs = NULL; - int i; - partattrs[attn] = 0; /* marks the column as expression */ *partexprs = lappend(*partexprs, expr); @@ -16594,41 +16639,6 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu * on the EXPR_KIND_ for partition expressions. */ - /* - * 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"))); - } - - /* - * Generated columns cannot work: They are computed after - * BEFORE triggers, but partition routing is done before all - * triggers. - */ - i = -1; - while ((i = bms_next_member(expr_attrs, i)) >= 0) - { - AttrNumber attno = i + FirstLowInvalidHeapAttributeNumber; - - 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))); - } - /* * Preprocess the expression before checking for mutability. * This is essential for the reasons described in diff --git a/src/test/regress/expected/generated.out b/src/test/regress/expected/generated.out index a63a2d7c0ab..157fee9f3de 100644 --- a/src/test/regress/expected/generated.out +++ b/src/test/regress/expected/generated.out @@ -693,11 +693,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_parent (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_parent (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_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_parent)); +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_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_parent 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/sql/generated.sql b/src/test/regress/sql/generated.sql index 7dcf07b38c2..533c98a003a 100644 --- a/src/test/regress/sql/generated.sql +++ b/src/test/regress/sql/generated.sql @@ -370,7 +370,10 @@ DROP TABLE gtest_parent; -- generated columns in partition key (not allowed) CREATE TABLE gtest_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3); +CREATE TABLE gtest_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3)); CREATE TABLE gtest_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3)); +CREATE TABLE gtest_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_parent)); +CREATE TABLE gtest_parent (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_parent is not null)); -- ALTER TABLE ... ADD COLUMN CREATE TABLE gtest25 (a int PRIMARY KEY); -- 2.47.3