Thanks Jian for your review.
On Tue, Apr 22, 2025 at 12:32 PM jian he <jian.universal...@gmail.com>
wrote:
> On Tue, Apr 22, 2025 at 11:45 AM Ashutosh Bapat
> <ashutosh.bapat....@gmail.com> wrote:
> >
> > I have included your original tests, but ended up rewriting code. Please
> let me know what do you think.
> >
>
> + if (attno < 0)
> + ereport(ERROR,
> + (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
> + errmsg("partition key expressions cannot contain system column
> references")));
> +
> + /*
> + * We do not check dropped columns explicitly since they will
> + * be eliminated when expanding the the whole row expression
> + * anyway.
> + */
> typo: "the the".
> I am confused by the above comments.
> ComputePartitionAttrs only called in DefineRelation.
> DefineRelation will only CREATE a table, there will be no dropped
> column via DefineRelation.
>
You are right! Fixed.
>
>
> + /*
> + * transformPartitionSpec() should have already rejected
> + * subqueries, aggregates, window functions, and SRFs, based
> + * on the EXPR_KIND_ for partition expressions.
> + */
> "EXPR_KIND_" can change to EXPR_KIND_PARTITION_EXPRESSION
> ?
>
That's an existing comment which appears to be moved in diff but it's
actually untouched by the patch. Maybe you are right, IDK, but since it's
not related to the fix, let's leave it untouched. I did wonder whether that
comment has any relation to the pull_varattnos call which has been moved
earlier since pull_varattnos doesn't expect any Query node in the tree. But
pondering more, I think the comment is related to the number of rows
participating in the partition key computation - there should be exactly
one key for one tuple. Having SRFs, subqueries in part expression means a
possibility of producing more than one set of partition keys, aggregates
and window functions OTOH will produce one partition key for more than one
row - all of them breaking 1:1 mapping between a tuple and a partition key.
Hence I think the comment is at the right place.
PFA revised patch.
--
Best Wishes,
Ashutosh Bapat
From 5467147f6b37898263c78ce64b086072c76caaad Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat....@gmail.com>
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. "(<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.
Discussion: https://postgr.es/m/CACJufxF=wdgthxsaqr9thyusfx_1_t9e6n8te3b8eqxcvov...@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
