v2-0003-fix-COPY-WHERE-clause-stored-generated-column-references.patch

From 9706a4ebbcfe186b2d6c9fca37187518e4f917c1 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
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