From 92bfdd62bfd0d420029020e9b10b273718c962dc Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Thu, 6 Nov 2025 15:10:26 +0800
Subject: [PATCH v1 1/1] fix COPY WHERE clause with tableoid field

discussion: https://postgr.es/m/30c39ee8-bb11-4b8f-9697-45f7e018a8d3@eisentraut.org
---
 src/backend/commands/copyfrom.c    | 32 +++++++++++++++++++++++++++---
 src/test/regress/expected/copy.out | 16 +++++++++++++++
 src/test/regress/sql/copy.sql      | 14 +++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c
index 12781963b4f..94aed15d379 100644
--- a/src/backend/commands/copyfrom.c
+++ b/src/backend/commands/copyfrom.c
@@ -1178,15 +1178,16 @@ CopyFrom(CopyFromState cstate)
 		ExecStoreVirtualTuple(myslot);
 
 		/*
-		 * Constraints and where clause might reference the tableoid column,
-		 * so (re-)initialize tts_tableOid before evaluating them.
+		 * where clause might reference the tableoid column, so (re-)initialize
+		 * tts_tableOid before evaluating them. It may change later if we are
+		 * COPY INTO partitioned table.
 		 */
 		myslot->tts_tableOid = RelationGetRelid(target_resultRelInfo->ri_RelationDesc);
 
 		/* Triggers and stuff need to be invoked in query context. */
 		MemoryContextSwitchTo(oldcontext);
 
-		if (cstate->whereClause)
+		if (proute == NULL && cstate->whereClause)
 		{
 			econtext->ecxt_scantuple = myslot;
 			/* Skip items that don't match COPY's WHERE clause */
@@ -1215,6 +1216,24 @@ CopyFrom(CopyFromState cstate)
 			resultRelInfo = ExecFindPartition(mtstate, target_resultRelInfo,
 											  proute, myslot, estate);
 
+			myslot->tts_tableOid = RelationGetRelid(resultRelInfo->ri_RelationDesc);
+
+			if (cstate->whereClause)
+			{
+				econtext->ecxt_scantuple = myslot;
+				/* Skip items that don't match COPY's WHERE clause */
+				if (!ExecQual(cstate->qualexpr, econtext))
+				{
+					/*
+					 * Report that this tuple was filtered out by the WHERE
+					 * clause.
+					 */
+					pgstat_progress_update_param(PROGRESS_COPY_TUPLES_EXCLUDED,
+												++excluded);
+					continue;
+				}
+			}
+
 			if (prevResultRelInfo != resultRelInfo)
 			{
 				/* Determine which triggers exist on this partition */
@@ -1343,6 +1362,13 @@ CopyFrom(CopyFromState cstate)
 			}
 			else
 			{
+				/*
+				 * Constraints and GENERATED expressions might reference the
+				 * tableoid column, so (re-)initialize tts_tableOid before
+				 * evaluating them.
+				 */
+				myslot->tts_tableOid = RelationGetRelid(resultRelInfo->ri_RelationDesc);
+
 				/* Compute stored generated columns */
 				if (resultRelInfo->ri_RelationDesc->rd_att->constr &&
 					resultRelInfo->ri_RelationDesc->rd_att->constr->has_generated_stored)
diff --git a/src/test/regress/expected/copy.out b/src/test/regress/expected/copy.out
index 24e0f472f14..0735160bc10 100644
--- a/src/test/regress/expected/copy.out
+++ b/src/test/regress/expected/copy.out
@@ -121,6 +121,22 @@ insert into parted_copytest select x,2,'Two' from generate_series(1001,1010) x;
 insert into parted_copytest select x,1,'One' from generate_series(1011,1020) x;
 \set filename :abs_builddir '/results/parted_copytest.csv'
 copy (select * from parted_copytest order by a) to :'filename';
+truncate parted_copytest;
+COPY parted_copytest FROM STDIN WHERE (tableoid in ('parted_copytest'::regclass, 'parted_copytest_a1'::regclass));
+SELECT * FROM parted_copytest; --expect one row
+ a | b | c 
+---+---+---
+ 1 | 1 | a
+(1 row)
+
+COPY parted_copytest FROM STDIN WHERE (tableoid in ('parted_copytest_a2'::regclass));
+SELECT * FROM parted_copytest; --expect two row
+ a | b | c 
+---+---+---
+ 1 | 1 | a
+ 2 | 2 | b
+(2 rows)
+
 truncate parted_copytest;
 copy parted_copytest from :'filename';
 -- Ensure COPY FREEZE errors for partitioned tables.
diff --git a/src/test/regress/sql/copy.sql b/src/test/regress/sql/copy.sql
index 676a8b342b5..a2d3194635a 100644
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -148,6 +148,20 @@ copy (select * from parted_copytest order by a) to :'filename';
 
 truncate parted_copytest;
 
+COPY parted_copytest FROM STDIN WHERE (tableoid in ('parted_copytest'::regclass, 'parted_copytest_a1'::regclass));
+1	1	a
+2	2	b
+\.
+
+SELECT * FROM parted_copytest; --expect one row
+COPY parted_copytest FROM STDIN WHERE (tableoid in ('parted_copytest_a2'::regclass));
+2	2	b
+1	1	a
+\.
+SELECT * FROM parted_copytest; --expect two row
+truncate parted_copytest;
+
+
 copy parted_copytest from :'filename';
 
 -- Ensure COPY FREEZE errors for partitioned tables.
-- 
2.34.1