Issue 362: support LIMIT in UPDATE statements.

h2/src/docsrc/help/help.csv

@@ -59,7 +59,7 @@ INSERT INTO TEST VALUES(1, 'Hello')
 "Commands (DML)","UPDATE","
 UPDATE tableName [ [ AS ] newTableAlias ]
 SET { columnName = { DEFAULT | expression } } [,...]
-[ WHERE expression ]
+[ WHERE expression ] [ LIMIT expression ]
 ","
 Updates data in a table.
 ","

h2/src/main/org/h2/command/Parser.java

@@ -701,6 +701,10 @@ public class Parser {
             Expression condition = readExpression();
             command.setCondition(condition);
         }
+        if (readIf("LIMIT")) {
+            Expression limit = readTerm().optimize(session);
+            command.setLimit(limit);
+        }
         setSQL(command, "UPDATE", start);
         return command;
     }

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

@@ -29,6 +29,7 @@ import org.h2.util.New;
 import org.h2.util.StatementBuilder;
 import org.h2.util.StringUtils;
 import org.h2.value.Value;
+import org.h2.value.ValueNull;
 
 /**
  * This class represents the statement
@@ -38,6 +39,10 @@ public class Update extends Prepared {
 
     private Expression condition;
     private TableFilter tableFilter;
+
+    /** The limit expression as specified in the LIMIT clause. */
+    private Expression limitExpr;
+
     private ArrayList<Column> columns = New.arrayList();
     private HashMap<Column, Expression> expressionMap  = New.hashMap();
 
@@ -86,8 +91,18 @@ public class Update extends Prepared {
             setCurrentRowNumber(0);
             int count = 0;
             Column[] columns = table.getColumns();
+            int limitRows = -1;
+            if (limitExpr != null) {
+                Value v = limitExpr.getValue(session);
+                if (v != ValueNull.INSTANCE) {
+                    limitRows = v.getInt();
+                }
+            }
             while (tableFilter.next()) {
                 setCurrentRowNumber(count+1);
+                if (limitRows >= 0 && count >= limitRows) {
+                    break;
+                }
                 if (condition == null || Boolean.TRUE.equals(condition.getBooleanValue(session))) {
                     Row oldRow = tableFilter.get();
                     Row newRow = table.getTemplateRow();
@@ -185,6 +200,10 @@ public class Update extends Prepared {
         return CommandInterface.UPDATE;
     }
 
+    public void setLimit(Expression limit) {
+        this.limitExpr = limit;
+    }
+
     public boolean isCacheable() {
         return true;
     }

h2/src/test/org/h2/test/jdbc/TestLimitUpdates.java

+/*
+ * Copyright 2004-2011 H2 Group. Multiple-Licensed under the H2 License,
+ * Version 1.0, and under the Eclipse Public License, Version 1.0
+ * (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.test.jdbc;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import org.h2.test.TestBase;
+
+/**
+ * Test for limit updates.
+ */
+public class TestLimitUpdates extends TestBase {
+
+    private static final String DATABASE_NAME = "limitUpdates";
+
+    /**
+     * Run just this test.
+     *
+     * @param a ignored
+     */
+    public static void main(String... a) throws Exception {
+        TestBase.createCaller().init().test();
+    }
+
+    public void test() throws SQLException {
+        testLimitUpdates();
+        deleteDb(DATABASE_NAME);
+    }
+
+    private void testLimitUpdates() throws SQLException {
+        deleteDb(DATABASE_NAME);
+        Connection conn = null;
+        PreparedStatement prep = null;
+
+        try {
+            conn = getConnection(DATABASE_NAME);
+            prep = conn.prepareStatement(
+                    "CREATE TABLE TEST(KEY_ID INT PRIMARY KEY, VALUE_ID INT)");
+            prep.executeUpdate();
+
+            prep.close();
+            prep = conn.prepareStatement("INSERT INTO TEST VALUES(?, ?)");
+            int numRows = 10;
+            for (int i = 0; i < numRows; ++i) {
+                prep.setInt(1, i);
+                prep.setInt(2, 0);
+                prep.execute();
+            }
+            assertEquals(numRows, countWhere(conn, 0));
+
+            // update all elements than available
+            prep.close();
+            prep = conn.prepareStatement("UPDATE TEST SET VALUE_ID = ?");
+            prep.setInt(1, 1);
+            prep.execute();
+            assertEquals(numRows, countWhere(conn, 1));
+
+            // update less elements than available
+            updateLimit(conn, 2, numRows / 2);
+            assertEquals(numRows / 2, countWhere(conn, 2));
+
+            // update more elements than available
+            updateLimit(conn, 3, numRows * 2);
+            assertEquals(numRows, countWhere(conn, 3));
+
+            // update no elements
+            updateLimit(conn, 4, 0);
+            assertEquals(0, countWhere(conn, 4));
+        } finally {
+            if (prep != null) {
+                prep.close();
+            }
+            if (conn != null) {
+                conn.close();
+            }
+        }
+    }
+
+    private static int countWhere(final Connection conn, final int where)
+            throws SQLException {
+        PreparedStatement prep = null;
+        ResultSet rs = null;
+        try {
+            prep = conn.prepareStatement(
+                    "SELECT COUNT(*) FROM TEST WHERE VALUE_ID = ?");
+            prep.setInt(1, where);
+            rs = prep.executeQuery();
+            rs.next();
+            return rs.getInt(1);
+        } finally {
+            if (rs != null) {
+                rs.close();
+            }
+            if (prep != null) {
+                prep.close();
+            }
+        }
+    }
+
+    private static void updateLimit(final Connection conn, final int value,
+            final int limit) throws SQLException {
+        PreparedStatement prep = null;
+        try {
+            prep = conn.prepareStatement(
+                    "UPDATE TEST SET VALUE_ID = ? LIMIT ?");
+            prep.setInt(1, value);
+            prep.setInt(2, limit);
+            prep.execute();
+        } finally {
+            if (prep != null) {
+                prep.close();
+            }
+        }
+    }
+}