Fix MERGE command tag for cross-partition updates.

deanrasheed · deanrasheed · commit 80a48e0f21c5 · 2023-02-22T09:39:09.000Z
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
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
@@ -2878,7 +2878,7 @@ ExecMergeMatched(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
 				}
 				ExecUpdatePrepareSlot(resultRelInfo, newslot, context->estate);
 				result = ExecUpdateAct(context, resultRelInfo, tupleid, NULL,
-									   newslot, mtstate->canSetTag, &updateCxt);
+									   newslot, false, &updateCxt);
 				if (result == TM_Ok && updateCxt.updated)
 				{
 					ExecUpdateEpilogue(context, &updateCxt, resultRelInfo,
diff --git a/src/test/regress/expected/merge.out b/src/test/regress/expected/merge.out
@@ -1582,13 +1582,29 @@ 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
   WHEN MATCHED THEN
     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
@@ -999,13 +999,24 @@ 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
   WHEN MATCHED THEN
     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;
 

Original file line number	Diff line number	Diff line change
`@@ -2878,7 +2878,7 @@ ExecMergeMatched(ModifyTableContext context, ResultRelInfo resultRelInfo,`
`2878`	`2878`	`}`
`2879`	`2879`	`ExecUpdatePrepareSlot(resultRelInfo, newslot, context->estate);`
`2880`	`2880`	`result = ExecUpdateAct(context, resultRelInfo, tupleid, NULL,`
`2881`		`- newslot, mtstate->canSetTag, &updateCxt);`
	`2881`	`+ newslot, false, &updateCxt);`
`2882`	`2882`	`if (result == TM_Ok && updateCxt.updated)`
`2883`	`2883`	`{`
`2884`	`2884`	`ExecUpdateEpilogue(context, &updateCxt, resultRelInfo,`