Forbid deletion of non-updated rows from the same clause in MERGE USING

h2/src/main/org/h2/command/dml/Delete.java

@@ -5,6 +5,8 @@
  */
 package org.h2.command.dml;
 
+import java.util.HashSet;
+
 import org.h2.api.Trigger;
 import org.h2.command.CommandInterface;
 import org.h2.command.Prepared;
@@ -40,6 +42,8 @@ public class Delete extends Prepared {
      */
     private TableFilter sourceTableFilter;
 
+    private HashSet<Long> keysFilter;
+
     public Delete(Session session) {
         super(session);
     }
@@ -56,6 +60,15 @@ public class Delete extends Prepared {
         return this.condition;
     }
 
+    /**
+     * Sets the keys filter.
+     *
+     * @param keysFilter the keys filter
+     */
+    public void setKeysFilter(HashSet<Long> keysFilter) {
+        this.keysFilter = keysFilter;
+    }
+
     @Override
     public int update() {
         targetTableFilter.startQuery(session);
@@ -79,16 +92,18 @@ public class Delete extends Prepared {
                 setCurrentRowNumber(rows.size() + 1);
                 if (condition == null || condition.getBooleanValue(session)) {
                     Row row = targetTableFilter.get();
-                    boolean done = false;
-                    if (table.fireRow()) {
-                        done = table.fireBeforeRow(session, row, null);
-                    }
-                    if (!done) {
-                        rows.add(row);
-                    }
-                    count++;
-                    if (limitRows >= 0 && count >= limitRows) {
-                        break;
+                    if (keysFilter == null || keysFilter.contains(row.getKey())) {
+                        boolean done = false;
+                        if (table.fireRow()) {
+                            done = table.fireBeforeRow(session, row, null);
+                        }
+                        if (!done) {
+                            rows.add(row);
+                        }
+                        count++;
+                        if (limitRows >= 0 && count >= limitRows) {
+                            break;
+                        }
                     }
                 }
             }

h2/src/main/org/h2/command/dml/MergeUsing.java

@@ -7,6 +7,7 @@ package org.h2.command.dml;
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 
 import org.h2.api.ErrorCode;
 import org.h2.api.Trigger;
@@ -113,6 +114,7 @@ public class MergeUsing extends Prepared {
     private int countUpdatedRows;
     private Select targetMatchQuery;
     private final HashMap<Value, Integer> targetRowidsRemembered = new HashMap<>();
+    private final HashSet<Long> updatedKeys = new HashSet<>();
     private int sourceQueryRowNumber;
 
 
@@ -142,7 +144,8 @@ public class MergeUsing extends Prepared {
         sourceQueryRowNumber = 0;
         checkRights();
         setCurrentRowNumber(0);
-
+        // Just to be sure
+        updatedKeys.clear();
         // process source select query data for row creation
         ResultInterface rows = query.query(0);
         targetTable.fire(session, evaluateTriggerMasks(), true);
@@ -207,6 +210,7 @@ public class MergeUsing extends Prepared {
             // allowed together
             if (deleteCommand != null) {
                 countUpdatedRows += deleteCommand.update();
+                updatedKeys.clear();
             }
         } else {
             if (insertCommand != null) {
@@ -328,6 +332,10 @@ public class MergeUsing extends Prepared {
             deleteCommand.setSourceTableFilter(sourceTableFilter);
             deleteCommand.setCondition(appendOnCondition(deleteCommand));
             deleteCommand.prepare();
+            if (updateCommand != null) {
+                updateCommand.setUpdatedKeysCollector(updatedKeys);
+                deleteCommand.setKeysFilter(updatedKeys);
+            }
         }
         if (insertCommand != null) {
             insertCommand.setSourceTableFilter(sourceTableFilter);

h2/src/main/org/h2/command/dml/Update.java

@@ -7,6 +7,7 @@ package org.h2.command.dml;
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Objects;
 
 import org.h2.api.ErrorCode;
@@ -53,6 +54,8 @@ public class Update extends Prepared {
     private final ArrayList<Column> columns = Utils.newSmallArrayList();
     private final HashMap<Column, Expression> expressionMap  = new HashMap<>();
 
+    private HashSet<Long> updatedKeysCollector;
+
     public Update(Session session) {
         super(session);
     }
@@ -88,6 +91,15 @@ public class Update extends Prepared {
         }
     }
 
+    /**
+     * Sets the collector of updated keys.
+     *
+     * @param updatedKeysCollector the collector of updated keys
+     */
+    public void setUpdatedKeysCollector(HashSet<Long> updatedKeysCollector) {
+        this.updatedKeysCollector = updatedKeysCollector;
+    }
+
     @Override
     public int update() {
         targetTableFilter.startQuery(session);
@@ -135,7 +147,8 @@ public class Update extends Prepared {
                         }
                         newRow.setValue(i, newValue);
                     }
-                    newRow.setKey(oldRow.getKey());
+                    long key = oldRow.getKey();
+                    newRow.setKey(key);
                     if (setOnUpdate || updateToCurrentValuesReturnsZero) {
                         setOnUpdate = false;
                         for (int i = 0; i < columnCount; i++) {
@@ -166,6 +179,9 @@ public class Update extends Prepared {
                     if (!done) {
                         rows.add(oldRow);
                         rows.add(newRow);
+                        if (updatedKeysCollector != null) {
+                            updatedKeysCollector.add(key);
+                        }
                     }
                     count++;
                 }

h2/src/test/org/h2/test/scripts/dml/mergeUsing.sql

@@ -259,5 +259,28 @@ SELECT * FROM TEST;
 > 3  50
 > rows: 3
 
+MERGE INTO TEST USING (SELECT 1) ON 1 = 1
+    WHEN MATCHED THEN UPDATE SET VALUE = 60 WHERE ID = 3 DELETE WHERE ID = 2;
+> update count: 1
+
+SELECT * FROM TEST;
+> ID VALUE
+> -- -----
+> 1  50
+> 2  50
+> 3  60
+> rows: 3
+
+MERGE INTO TEST USING (SELECT 1) ON 1 = 1
+    WHEN MATCHED THEN DELETE WHERE ID = 2;
+> update count: 1
+
+SELECT * FROM TEST;
+> ID VALUE
+> -- -----
+> 1  50
+> 3  60
+> rows: 2
+
 DROP TABLE TEST;
 > ok