Skip to content

Commit 040cc5f

Browse files
peterejianhe-funashutosh-bapat
committed
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
1 parent a95e3d8 commit 040cc5f

File tree

5 files changed

+86
-40
lines changed

5 files changed

+86
-40
lines changed

‎src/backend/commands/tablecmds.c‎

Lines changed: 50 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -19835,6 +19835,8 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
1983519835
/* Expression */
1983619836
Node *expr = pelem->expr;
1983719837
char partattname[16];
19838+
Bitmapset *expr_attrs = NULL;
19839+
int i;
1983819840

1983919841
Assert(expr != NULL);
1984019842
atttype = exprType(expr);
@@ -19858,43 +19860,36 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
1985819860
while (IsA(expr, CollateExpr))
1985919861
expr = (Node *) ((CollateExpr *) expr)->arg;
1986019862

19861-
if (IsA(expr, Var) &&
19862-
((Var *) expr)->varattno > 0)
19863+
/*
19864+
* Examine all the columns in the partition key expression. When
19865+
* the whole-row reference is present, examine all the columns of
19866+
* the partitioned table.
19867+
*/
19868+
pull_varattnos(expr, 1, &expr_attrs);
19869+
if (bms_is_member(0 - FirstLowInvalidHeapAttributeNumber, expr_attrs))
1986319870
{
19864-
/*
19865-
* User wrote "(column)" or "(column COLLATE something)".
19866-
* Treat it like simple attribute anyway.
19867-
*/
19868-
partattrs[attn] = ((Var *) expr)->varattno;
19871+
expr_attrs = bms_add_range(expr_attrs,
19872+
1 - FirstLowInvalidHeapAttributeNumber,
19873+
RelationGetNumberOfAttributes(rel) - FirstLowInvalidHeapAttributeNumber);
19874+
expr_attrs = bms_del_member(expr_attrs, 0 - FirstLowInvalidHeapAttributeNumber);
1986919875
}
19870-
else
19871-
{
19872-
Bitmapset *expr_attrs = NULL;
19873-
int i;
1987419876

19875-
partattrs[attn] = 0; /* marks the column as expression */
19876-
*partexprs = lappend(*partexprs, expr);
19877+
i = -1;
19878+
while ((i = bms_next_member(expr_attrs, i)) >= 0)
19879+
{
19880+
AttrNumber attno = i + FirstLowInvalidHeapAttributeNumber;
1987719881

19878-
/*
19879-
* transformPartitionSpec() should have already rejected
19880-
* subqueries, aggregates, window functions, and SRFs, based
19881-
* on the EXPR_KIND_ for partition expressions.
19882-
*/
19882+
Assert(attno != 0);
1988319883

1988419884
/*
1988519885
* Cannot allow system column references, since that would
1988619886
* make partition routing impossible: their values won't be
1988719887
* known yet when we need to do that.
1988819888
*/
19889-
pull_varattnos(expr, 1, &expr_attrs);
19890-
for (i = FirstLowInvalidHeapAttributeNumber; i < 0; i++)
19891-
{
19892-
if (bms_is_member(i - FirstLowInvalidHeapAttributeNumber,
19893-
expr_attrs))
19894-
ereport(ERROR,
19895-
(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
19896-
errmsg("partition key expressions cannot contain system column references")));
19897-
}
19889+
if (attno < 0)
19890+
ereport(ERROR,
19891+
(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
19892+
errmsg("partition key expressions cannot contain system column references")));
1989819893

1989919894
/*
1990019895
* Stored generated columns cannot work: They are computed
@@ -19904,20 +19899,35 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
1990419899
* SET EXPRESSION would need to check whether the column is
1990519900
* used in partition keys). Seems safer to prohibit for now.
1990619901
*/
19907-
i = -1;
19908-
while ((i = bms_next_member(expr_attrs, i)) >= 0)
19909-
{
19910-
AttrNumber attno = i + FirstLowInvalidHeapAttributeNumber;
19902+
if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
19903+
ereport(ERROR,
19904+
(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
19905+
errmsg("cannot use generated column in partition key"),
19906+
errdetail("Column \"%s\" is a generated column.",
19907+
get_attname(RelationGetRelid(rel), attno, false)),
19908+
parser_errposition(pstate, pelem->location)));
19909+
}
1991119910

19912-
if (attno > 0 &&
19913-
TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
19914-
ereport(ERROR,
19915-
(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
19916-
errmsg("cannot use generated column in partition key"),
19917-
errdetail("Column \"%s\" is a generated column.",
19918-
get_attname(RelationGetRelid(rel), attno, false)),
19919-
parser_errposition(pstate, pelem->location)));
19920-
}
19911+
if (IsA(expr, Var) &&
19912+
((Var *) expr)->varattno > 0)
19913+
{
19914+
19915+
/*
19916+
* User wrote "(column)" or "(column COLLATE something)".
19917+
* Treat it like simple attribute anyway.
19918+
*/
19919+
partattrs[attn] = ((Var *) expr)->varattno;
19920+
}
19921+
else
19922+
{
19923+
partattrs[attn] = 0; /* marks the column as expression */
19924+
*partexprs = lappend(*partexprs, expr);
19925+
19926+
/*
19927+
* transformPartitionSpec() should have already rejected
19928+
* subqueries, aggregates, window functions, and SRFs, based
19929+
* on the EXPR_KIND_ for partition expressions.
19930+
*/
1992119931

1992219932
/*
1992319933
* Preprocess the expression before checking for mutability.

‎src/test/regress/expected/generated_stored.out‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1074,11 +1074,26 @@ ERROR: cannot use generated column in partition key
10741074
LINE 1: ...ENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
10751075
^
10761076
DETAIL: Column "f3" is a generated column.
1077+
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
1078+
ERROR: cannot use generated column in partition key
1079+
LINE 1: ...ERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
1080+
^
1081+
DETAIL: Column "f3" is a generated column.
10771082
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
10781083
ERROR: cannot use generated column in partition key
10791084
LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
10801085
^
10811086
DETAIL: Column "f3" is a generated column.
1087+
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));
1088+
ERROR: cannot use generated column in partition key
1089+
LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_par...
1090+
^
1091+
DETAIL: Column "f3" is a generated column.
1092+
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));
1093+
ERROR: cannot use generated column in partition key
1094+
LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_par...
1095+
^
1096+
DETAIL: Column "f3" is a generated column.
10821097
-- ALTER TABLE ... ADD COLUMN
10831098
CREATE TABLE gtest25 (a int PRIMARY KEY);
10841099
INSERT INTO gtest25 VALUES (3), (4);

‎src/test/regress/expected/generated_virtual.out‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1036,11 +1036,26 @@ ERROR: cannot use generated column in partition key
10361036
LINE 1: ...NERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
10371037
^
10381038
DETAIL: Column "f3" is a generated column.
1039+
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
1040+
ERROR: cannot use generated column in partition key
1041+
LINE 1: ...RATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
1042+
^
1043+
DETAIL: Column "f3" is a generated column.
10391044
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
10401045
ERROR: cannot use generated column in partition key
10411046
LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
10421047
^
10431048
DETAIL: Column "f3" is a generated column.
1049+
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));
1050+
ERROR: cannot use generated column in partition key
1051+
LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
1052+
^
1053+
DETAIL: Column "f3" is a generated column.
1054+
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));
1055+
ERROR: cannot use generated column in partition key
1056+
LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
1057+
^
1058+
DETAIL: Column "f3" is a generated column.
10441059
-- ALTER TABLE ... ADD COLUMN
10451060
CREATE TABLE gtest25 (a int PRIMARY KEY);
10461061
INSERT INTO gtest25 VALUES (3), (4);

‎src/test/regress/sql/generated_stored.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,7 +500,10 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
500500

501501
-- generated columns in partition key (not allowed)
502502
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
503+
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
503504
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
505+
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));
506+
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));
504507

505508
-- ALTER TABLE ... ADD COLUMN
506509
CREATE TABLE gtest25 (a int PRIMARY KEY);

‎src/test/regress/sql/generated_virtual.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -543,7 +543,10 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
543543

544544
-- generated columns in partition key (not allowed)
545545
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
546+
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
546547
CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
548+
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));
549+
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));
547550

548551
-- ALTER TABLE ... ADD COLUMN
549552
CREATE TABLE gtest25 (a int PRIMARY KEY);

0 commit comments

Comments
 (0)