Add support for generated keys from INSERT ... SELECT

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

@@ -144,11 +144,11 @@ public class Insert extends Prepared implements ResultTarget {
         setCurrentRowNumber(0);
         table.fire(session, Trigger.INSERT, true);
         rowNumber = 0;
+        GeneratedKeys generatedKeys = session.getGeneratedKeys();
+        generatedKeys.initialize(table);
         int listSize = list.size();
         if (listSize > 0) {
             int columnLen = columns.length;
-            GeneratedKeys generatedKeys = session.getGeneratedKeys();
-            generatedKeys.initialize(table);
             for (int x = 0; x < listSize; x++) {
                 session.startStatementWithinTransaction();
                 generatedKeys.nextRow();
@@ -199,8 +199,12 @@ public class Insert extends Prepared implements ResultTarget {
             } else {
                 ResultInterface rows = query.query(0);
                 while (rows.next()) {
+                    generatedKeys.nextRow();
                     Value[] r = rows.currentRow();
-                    addRow(r);
+                    Row newRow = addRowImpl(r);
+                    if (newRow != null) {
+                        generatedKeys.confirmRow(newRow);
+                    }
                 }
                 rows.close();
             }
@@ -211,6 +215,10 @@ public class Insert extends Prepared implements ResultTarget {
 
     @Override
     public void addRow(Value[] values) {
+        addRowImpl(values);
+    }
+
+    private Row addRowImpl(Value[] values) {
         Row newRow = table.getTemplateRow();
         setCurrentRowNumber(++rowNumber);
         for (int j = 0, len = columns.length; j < len; j++) {
@@ -229,7 +237,9 @@ public class Insert extends Prepared implements ResultTarget {
             table.addRow(session, newRow);
             session.log(table, UndoLogRecord.INSERT, newRow);
             table.fireAfterRow(session, null, newRow, false);
+            return newRow;
         }
+        return null;
     }
 
     @Override

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

@@ -62,6 +62,7 @@ public class TestGetGeneratedKeys extends TestBase {
         Connection conn = getConnection("getGeneratedKeys");
         testBatchAndMergeInto(conn);
         testCalledSequenses(conn);
+        testInsertWithSelect(conn);
         testMergeUsing(conn);
         testMultithreaded(conn);
         testNameCase(conn);
@@ -231,6 +232,29 @@ public class TestGetGeneratedKeys extends TestBase {
         stat.execute("DROP SEQUENCE SEQ");
     }
 
+    /**
+     * Test method for INSERT ... SELECT operator.
+     *
+     * @param conn
+     *            connection
+     * @throws Exception
+     *             on exception
+     */
+    private void testInsertWithSelect(Connection conn) throws Exception {
+        Statement stat = conn.createStatement();
+        stat.execute("CREATE TABLE TEST (ID BIGINT PRIMARY KEY AUTO_INCREMENT, VALUE INT NOT NULL)");
+
+        PreparedStatement prep = conn.prepareStatement("INSERT INTO TEST(VALUE) SELECT 10",
+                Statement.RETURN_GENERATED_KEYS);
+        prep.executeUpdate();
+        ResultSet rs = prep.getGeneratedKeys();
+        assertTrue(rs.next());
+        assertEquals(1, rs.getLong(1));
+        rs.close();
+
+        stat.execute("DROP TABLE TEST");
+    }
+
     /**
      * Test method for MERGE USING operator.
      *