From 4e6b62fee8cda9a9fe7542f731e688b16cee2742 Mon Sep 17 00:00:00 2001
From: naisila <nicypp@gmail.com>
Date: Thu, 19 Dec 2024 13:13:27 +0300
Subject: [PATCH] Call ErrorIfMergeInCopy only if Citus tables are involved,
 fix tests

---
 src/backend/distributed/commands/multi_copy.c   | 9 +++++++--
 src/test/regress/expected/merge_unsupported.out | 4 ++--
 src/test/regress/expected/pgmerge.out           | 6 ------
 src/test/regress/sql/pgmerge.sql                | 5 -----
 4 files changed, 9 insertions(+), 15 deletions(-)

diff --git a/src/backend/distributed/commands/multi_copy.c b/src/backend/distributed/commands/multi_copy.c
index e8083dba9..268010e7a 100644
--- a/src/backend/distributed/commands/multi_copy.c
+++ b/src/backend/distributed/commands/multi_copy.c
@@ -2842,6 +2842,11 @@ ErrorIfMergeInCopy(CopyStmt *copyStatement)
 #else
 	if (!copyStatement->relation && (IsA(copyStatement->query, MergeStmt)))
 	{
+		/*
+		 * This path is currently not reachable because Merge in COPY can
+		 * only work with a RETURNING clause, and a RETURNING check
+		 * will error out sooner for Citus
+		 */
 		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						errmsg("MERGE with Citus tables "
 							   "is not yet supported in COPY")));
@@ -2860,8 +2865,6 @@ Node *
 ProcessCopyStmt(CopyStmt *copyStatement, QueryCompletion *completionTag, const
 				char *queryString)
 {
-	ErrorIfMergeInCopy(copyStatement);
-
 	/*
 	 * Handle special COPY "resultid" FROM STDIN WITH (format result) commands
 	 * for sending intermediate results to workers.
@@ -2889,6 +2892,8 @@ ProcessCopyStmt(CopyStmt *copyStatement, QueryCompletion *completionTag, const
 	 */
 	if (copyStatement->relation != NULL)
 	{
+		ErrorIfMergeInCopy(copyStatement);
+
 		bool isFrom = copyStatement->is_from;
 
 		/* consider using RangeVarGetRelidExtended to check perms before locking */
diff --git a/src/test/regress/expected/merge_unsupported.out b/src/test/regress/expected/merge_unsupported.out
index 669ae0be7..62f51a679 100644
--- a/src/test/regress/expected/merge_unsupported.out
+++ b/src/test/regress/expected/merge_unsupported.out
@@ -40,7 +40,7 @@ COPY (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE
 ) TO stdout;
-ERROR:  MERGE with Citus tables is not yet supported in COPY
+ERROR:  COPY query must have a RETURNING clause
 -- used in a CTE with RETURNING
 WITH foo AS (
   MERGE INTO target USING source ON (true)
@@ -52,7 +52,7 @@ COPY (
   MERGE INTO target USING source ON (true)
   WHEN MATCHED THEN DELETE RETURNING target.*
 ) TO stdout;
-ERROR:  MERGE with Citus tables is not yet supported in COPY
+ERROR:  MERGE with RETURNING is not yet supported for Citus tables
 -- unsupported relation types
 -- view
 CREATE VIEW tv AS SELECT count(tid) AS tid FROM target;
diff --git a/src/test/regress/expected/pgmerge.out b/src/test/regress/expected/pgmerge.out
index ea7cc42c5..0c2f9b741 100644
--- a/src/test/regress/expected/pgmerge.out
+++ b/src/test/regress/expected/pgmerge.out
@@ -162,12 +162,6 @@ ON tid = tid
 WHEN MATCHED THEN DO NOTHING;
 ERROR:  name "target" specified more than once
 DETAIL:  The name is used both as MERGE target table and data source.
--- used in COPY
-COPY (
-  MERGE INTO target USING source ON (true)
-  WHEN MATCHED THEN DELETE
-) TO stdout;
-ERROR:  MERGE with Citus tables is not yet supported in COPY
 -- unsupported relation types
 -- materialized view
 CREATE MATERIALIZED VIEW mv AS SELECT * FROM target;
diff --git a/src/test/regress/sql/pgmerge.sql b/src/test/regress/sql/pgmerge.sql
index 52d2cc800..69a0210bc 100644
--- a/src/test/regress/sql/pgmerge.sql
+++ b/src/test/regress/sql/pgmerge.sql
@@ -116,11 +116,6 @@ MERGE INTO target
 USING target
 ON tid = tid
 WHEN MATCHED THEN DO NOTHING;
--- used in COPY
-COPY (
-  MERGE INTO target USING source ON (true)
-  WHEN MATCHED THEN DELETE
-) TO stdout;
 
 -- unsupported relation types
 -- materialized view