Fix MERGE command tag for cross-partition updates.

This ensures that the row count in the command tag for a MERGE is
correctly computed. Previously, if MERGE updated a partitioned table,
the row count would be incorrect if any row was moved to a different
partition, since such updates were counted twice.

Back-patch to v15, where MERGE was introduced.

Discussion: https://postgr.es/m/CAEZATCWRMG7XX2QEsVL1LswmNo2d_YG8tKTLkpD3=Lp644S7rg@mail.gmail.com

2 files changed, 27 insertions, 0 deletions

diff --git a/src/test/regress/expected/merge.out b/src/test/regress/expected/merge.out
index 2cf1409470b..c298bfa314c 100644
--- a/src/test/regress/expected/merge.out
+++ b/src/test/regress/expected/merge.out
@@ -1582,6 +1582,10 @@ SELECT * FROM pa_target ORDER BY tid;
 ROLLBACK;
 -- try updating the partition key column
 BEGIN;
+CREATE FUNCTION merge_func() RETURNS integer LANGUAGE plpgsql AS $$
+DECLARE
+  result integer;
+BEGIN
 MERGE INTO pa_target t
   USING pa_source s
   ON t.tid = s.sid
@@ -1589,6 +1593,18 @@ MERGE INTO pa_target t
     UPDATE SET tid = tid + 1, balance = balance + delta, val = val || ' updated by merge'
   WHEN NOT MATCHED THEN
     INSERT VALUES (sid, delta, 'inserted by merge');
+IF FOUND THEN
+  GET DIAGNOSTICS result := ROW_COUNT;
+END IF;
+RETURN result;
+END;
+$$;
+SELECT merge_func();
+ merge_func 
+------------
+         14
+(1 row)
+
 SELECT * FROM pa_target ORDER BY tid;
  tid | balance |           val            
 -----+---------+--------------------------
diff --git a/src/test/regress/sql/merge.sql b/src/test/regress/sql/merge.sql
index cef37e57d92..9ed648ad1f6 100644
--- a/src/test/regress/sql/merge.sql
+++ b/src/test/regress/sql/merge.sql
@@ -999,6 +999,10 @@ ROLLBACK;
 
 -- try updating the partition key column
 BEGIN;
+CREATE FUNCTION merge_func() RETURNS integer LANGUAGE plpgsql AS $$
+DECLARE
+  result integer;
+BEGIN
 MERGE INTO pa_target t
   USING pa_source s
   ON t.tid = s.sid
@@ -1006,6 +1010,13 @@ MERGE INTO pa_target t
     UPDATE SET tid = tid + 1, balance = balance + delta, val = val || ' updated by merge'
   WHEN NOT MATCHED THEN
     INSERT VALUES (sid, delta, 'inserted by merge');
+IF FOUND THEN
+  GET DIAGNOSTICS result := ROW_COUNT;
+END IF;
+RETURN result;
+END;
+$$;
+SELECT merge_func();
 SELECT * FROM pa_target ORDER BY tid;
 ROLLBACK;