From 9706a4ebbcfe186b2d6c9fca37187518e4f917c1 Mon Sep 17 00:00:00 2001 From: jian he Date: Tue, 4 Nov 2025 19:41:56 +0800 Subject: [PATCH v2 3/3] fix COPY WHERE clause stored generated column references discussion: https://postgr.es/m/CACJufxHb8YPQ095R_pYDr77W9XKNaXg5Rzy-WP525mkq+hRM3g@mail.gmail.com --- src/backend/commands/copyfrom.c | 56 ++++++++++++++++++- .../regress/expected/generated_stored.out | 14 ++++- src/test/regress/sql/generated_stored.sql | 13 +++++ 3 files changed, 81 insertions(+), 2 deletions(-) diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c index 12781963b4f..7f1a4728f93 100644 --- a/src/backend/commands/copyfrom.c +++ b/src/backend/commands/copyfrom.c @@ -785,6 +785,7 @@ CopyFrom(CopyFromState cstate) ModifyTableState *mtstate; ExprContext *econtext; TupleTableSlot *singleslot = NULL; + TupleTableSlot *tmpslot = NULL; MemoryContext oldcontext = CurrentMemoryContext; PartitionTupleRouting *proute = NULL; @@ -798,6 +799,7 @@ CopyFrom(CopyFromState cstate) int64 excluded = 0; bool has_before_insert_row_trig; bool has_instead_insert_row_trig; + bool has_generated_stored = false; bool leafpart_use_multi_insert = false; Assert(cstate->rel); @@ -908,6 +910,34 @@ CopyFrom(CopyFromState cstate) ti_options |= TABLE_INSERT_FROZEN; } + if (cstate->whereClause) + { + TupleDesc tupDesc = RelationGetDescr(cstate->rel); + + if (tupDesc->constr && tupDesc->constr->has_generated_stored) + { + int k = -1; + Bitmapset *attnums = NULL; + + pull_varattnos(cstate->whereClause, 1, &attnums); + while ((k = bms_next_member(attnums, k)) >= 0) + { + Form_pg_attribute col; + AttrNumber attnum = k + FirstLowInvalidHeapAttributeNumber; + + col = TupleDescAttr(tupDesc, attnum - 1); + if (col->attgenerated == ATTRIBUTE_GENERATED_STORED) + { + has_generated_stored = true; + break; + } + } + } + + if (has_generated_stored) + tmpslot = table_slot_create(cstate->rel, NULL); + } + /* * We need a ResultRelInfo so we can use the regular executor's * index-entry-making machinery. (There used to be a huge amount of code @@ -1188,7 +1218,20 @@ CopyFrom(CopyFromState cstate) if (cstate->whereClause) { - econtext->ecxt_scantuple = myslot; + if (has_generated_stored) + { + ExecClearTuple(tmpslot); + + ExecCopySlot(tmpslot, myslot); + + ExecComputeStoredGenerated(resultRelInfo, estate, tmpslot, + CMD_INSERT); + + econtext->ecxt_scantuple = tmpslot; + } + else + econtext->ecxt_scantuple = myslot; + /* Skip items that don't match COPY's WHERE clause */ if (!ExecQual(cstate->qualexpr, econtext)) { @@ -1489,6 +1532,17 @@ CopyFrom(CopyFromState cstate) ExecResetTupleTable(estate->es_tupleTable, false); + if (has_generated_stored) + { + ExecClearTuple(tmpslot); + tmpslot->tts_ops->release(tmpslot); + if (tmpslot->tts_tupleDescriptor) + { + ReleaseTupleDesc(tmpslot->tts_tupleDescriptor); + tmpslot->tts_tupleDescriptor = NULL; + } + } + /* Allow the FDW to shut down */ if (target_resultRelInfo->ri_FdwRoutine != NULL && target_resultRelInfo->ri_FdwRoutine->EndForeignInsert != NULL) diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out index adac2cedfb2..d91989ae4cb 100644 --- a/src/test/regress/expected/generated_stored.out +++ b/src/test/regress/expected/generated_stored.out @@ -520,6 +520,7 @@ COPY gtest3 (a, b) TO stdout; ERROR: column "b" is a generated column DETAIL: Generated columns cannot be used in COPY. COPY gtest3 FROM stdin; +COPY gtest3 FROM stdin WHERE (b <> 15); COPY gtest3 (a, b) FROM stdin; ERROR: column "b" is a generated column DETAIL: Generated columns cannot be used in COPY. @@ -530,7 +531,8 @@ SELECT * FROM gtest3 ORDER BY a; 2 | 6 3 | 9 4 | 12 -(4 rows) + 6 | 18 +(5 rows) -- null values CREATE TABLE gtest2 (a int PRIMARY KEY, b int GENERATED ALWAYS AS (NULL) STORED); @@ -1067,6 +1069,16 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3; gtest_child3 | 09-13-2016 | 1 | 2 (3 rows) +COPY gtest_parent FROM STDIN WITH DELIMITER ',' WHERE f3 <> 2; +SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3; + tableoid | f1 | f2 | f3 +--------------+------------+----+---- + gtest_child | 07-15-2016 | 2 | 4 + gtest_child | 07-16-2016 | 4 | 8 + gtest_child2 | 08-15-2016 | 3 | 6 + gtest_child3 | 09-13-2016 | 1 | 2 +(4 rows) + -- we leave these tables around for purposes of testing dump/reload/upgrade -- 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); diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql index f56fde8d4e5..60770a54d89 100644 --- a/src/test/regress/sql/generated_stored.sql +++ b/src/test/regress/sql/generated_stored.sql @@ -231,6 +231,12 @@ COPY gtest3 FROM stdin; 4 \. +COPY gtest3 FROM stdin WHERE (b <> 15); +5 +6 +\. + + COPY gtest3 (a, b) FROM stdin; SELECT * FROM gtest3 ORDER BY a; @@ -496,6 +502,13 @@ ALTER TABLE gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2); \d gtest_child2 \d gtest_child3 SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3; + +COPY gtest_parent FROM STDIN WITH DELIMITER ',' WHERE f3 <> 2; +2016-07-15,1 +2016-07-16,4 +\. +SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3; + -- we leave these tables around for purposes of testing dump/reload/upgrade -- generated columns in partition key (not allowed) -- 2.34.1