Merge pull request #858 from katzyn/generated_keys

Return all generated rows and columns from getGeneratedKeys()

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

@@ -15,6 +15,7 @@ import org.h2.expression.ParameterInterface;
 import org.h2.message.DbException;
 import org.h2.message.Trace;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.util.MathUtils;
 
 /**
@@ -149,7 +150,7 @@ public abstract class Command implements CommandInterface {
     @Override
     public void stop() {
         session.endStatement();
-        session.setCurrentCommand(null);
+        session.setCurrentCommand(null, false);
         if (!isTransactional()) {
             session.commit(true);
         } else if (session.getAutoCommit()) {
@@ -193,7 +194,7 @@ public abstract class Command implements CommandInterface {
             }
         }
         synchronized (sync) {
-            session.setCurrentCommand(this);
+            session.setCurrentCommand(this, false);
             try {
                 while (true) {
                     database.checkPowerOff();
@@ -238,7 +239,7 @@ public abstract class Command implements CommandInterface {
     }
 
     @Override
-    public int executeUpdate() {
+    public ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest) {
         long start = 0;
         Database database = session.getDatabase();
         Object sync = database.isMultiThreaded() ? (Object) session : (Object) database;
@@ -252,12 +253,17 @@ public abstract class Command implements CommandInterface {
         }
         synchronized (sync) {
             Session.Savepoint rollback = session.setSavepoint();
-            session.setCurrentCommand(this);
+            session.setCurrentCommand(this, generatedKeysRequest);
             try {
                 while (true) {
                     database.checkPowerOff();
                     try {
-                        return update();
+                        int updateCount = update();
+                        if (!Boolean.FALSE.equals(generatedKeysRequest)) {
+                            return new ResultWithGeneratedKeys.WithKeys(updateCount,
+                                    session.getGeneratedKeys().getKeys(session));
+                        }
+                        return ResultWithGeneratedKeys.of(updateCount);
                     } catch (DbException e) {
                         start = filterConcurrentUpdate(e, start);
                     } catch (OutOfMemoryError e) {

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

@@ -8,6 +8,7 @@ package org.h2.command;
 import java.util.ArrayList;
 import org.h2.expression.ParameterInterface;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
 
 /**
  * Represents a SQL statement.
@@ -510,9 +511,16 @@ public interface CommandInterface {
     /**
      * Execute the statement
      *
+     * @param generatedKeysRequest
+     *            {@code false} if generated keys are not needed, {@code true} if
+     *            generated keys should be configured automatically, {@code int[]}
+     *            to specify column indices to return generated keys from, or
+     *            {@code String[]} to specify column names to return generated keys
+     *            from
+     *
      * @return the update count
      */
-    int executeUpdate();
+    ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest);
 
     /**
      * Stop the command execution, release all locks and resources

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

@@ -39,7 +39,7 @@ class CommandList extends Command {
 
     @Override
     public int update() {
-        int updateCount = command.executeUpdate();
+        int updateCount = command.executeUpdate(false).getUpdateCount();
         executeRemaining();
         return updateCount;
     }

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

@@ -9,6 +9,7 @@ import java.io.IOException;
 import java.util.ArrayList;
 
 import org.h2.engine.Constants;
+import org.h2.engine.GeneratedKeysMode;
 import org.h2.engine.SessionRemote;
 import org.h2.engine.SysProperties;
 import org.h2.expression.ParameterInterface;
@@ -17,6 +18,7 @@ import org.h2.message.DbException;
 import org.h2.message.Trace;
 import org.h2.result.ResultInterface;
 import org.h2.result.ResultRemote;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.util.New;
 import org.h2.value.Transfer;
 import org.h2.value.Value;
@@ -194,10 +196,14 @@ public class CommandRemote implements CommandInterface {
     }
 
     @Override
-    public int executeUpdate() {
+    public ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest) {
         checkParameters();
+        boolean supportsGeneratedKeys = session.isSupportsGeneratedKeys();
+        boolean readGeneratedKeys = supportsGeneratedKeys && !Boolean.FALSE.equals(generatedKeysRequest);
+        int objectId = readGeneratedKeys ? session.getNextId() : 0;
         synchronized (session) {
             int updateCount = 0;
+            ResultRemote generatedKeys = null;
             boolean autoCommit = false;
             for (int i = 0, count = 0; i < transferList.size(); i++) {
                 prepareIfRequired();
@@ -206,9 +212,39 @@ public class CommandRemote implements CommandInterface {
                     session.traceOperation("COMMAND_EXECUTE_UPDATE", id);
                     transfer.writeInt(SessionRemote.COMMAND_EXECUTE_UPDATE).writeInt(id);
                     sendParameters(transfer);
+                    if (supportsGeneratedKeys) {
+                        int mode = GeneratedKeysMode.valueOf(generatedKeysRequest);
+                        transfer.writeInt(mode);
+                        switch (mode) {
+                        case GeneratedKeysMode.COLUMN_NUMBERS: {
+                            int[] keys = (int[]) generatedKeysRequest;
+                            transfer.writeInt(keys.length);
+                            for (int key : keys) {
+                                transfer.writeInt(key);
+                            }
+                            break;
+                        }
+                        case GeneratedKeysMode.COLUMN_NAMES: {
+                            String[] keys = (String[]) generatedKeysRequest;
+                            transfer.writeInt(keys.length);
+                            for (String key : keys) {
+                                transfer.writeString(key);
+                            }
+                            break;
+                        }
+                        }
+                    }
                     session.done(transfer);
                     updateCount = transfer.readInt();
                     autoCommit = transfer.readBoolean();
+                    if (readGeneratedKeys) {
+                        int columnCount = transfer.readInt();
+                        if (generatedKeys != null) {
+                            generatedKeys.close();
+                            generatedKeys = null;
+                        }
+                        generatedKeys = new ResultRemote(session, transfer, objectId, columnCount, Integer.MAX_VALUE);
+                    }
                 } catch (IOException e) {
                     session.removeServer(e, i--, ++count);
                 }
@@ -216,7 +252,10 @@ public class CommandRemote implements CommandInterface {
             session.setAutoCommitFromServer(autoCommit);
             session.autoCommitIfCluster();
             session.readSessionState();
-            return updateCount;
+            if (generatedKeys != null) {
+                return new ResultWithGeneratedKeys.WithKeys(updateCount, generatedKeys);
+            }
+            return ResultWithGeneratedKeys.of(updateCount);
         }
     }
 

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

@@ -5667,6 +5667,10 @@ public class Parser {
             readIfEqualOrTo();
             read();
             return new NoOperation(session);
+        } else if (readIf("SCOPE_GENERATED_KEYS")) {
+            readIfEqualOrTo();
+            read();
+            return new NoOperation(session);
         } else if (readIf("SCHEMA")) {
             readIfEqualOrTo();
             Set command = new Set(session, SetTypes.SCHEMA);

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

@@ -12,6 +12,7 @@ import org.h2.api.Trigger;
 import org.h2.command.Command;
 import org.h2.command.CommandInterface;
 import org.h2.command.Prepared;
+import org.h2.engine.GeneratedKeys;
 import org.h2.engine.Right;
 import org.h2.engine.Session;
 import org.h2.engine.UndoLogRecord;
@@ -20,6 +21,7 @@ import org.h2.expression.ConditionAndOr;
 import org.h2.expression.Expression;
 import org.h2.expression.ExpressionColumn;
 import org.h2.expression.Parameter;
+import org.h2.expression.SequenceValue;
 import org.h2.index.Index;
 import org.h2.message.DbException;
 import org.h2.mvstore.db.MVPrimaryIndex;
@@ -142,11 +144,14 @@ 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;
             for (int x = 0; x < listSize; x++) {
                 session.startStatementWithinTransaction();
+                generatedKeys.nextRow();
                 Row newRow = table.getTemplateRow();
                 Expression[] expr = list.get(x);
                 setCurrentRowNumber(x + 1);
@@ -160,6 +165,9 @@ public class Insert extends Prepared implements ResultTarget {
                         try {
                             Value v = c.convert(e.getValue(session), session.getDatabase().getMode());
                             newRow.setValue(index, v);
+                            if (e instanceof SequenceValue) {
+                                generatedKeys.add(c);
+                            }
                         } catch (DbException ex) {
                             throw setRow(ex, x, getSQL(expr));
                         }
@@ -179,6 +187,7 @@ public class Insert extends Prepared implements ResultTarget {
                             continue;
                         }
                     }
+                    generatedKeys.confirmRow(newRow);
                     session.log(table, UndoLogRecord.INSERT, newRow);
                     table.fireAfterRow(session, null, newRow, false);
                 }
@@ -190,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();
             }
@@ -202,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++) {
@@ -220,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/main/org/h2/command/dml/Merge.java

@@ -11,11 +11,13 @@ import org.h2.api.Trigger;
 import org.h2.command.Command;
 import org.h2.command.CommandInterface;
 import org.h2.command.Prepared;
+import org.h2.engine.GeneratedKeys;
 import org.h2.engine.Right;
 import org.h2.engine.Session;
 import org.h2.engine.UndoLogRecord;
 import org.h2.expression.Expression;
 import org.h2.expression.Parameter;
+import org.h2.expression.SequenceValue;
 import org.h2.index.Index;
 import org.h2.message.DbException;
 import org.h2.result.ResultInterface;
@@ -84,11 +86,14 @@ public class Merge extends Prepared {
         session.getUser().checkRight(targetTable, Right.INSERT);
         session.getUser().checkRight(targetTable, Right.UPDATE);
         setCurrentRowNumber(0);
+        GeneratedKeys generatedKeys = session.getGeneratedKeys();
         if (!valuesExpressionList.isEmpty()) {
             // process values in list
             count = 0;
+            generatedKeys.initialize(targetTable);
             for (int x = 0, size = valuesExpressionList.size(); x < size; x++) {
                 setCurrentRowNumber(x + 1);
+                generatedKeys.nextRow();
                 Expression[] expr = valuesExpressionList.get(x);
                 Row newRow = targetTable.getTemplateRow();
                 for (int i = 0, len = columns.length; i < len; i++) {
@@ -100,6 +105,9 @@ public class Merge extends Prepared {
                         try {
                             Value v = c.convert(e.getValue(session));
                             newRow.setValue(index, v);
+                            if (e instanceof SequenceValue) {
+                                generatedKeys.add(c);
+                            }
                         } catch (DbException ex) {
                             throw setRow(ex, count, getSQL(expr));
                         }
@@ -116,6 +124,7 @@ public class Merge extends Prepared {
             targetTable.lock(session, true, false);
             while (rows.next()) {
                 count++;
+                generatedKeys.nextRow();
                 Value[] r = rows.currentRow();
                 Row newRow = targetTable.getTemplateRow();
                 setCurrentRowNumber(count);
@@ -171,6 +180,7 @@ public class Merge extends Prepared {
                 if (!done) {
                     targetTable.lock(session, true, false);
                     targetTable.addRow(session, row);
+                    session.getGeneratedKeys().confirmRow(row);
                     session.log(targetTable, UndoLogRecord.INSERT, row);
                     targetTable.fireAfterRow(session, null, row, false);
                 }

h2/src/main/org/h2/engine/ConnectionInfo.java

@@ -96,7 +96,8 @@ public class ConnectionInfo implements Cloneable {
                 "CREATE", "CACHE_TYPE", "FILE_LOCK", "IGNORE_UNKNOWN_SETTINGS",
                 "IFEXISTS", "INIT", "PASSWORD", "RECOVER", "RECOVER_TEST",
                 "USER", "AUTO_SERVER", "AUTO_SERVER_PORT", "NO_UPGRADE",
-                "AUTO_RECONNECT", "OPEN_NEW", "PAGE_SIZE", "PASSWORD_HASH", "JMX" };
+                "AUTO_RECONNECT", "OPEN_NEW", "PAGE_SIZE", "PASSWORD_HASH", "JMX",
+                "SCOPE_GENERATED_KEYS" };
         HashSet<String> set = new HashSet<>(list.size() + connectionTime.length);
         set.addAll(list);
         for (String key : connectionTime) {

h2/src/main/org/h2/engine/Constants.java

@@ -100,6 +100,21 @@ public class Constants {
      */
     public static final int TCP_PROTOCOL_VERSION_16 = 16;
 
+    /**
+     * The TCP protocol version number 17.
+     */
+    public static final int TCP_PROTOCOL_VERSION_17 = 17;
+
+    /**
+     * Minimum supported version of TCP protocol.
+     */
+    public static final int TCP_PROTOCOL_VERSION_MIN_SUPPORTED = TCP_PROTOCOL_VERSION_6;
+
+    /**
+     * Maximum supported version of TCP protocol.
+     */
+    public static final int TCP_PROTOCOL_VERSION_MAX_SUPPORTED = TCP_PROTOCOL_VERSION_17;
+
     /**
      * The major version of this database.
      */

h2/src/main/org/h2/engine/Engine.java

@@ -206,7 +206,7 @@ public class Engine implements SessionFactory {
                     CommandInterface command = session.prepareCommand(
                             "SET " + Parser.quoteIdentifier(setting) + " " + value,
                             Integer.MAX_VALUE);
-                    command.executeUpdate();
+                    command.executeUpdate(false);
                 } catch (DbException e) {
                     if (e.getErrorCode() == ErrorCode.ADMIN_RIGHTS_REQUIRED) {
                         session.getTrace().error(e, "admin rights required; user: \"" +
@@ -224,7 +224,7 @@ public class Engine implements SessionFactory {
                 try {
                     CommandInterface command = session.prepareCommand(init,
                             Integer.MAX_VALUE);
-                    command.executeUpdate();
+                    command.executeUpdate(false);
                 } catch (DbException e) {
                     if (!ignoreUnknownSetting) {
                         session.close();

h2/src/main/org/h2/engine/GeneratedKeys.java

+/*
+ * Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
+ * and the EPL 1.0 (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.engine;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.h2.expression.Expression;
+import org.h2.expression.ExpressionColumn;
+import org.h2.result.LocalResult;
+import org.h2.result.Row;
+import org.h2.table.Column;
+import org.h2.table.Table;
+import org.h2.util.New;
+import org.h2.util.StringUtils;
+import org.h2.value.Value;
+import org.h2.value.ValueNull;
+
+/**
+ * Class for gathering and processing of generated keys.
+ */
+public final class GeneratedKeys {
+    /**
+     * Data for result set with generated keys.
+     */
+    private final ArrayList<Map<Column, Value>> data = New.arrayList();
+
+    /**
+     * Columns with generated keys in the current row.
+     */
+    private final ArrayList<Column> row = New.arrayList();
+
+    /**
+     * All columns with generated keys.
+     */
+    private final ArrayList<Column> allColumns = New.arrayList();
+
+    /**
+     * Request for keys gathering. {@code false} if generated keys are not needed,
+     * {@code true} if generated keys should be configured automatically,
+     * {@code int[]} to specify column indices to return generated keys from, or
+     * {@code String[]} to specify column names to return generated keys from.
+     */
+    private Object generatedKeysRequest;
+
+    /**
+     * Processed table.
+     */
+    private Table table;
+
+    /**
+     * Remembers columns with generated keys.
+     *
+     * @param column
+     *            table column
+     */
+    public void add(Column column) {
+        if (Boolean.FALSE.equals(generatedKeysRequest)) {
+            return;
+        }
+        row.add(column);
+    }
+
+    /**
+     * Clears all information from previous runs and sets a new request for
+     * gathering of generated keys.
+     *
+     * @param generatedKeysRequest
+     *            {@code false} if generated keys are not needed, {@code true} if
+     *            generated keys should be configured automatically, {@code int[]}
+     *            to specify column indices to return generated keys from, or
+     *            {@code String[]} to specify column names to return generated keys
+     *            from
+     */
+    public void clear(Object generatedKeysRequest) {
+        this.generatedKeysRequest = generatedKeysRequest;
+        data.clear();
+        row.clear();
+        allColumns.clear();
+        table = null;
+    }
+
+    /**
+     * Saves row with generated keys if any.
+     *
+     * @param tableRow
+     *            table row that was inserted
+     */
+    public void confirmRow(Row tableRow) {
+        if (Boolean.FALSE.equals(generatedKeysRequest)) {
+            return;
+        }
+        int size = row.size();
+        if (size > 0) {
+            if (size == 1) {
+                Column column = row.get(0);
+                data.add(Collections.singletonMap(column, tableRow.getValue(column.getColumnId())));
+                if (!allColumns.contains(column)) {
+                    allColumns.add(column);
+                }
+            } else {
+                HashMap<Column, Value> map = new HashMap<>();
+                for (Column column : row) {
+                    map.put(column, tableRow.getValue(column.getColumnId()));
+                    if (!allColumns.contains(column)) {
+                        allColumns.add(column);
+                    }
+                }
+                data.add(map);
+            }
+            row.clear();
+        }
+    }
+
+    /**
+     * Returns generated keys.
+     *
+     * @return local result with generated keys
+     */
+    public LocalResult getKeys(Session session) {
+        Database db = session == null ? null : session.getDatabase();
+        if (Boolean.FALSE.equals(generatedKeysRequest)) {
+            clear(null);
+            return new LocalResult();
+        }
+        ArrayList<ExpressionColumn> expressionColumns;
+        if (Boolean.TRUE.equals(generatedKeysRequest)) {
+            expressionColumns = new ArrayList<>(allColumns.size());
+            for (Column column : allColumns) {
+                expressionColumns.add(new ExpressionColumn(db, column));
+            }
+        } else if (generatedKeysRequest instanceof int[]) {
+            if (table != null) {
+                int[] indices = (int[]) generatedKeysRequest;
+                Column[] columns = table.getColumns();
+                int cnt = columns.length;
+                allColumns.clear();
+                expressionColumns = new ArrayList<>(indices.length);
+                for (int idx : indices) {
+                    if (idx >= 1 && idx <= cnt) {
+                        Column column = columns[idx - 1];
+                        expressionColumns.add(new ExpressionColumn(db, column));
+                        allColumns.add(column);
+                    }
+                }
+            } else {
+                clear(null);
+                return new LocalResult();
+            }
+        } else if (generatedKeysRequest instanceof String[]) {
+            if (table != null) {
+                String[] names = (String[]) generatedKeysRequest;
+                allColumns.clear();
+                expressionColumns = new ArrayList<>(names.length);
+                for (String name : names) {
+                    Column column;
+                    search: if (table.doesColumnExist(name)) {
+                        column = table.getColumn(name);
+                    } else {
+                        name = StringUtils.toUpperEnglish(name);
+                        if (table.doesColumnExist(name)) {
+                            column = table.getColumn(name);
+                        } else {
+                            for (Column c : table.getColumns()) {
+                                if (c.getName().equalsIgnoreCase(name)) {
+                                    column = c;
+                                    break search;
+                                }
+                            }
+                            continue;
+                        }
+                    }
+                    expressionColumns.add(new ExpressionColumn(db, column));
+                    allColumns.add(column);
+                }
+            } else {
+                clear(null);
+                return new LocalResult();
+            }
+        } else {
+            clear(null);
+            return new LocalResult();
+        }
+        int columnCount = expressionColumns.size();
+        if (columnCount == 0) {
+            clear(null);
+            return new LocalResult();
+        }
+        LocalResult result = new LocalResult(session, expressionColumns.toArray(new Expression[0]), columnCount);
+        for (Map<Column, Value> map : data) {
+            Value[] row = new Value[columnCount];
+            for (Map.Entry<Column, Value> entry : map.entrySet()) {
+                int idx = allColumns.indexOf(entry.getKey());
+                if (idx >= 0) {
+                    row[idx] = entry.getValue();
+                }
+            }
+            for (int i = 0; i < columnCount; i++) {
+                if (row[i] == null) {
+                    row[i] = ValueNull.INSTANCE;
+                }
+            }
+            result.addRow(row);
+        }
+        clear(null);
+        return result;
+    }
+
+    /**
+     * Initializes processing of the specified table. Should be called after
+     * {@code clear()}, but before other methods.
+     *
+     * @param table
+     *            table
+     */
+    public void initialize(Table table) {
+        this.table = table;
+    }
+
+    /**
+     * Clears unsaved information about previous row, if any. Should be called
+     * before processing of a new row if previous row was not confirmed or simply
+     * always before each row.
+     */
+    public void nextRow() {
+        row.clear();
+    }
+
+    @Override
+    public String toString() {
+        return allColumns + ": " + data.size();
+    }
+
+}

h2/src/main/org/h2/engine/GeneratedKeysMode.java

+/*
+ * Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
+ * and the EPL 1.0 (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.engine;
+
+import org.h2.api.ErrorCode;
+import org.h2.message.DbException;
+
+/**
+ * Modes of generated keys' gathering.
+ */
+public final class GeneratedKeysMode {
+
+    /**
+     * Generated keys are not needed.
+     */
+    public static final int NONE = 0;
+
+    /**
+     * Generated keys should be configured automatically.
+     */
+    public static final int AUTO = 1;
+
+    /**
+     * Use specified column indices to return generated keys from.
+     */
+    public static final int COLUMN_NUMBERS = 2;
+
+    /**
+     * Use specified column names to return generated keys from.
+     */
+    public static final int COLUMN_NAMES = 3;
+
+    /**
+     * Determines mode of generated keys' gathering.
+     *
+     * @param generatedKeysRequest
+     *            {@code false} if generated keys are not needed, {@code true} if
+     *            generated keys should be configured automatically, {@code int[]}
+     *            to specify column indices to return generated keys from, or
+     *            {@code String[]} to specify column names to return generated keys
+     *            from
+     * @return mode for the specified generated keys request
+     */
+    public static int valueOf(Object generatedKeysRequest) {
+        if (Boolean.FALSE.equals(generatedKeysRequest)) {
+            return NONE;
+        }
+        if (Boolean.TRUE.equals(generatedKeysRequest)) {
+            return AUTO;
+        }
+        if (generatedKeysRequest instanceof int[]) {
+            return COLUMN_NUMBERS;
+        }
+        if (generatedKeysRequest instanceof String[]) {
+            return COLUMN_NAMES;
+        }
+        throw DbException.get(ErrorCode.INVALID_VALUE_2,
+                generatedKeysRequest == null ? "null" : generatedKeysRequest.toString());
+    }
+
+    private GeneratedKeysMode() {
+    }
+
+}

h2/src/main/org/h2/engine/Session.java

@@ -85,6 +85,7 @@ public class Session extends SessionWithState {
     private Value lastIdentity = ValueLong.get(0);
     private Value lastScopeIdentity = ValueLong.get(0);
     private Value lastTriggerIdentity;
+    private GeneratedKeys generatedKeys;
     private int firstUncommittedLog = Session.LOG_WRITTEN;
     private int firstUncommittedPos = Session.LOG_WRITTEN;
     private HashMap<String, Savepoint> savepoints;
@@ -1075,6 +1076,13 @@ public class Session extends SessionWithState {
         return lastTriggerIdentity;
     }
 
+    public GeneratedKeys getGeneratedKeys() {
+        if (generatedKeys == null) {
+            generatedKeys = new GeneratedKeys();
+        }
+        return generatedKeys;
+    }
+
     /**
      * Called when a log entry for this session is added. The session keeps
      * track of the first entry in the transaction log that is not yet
@@ -1237,9 +1245,20 @@ public class Session extends SessionWithState {
      * executing the statement.
      *
      * @param command the command
-     */
-    public void setCurrentCommand(Command command) {
+     * @param generatedKeysRequest
+     *            {@code false} if generated keys are not needed, {@code true} if
+     *            generated keys should be configured automatically, {@code int[]}
+     *            to specify column indices to return generated keys from, or
+     *            {@code String[]} to specify column names to return generated keys
+     *            from
+     */
+    public void setCurrentCommand(Command command, Object generatedKeysRequest) {
         this.currentCommand = command;
+        // Preserve generated keys in case of a new query due to possible nested
+        // queries in update
+        if (command != null && !command.isQuery()) {
+            getGeneratedKeys().clear(generatedKeysRequest);
+        }
         if (queryTimeout > 0 && command != null) {
             currentCommandStart = System.currentTimeMillis();
             long now = System.nanoTime();
@@ -1790,4 +1809,10 @@ public class Session extends SessionWithState {
     public void setColumnNamerConfiguration(ColumnNamerConfiguration columnNamerConfiguration) {
         this.columnNamerConfiguration = columnNamerConfiguration;
     }
+
+    @Override
+    public boolean isSupportsGeneratedKeys() {
+        return true;
+    }
+
 }

h2/src/main/org/h2/engine/SessionInterface.java

@@ -153,4 +153,13 @@ public interface SessionInterface extends Closeable {
      * @return the current schema name
      */
     String getCurrentSchemaName();
+
+    /**
+     * Returns is this session supports generated keys.
+     *
+     * @return {@code true} if generated keys are supported, {@code false} if only
+     *         {@code SCOPE_IDENTITY()} is supported
+     */
+    boolean isSupportsGeneratedKeys();
+
 }

h2/src/main/org/h2/engine/SessionRemote.java

@@ -117,8 +117,8 @@ public class SessionRemote extends SessionWithState implements DataHandler {
         Transfer trans = new Transfer(this, socket);
         trans.setSSL(ci.isSSL());
         trans.init();
-        trans.writeInt(Constants.TCP_PROTOCOL_VERSION_6);
-        trans.writeInt(Constants.TCP_PROTOCOL_VERSION_16);
+        trans.writeInt(Constants.TCP_PROTOCOL_VERSION_MIN_SUPPORTED);
+        trans.writeInt(Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED);
         trans.writeString(db);
         trans.writeString(ci.getOriginalURL());
         trans.writeString(ci.getUserName());
@@ -210,7 +210,7 @@ public class SessionRemote extends SessionWithState implements DataHandler {
             CommandInterface c = prepareCommand(
                     "SET CLUSTER " + serverList, Integer.MAX_VALUE);
             // this will set autoCommit to false
-            c.executeUpdate();
+            c.executeUpdate(false);
             // so we need to switch it on
             autoCommit = true;
             cluster = true;
@@ -265,13 +265,13 @@ public class SessionRemote extends SessionWithState implements DataHandler {
                     autoCommitTrue = prepareCommand(
                             "SET AUTOCOMMIT TRUE", Integer.MAX_VALUE);
                 }
-                autoCommitTrue.executeUpdate();
+                autoCommitTrue.executeUpdate(false);
             } else {
                 if (autoCommitFalse == null) {
                     autoCommitFalse = prepareCommand(
                             "SET AUTOCOMMIT FALSE", Integer.MAX_VALUE);
                 }
-                autoCommitFalse.executeUpdate();
+                autoCommitFalse.executeUpdate(false);
             }
         }
     }
@@ -468,7 +468,7 @@ public class SessionRemote extends SessionWithState implements DataHandler {
 
     private void switchOffCluster() {
         CommandInterface ci = prepareCommand("SET CLUSTER ''", Integer.MAX_VALUE);
-        ci.executeUpdate();
+        ci.executeUpdate(false);
     }
 
     /**
@@ -869,4 +869,10 @@ public class SessionRemote extends SessionWithState implements DataHandler {
     public void setCurrentSchemaName(String schema) {
         throw DbException.getUnsupportedException("setSchema && remote session");
     }
+
+    @Override
+    public boolean isSupportsGeneratedKeys() {
+        return getClientVersion() >= Constants.TCP_PROTOCOL_VERSION_17;
+    }
+
 }

h2/src/main/org/h2/engine/SessionWithState.java

@@ -29,7 +29,7 @@ abstract class SessionWithState implements SessionInterface {
             try {
                 for (String sql : sessionState) {
                     CommandInterface ci = prepareCommand(sql, Integer.MAX_VALUE);
-                    ci.executeUpdate();
+                    ci.executeUpdate(false);
                 }
             } finally {
                 sessionStateUpdating = false;

h2/src/main/org/h2/fulltext/FullText.java

@@ -950,7 +950,8 @@ public class FullText {
             useOwnConnection = isMultiThread(conn);
             if(!useOwnConnection) {
                 for (int i = 0; i < SQL.length; i++) {
-                    prepStatements[i] = conn.prepareStatement(SQL[i]);
+                    prepStatements[i] = conn.prepareStatement(SQL[i],
+                            Statement.RETURN_GENERATED_KEYS);
                 }
             }
         }
@@ -1154,7 +1155,9 @@ public class FullText {
         }
 
         private PreparedStatement getStatement(Connection conn, int index) throws SQLException {
-            return useOwnConnection ? conn.prepareStatement(SQL[index]) : prepStatements[index];
+            return useOwnConnection ?
+                    conn.prepareStatement(SQL[index], Statement.RETURN_GENERATED_KEYS)
+                    : prepStatements[index];
         }
 
     }

h2/src/main/org/h2/jdbc/JdbcCallableStatement.java

@@ -47,7 +47,7 @@ public class JdbcCallableStatement extends JdbcPreparedStatement implements
 
     JdbcCallableStatement(JdbcConnection conn, String sql, int id,
             int resultSetType, int resultSetConcurrency) {
-        super(conn, sql, id, resultSetType, resultSetConcurrency, false);
+        super(conn, sql, id, resultSetType, resultSetConcurrency, false, false);
         setTrace(session.getTrace(), TraceObject.CALLABLE_STATEMENT, id);
     }
 

h2/src/main/org/h2/jdbc/JdbcPreparedStatement.java

@@ -31,9 +31,10 @@ import org.h2.expression.ParameterInterface;
 import org.h2.message.DbException;
 import org.h2.message.TraceObject;
 import org.h2.result.ResultInterface;
-import org.h2.tools.SimpleResultSet;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.util.DateTimeUtils;
 import org.h2.util.IOUtils;
+import org.h2.util.MergedResultSet;
 import org.h2.util.New;
 import org.h2.value.DataType;
 import org.h2.value.Value;
@@ -61,13 +62,15 @@ public class JdbcPreparedStatement extends JdbcStatement implements
     protected CommandInterface command;
     private final String sqlStatement;
     private ArrayList<Value[]> batchParameters;
-    private ArrayList<Object> batchIdentities;
+    private MergedResultSet batchIdentities;
     private HashMap<String, Integer> cachedColumnLabelMap;
+    private final Object generatedKeysRequest;
 
     JdbcPreparedStatement(JdbcConnection conn, String sql, int id,
             int resultSetType, int resultSetConcurrency,
-            boolean closeWithResultSet) {
+            boolean closeWithResultSet, Object generatedKeysRequest) {
         super(conn, id, resultSetType, resultSetConcurrency, closeWithResultSet);
+        this.generatedKeysRequest = conn.scopeGeneratedKeys() ? false : generatedKeysRequest;
         setTrace(session.getTrace(), TraceObject.PREPARED_STATEMENT, id);
         this.sqlStatement = sql;
         command = conn.prepareCommand(sql, fetchSize);
@@ -193,7 +196,14 @@ public class JdbcPreparedStatement extends JdbcStatement implements
         synchronized (session) {
             try {
                 setExecutingStatement(command);
-                updateCount = command.executeUpdate();
+                ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
+                updateCount = result.getUpdateCount();
+                ResultInterface gk = result.getGeneratedKeys();
+                if (gk != null) {
+                    int id = getNextId(TraceObject.RESULT_SET);
+                    generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
+                            false, true, false);
+                }
             } finally {
                 setExecutingStatement(null);
             }
@@ -236,7 +246,13 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                                     updatable, cachedColumnLabelMap);
                         } else {
                             returnsResultSet = false;
-                            updateCount = command.executeUpdate();
+                            ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
+                            updateCount = result.getUpdateCount();
+                            ResultInterface gk = result.getGeneratedKeys();
+                            if (gk != null) {
+                                generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
+                                        false, true, false);
+                            }
                         }
                     } finally {
                         if (!lazy) {
@@ -1243,7 +1259,7 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                 // set
                 batchParameters = New.arrayList();
             }
-            batchIdentities = New.arrayList();
+            batchIdentities = new MergedResultSet();
             int size = batchParameters.size();
             int[] result = new int[size];
             boolean error = false;
@@ -1261,10 +1277,9 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                     }
                     try {
                         result[i] = executeUpdateInternal();
-                        ResultSet rs = conn.getGeneratedKeys(this, id);
-                        while (rs.next()) {
-                            batchIdentities.add(rs.getObject(1));
-                        }
+                        // Cannot use own implementation, it returns batch identities
+                        ResultSet rs = super.getGeneratedKeys();
+                        batchIdentities.add(rs);
                     } catch (Exception re) {
                         SQLException e = logAndConvert(re);
                         if (next == null) {
@@ -1293,14 +1308,8 @@ public class JdbcPreparedStatement extends JdbcStatement implements
 
     @Override
     public ResultSet getGeneratedKeys() throws SQLException {
-        if (batchIdentities != null && !batchIdentities.isEmpty()) {
-            SimpleResultSet rs = new SimpleResultSet();
-            rs.addColumn("identity", java.sql.Types.INTEGER,
-                    10, 0);
-            for (Object o : batchIdentities) {
-                rs.addRow(o);
-            }
-            return rs;
+        if (batchIdentities != null) {
+            return batchIdentities.getResult();
         }
         return super.getGeneratedKeys();
     }

h2/src/main/org/h2/jdbc/JdbcSavepoint.java

@@ -65,7 +65,7 @@ public class JdbcSavepoint extends TraceObject implements Savepoint {
         checkValid();
         conn.prepareCommand(
                 "ROLLBACK TO SAVEPOINT " + getName(name, savepointId),
-                Integer.MAX_VALUE).executeUpdate();
+                Integer.MAX_VALUE).executeUpdate(false);
     }
 
     private void checkValid() {

h2/src/main/org/h2/result/ResultWithGeneratedKeys.java

+/*
+ * Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
+ * and the EPL 1.0 (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.result;
+
+/**
+ * Result of update command with optional generated keys.
+ */
+public class ResultWithGeneratedKeys {
+    /**
+     * Result of update command with generated keys;
+     */
+    public static final class WithKeys extends ResultWithGeneratedKeys {
+        private final ResultInterface generatedKeys;
+
+        /**
+         * Creates a result with update count and generated keys.
+         *
+         * @param updateCount
+         *            update count
+         * @param generatedKeys
+         *            generated keys
+         */
+        public WithKeys(int updateCount, ResultInterface generatedKeys) {
+            super(updateCount);
+            this.generatedKeys = generatedKeys;
+        }
+
+        @Override
+        public ResultInterface getGeneratedKeys() {
+            return generatedKeys;
+        }
+    }
+
+    /**
+     * Returns a result with only update count.
+     *
+     * @param updateCount
+     *            update count
+     */
+    public static ResultWithGeneratedKeys of(int updateCount) {
+        return new ResultWithGeneratedKeys(updateCount);
+    }
+
+    private final int updateCount;
+
+    ResultWithGeneratedKeys(int updateCount) {
+        this.updateCount = updateCount;
+    }
+
+    /**
+     * Returns generated keys, or {@code null}.
+     *
+     * @return generated keys, or {@code null}
+     */
+    public ResultInterface getGeneratedKeys() {
+        return null;
+    }
+
+    /**
+     * Returns update count.
+     *
+     * @return update count
+     */
+    public int getUpdateCount() {
+        return updateCount;
+    }
+
+}

h2/src/main/org/h2/schema/TriggerObject.java

@@ -203,6 +203,7 @@ public class TriggerObject extends SchemaObjectBase {
      * times for each statement.
      *
      * @param session the session
+     * @param table the table
      * @param oldRow the old row
      * @param newRow the new row
      * @param beforeAction true if this method is called before the operation is
@@ -210,7 +211,7 @@ public class TriggerObject extends SchemaObjectBase {
      * @param rollback when the operation occurred within a rollback
      * @return true if no further action is required (for 'instead of' triggers)
      */
-    public boolean fireRow(Session session, Row oldRow, Row newRow,
+    public boolean fireRow(Session session, Table table, Row oldRow, Row newRow,
             boolean beforeAction, boolean rollback) {
         if (!rowBased || before != beforeAction) {
             return false;
@@ -260,6 +261,7 @@ public class TriggerObject extends SchemaObjectBase {
                     Object o = newList[i];
                     if (o != newListBackup[i]) {
                         Value v = DataType.convertToValue(session, o, Value.UNKNOWN);
+                        session.getGeneratedKeys().add(table.getColumn(i));
                         newRow.setValue(i, v);
                     }
                 }

h2/src/main/org/h2/server/TcpServerThread.java

@@ -21,6 +21,7 @@ import org.h2.command.Command;
 import org.h2.engine.ConnectionInfo;
 import org.h2.engine.Constants;
 import org.h2.engine.Engine;
+import org.h2.engine.GeneratedKeysMode;
 import org.h2.engine.Session;
 import org.h2.engine.SessionRemote;
 import org.h2.engine.SysProperties;
@@ -31,6 +32,7 @@ import org.h2.jdbc.JdbcSQLException;
 import org.h2.message.DbException;
 import org.h2.result.ResultColumn;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.store.LobStorageInterface;
 import org.h2.util.IOUtils;
 import org.h2.util.SmallLRUCache;
@@ -88,15 +90,15 @@ public class TcpServerThread implements Runnable {
                 }
                 int minClientVersion = transfer.readInt();
                 int maxClientVersion = transfer.readInt();
-                if (maxClientVersion < Constants.TCP_PROTOCOL_VERSION_6) {
+                if (maxClientVersion < Constants.TCP_PROTOCOL_VERSION_MIN_SUPPORTED) {
                     throw DbException.get(ErrorCode.DRIVER_VERSION_ERROR_2,
-                            "" + clientVersion, "" + Constants.TCP_PROTOCOL_VERSION_6);
-                } else if (minClientVersion > Constants.TCP_PROTOCOL_VERSION_16) {
+                            "" + clientVersion, "" + Constants.TCP_PROTOCOL_VERSION_MIN_SUPPORTED);
+                } else if (minClientVersion > Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED) {
                     throw DbException.get(ErrorCode.DRIVER_VERSION_ERROR_2,
-                            "" + clientVersion, "" + Constants.TCP_PROTOCOL_VERSION_16);
+                            "" + clientVersion, "" + Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED);
                 }
-                if (maxClientVersion >= Constants.TCP_PROTOCOL_VERSION_16) {
-                    clientVersion = Constants.TCP_PROTOCOL_VERSION_16;
+                if (maxClientVersion >= Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED) {
+                    clientVersion = Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED;
                 } else {
                     clientVersion = maxClientVersion;
                 }
@@ -178,7 +180,7 @@ public class TcpServerThread implements Runnable {
             RuntimeException closeError = null;
             try {
                 Command rollback = session.prepareLocal("ROLLBACK");
-                rollback.executeUpdate();
+                rollback.executeUpdate(false);
             } catch (RuntimeException e) {
                 closeError = e;
                 server.traceError(e);
@@ -302,7 +304,7 @@ public class TcpServerThread implements Runnable {
                 commit = session.prepareLocal("COMMIT");
             }
             int old = session.getModificationId();
-            commit.executeUpdate();
+            commit.executeUpdate(false);
             transfer.writeInt(getState(old)).flush();
             break;
         }
@@ -353,10 +355,48 @@ public class TcpServerThread implements Runnable {
             int id = transfer.readInt();
             Command command = (Command) cache.getObject(id, false);
             setParameters(command);
+            boolean supportsGeneratedKeys = clientVersion >= Constants.TCP_PROTOCOL_VERSION_17;
+            boolean writeGeneratedKeys = supportsGeneratedKeys;
+            Object generatedKeysRequest;
+            if (supportsGeneratedKeys) {
+                int mode = transfer.readInt();
+                switch (mode) {
+                case GeneratedKeysMode.NONE:
+                    generatedKeysRequest = false;
+                    writeGeneratedKeys = false;
+                    break;
+                case GeneratedKeysMode.AUTO:
+                    generatedKeysRequest = true;
+                    break;
+                case GeneratedKeysMode.COLUMN_NUMBERS: {
+                    int len = transfer.readInt();
+                    int[] keys = new int[len];
+                    for (int i = 0; i < len; i++) {
+                        keys[i] = transfer.readInt();
+                    }
+                    generatedKeysRequest = keys;
+                    break;
+                }
+                case GeneratedKeysMode.COLUMN_NAMES: {
+                    int len = transfer.readInt();
+                    String[] keys = new String[len];
+                    for (int i = 0; i < len; i++) {
+                        keys[i] = transfer.readString();
+                    }
+                    generatedKeysRequest = keys;
+                    break;
+                }
+                default:
+                    throw DbException.get(ErrorCode.CONNECTION_BROKEN_1,
+                            "Unsupported generated keys' mode " + mode);
+                }
+            } else {
+                generatedKeysRequest = false;
+            }
             int old = session.getModificationId();
-            int updateCount;
+            ResultWithGeneratedKeys result;
             synchronized (session) {
-                updateCount = command.executeUpdate();
+                result = command.executeUpdate(generatedKeysRequest);
             }
             int status;
             if (session.isClosed()) {
@@ -365,8 +405,22 @@ public class TcpServerThread implements Runnable {
             } else {
                 status = getState(old);
             }
-            transfer.writeInt(status).writeInt(updateCount).
+            transfer.writeInt(status).writeInt(result.getUpdateCount()).
                     writeBoolean(session.getAutoCommit());
+            if (writeGeneratedKeys) {
+                ResultInterface generatedKeys = result.getGeneratedKeys();
+                int columnCount = generatedKeys.getVisibleColumnCount();
+                transfer.writeInt(columnCount);
+                int rowCount = generatedKeys.getRowCount();
+                transfer.writeInt(rowCount);
+                for (int i = 0; i < columnCount; i++) {
+                    ResultColumn.writeColumn(transfer, generatedKeys, i);
+                }
+                for (int i = 0; i < rowCount; i++) {
+                    sendRow(generatedKeys);
+                }
+                generatedKeys.close();
+            }
             transfer.flush();
             break;
         }

h2/src/main/org/h2/server/web/WebApp.java

@@ -1284,7 +1284,7 @@ public class WebApp {
             ResultSet rs;
             long time = System.currentTimeMillis();
             boolean metadata = false;
-            boolean generatedKeys = false;
+            int generatedKeys = Statement.NO_GENERATED_KEYS;
             boolean edit = false;
             boolean list = false;
             if (isBuiltIn(sql, "@autocommit_true")) {
@@ -1316,7 +1316,7 @@ public class WebApp {
                 sql = sql.substring("@meta".length()).trim();
             }
             if (isBuiltIn(sql, "@generated")) {
-                generatedKeys = true;
+                generatedKeys = Statement.RETURN_GENERATED_KEYS;
                 sql = sql.substring("@generated".length()).trim();
             } else if (isBuiltIn(sql, "@history")) {
                 buff.append(getCommandHistoryString());
@@ -1385,9 +1385,9 @@ public class WebApp {
                 int maxrows = getMaxrows();
                 stat.setMaxRows(maxrows);
                 session.executingStatement = stat;
-                boolean isResultSet = stat.execute(sql);
+                boolean isResultSet = stat.execute(sql, generatedKeys);
                 session.addCommand(sql);
-                if (generatedKeys) {
+                if (generatedKeys == Statement.RETURN_GENERATED_KEYS) {
                     rs = null;
                     rs = stat.getGeneratedKeys();
                 } else {

h2/src/main/org/h2/store/FileLock.java

@@ -210,8 +210,8 @@ public class FileLock implements Runnable {
                     Constants.DEFAULT_TCP_PORT, false);
             Transfer transfer = new Transfer(null, socket);
             transfer.init();
-            transfer.writeInt(Constants.TCP_PROTOCOL_VERSION_6);
-            transfer.writeInt(Constants.TCP_PROTOCOL_VERSION_16);
+            transfer.writeInt(Constants.TCP_PROTOCOL_VERSION_MIN_SUPPORTED);
+            transfer.writeInt(Constants.TCP_PROTOCOL_VERSION_MAX_SUPPORTED);
             transfer.writeString(null);
             transfer.writeString(null);
             transfer.writeString(id);

h2/src/main/org/h2/table/Column.java

@@ -321,6 +321,7 @@ public class Column {
                 value = ValueNull.INSTANCE;
             } else {
                 value = localDefaultExpression.getValue(session).convertTo(type);
+                session.getGeneratedKeys().add(this);
                 if (primaryKey) {
                     session.setLastIdentity(value);
                 }
@@ -330,6 +331,7 @@ public class Column {
         if (value == ValueNull.INSTANCE) {
             if (convertNullToDefault) {
                 value = localDefaultExpression.getValue(session).convertTo(type);
+                session.getGeneratedKeys().add(this);
             }
             if (value == ValueNull.INSTANCE && !nullable) {
                 if (mode.convertInsertNullToZero) {

h2/src/main/org/h2/table/Table.java

@@ -1026,7 +1026,7 @@ public abstract class Table extends SchemaObjectBase {
             boolean beforeAction, boolean rollback) {
         if (triggers != null) {
             for (TriggerObject trigger : triggers) {
-                boolean done = trigger.fireRow(session, oldRow, newRow, beforeAction, rollback);
+                boolean done = trigger.fireRow(session, this, oldRow, newRow, beforeAction, rollback);
                 if (done) {
                     return true;
                 }

h2/src/main/org/h2/util/MergedResultSet.java

+/*
+ * Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
+ * and the EPL 1.0 (http://h2database.com/html/license.html).
+ * Initial Developer: H2 Group
+ */
+package org.h2.util;
+
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.h2.tools.SimpleResultSet;
+
+/**
+ * Merged result set. Used to combine several result sets into one. Merged
+ * result set will contain rows from all appended result sets. Result sets are
+ * not required to have the same lists of columns, but required to have
+ * compatible column definitions, for example, if one result set has a
+ * {@link java.sql.Types#VARCHAR} column {@code NAME} then another results sets
+ * that have {@code NAME} column should also define it with the same type.
+ */
+public final class MergedResultSet {
+    /**
+     * Metadata of a column.
+     */
+    private static final class ColumnInfo {
+        final String name;
+
+        final int type;
+
+        final int precision;
+
+        final int scale;
+
+        /**
+         * Creates metadata.
+         *
+         * @param name
+         *            name of the column
+         * @param type
+         *            type of the column, see {@link java.sql.Types}
+         * @param precision
+         *            precision of the column
+         * @param scale
+         *            scale of the column
+         */
+        ColumnInfo(String name, int type, int precision, int scale) {
+            this.name = name;
+            this.type = type;
+            this.precision = precision;
+            this.scale = scale;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (this == obj)
+                return true;
+            if (obj == null || getClass() != obj.getClass())
+                return false;
+            ColumnInfo other = (ColumnInfo) obj;
+            return name.equals(other.name);
+        }
+
+        @Override
+        public int hashCode() {
+            return name.hashCode();
+        }
+
+    }
+
+    private final ArrayList<Map<ColumnInfo, Object>> data = New.arrayList();
+
+    private final ArrayList<ColumnInfo> columns = New.arrayList();
+
+    /**
+     * Appends a result set.
+     *
+     * @param rs
+     *            result set to append
+     * @throws SQLException
+     *             on SQL exception
+     */
+    public void add(ResultSet rs) throws SQLException {
+        ResultSetMetaData meta = rs.getMetaData();
+        int cols = meta.getColumnCount();
+        if (cols == 0) {
+            return;
+        }
+        ColumnInfo[] info = new ColumnInfo[cols];
+        for (int i = 1; i <= cols; i++) {
+            ColumnInfo ci = new ColumnInfo(meta.getColumnName(i), meta.getColumnType(i), meta.getPrecision(i),
+                    meta.getScale(i));
+            info[i - 1] = ci;
+            if (!columns.contains(ci)) {
+                columns.add(ci);
+            }
+        }
+        while (rs.next()) {
+            if (cols == 1) {
+                data.add(Collections.singletonMap(info[0], rs.getObject(1)));
+            } else {
+                HashMap<ColumnInfo, Object> map = new HashMap<>();
+                for (int i = 1; i <= cols; i++) {
+                    ColumnInfo ci = info[i - 1];
+                    map.put(ci, rs.getObject(i));
+                }
+                data.add(map);
+            }
+        }
+    }
+
+    /**
+     * Returns merged results set.
+     *
+     * @return result set with rows from all appended result sets
+     */
+    public SimpleResultSet getResult() {
+        SimpleResultSet rs = new SimpleResultSet();
+        for (ColumnInfo ci : columns) {
+            rs.addColumn(ci.name, ci.type, ci.precision, ci.scale);
+        }
+        for (Map<ColumnInfo, Object> map : data) {
+            Object[] row = new Object[columns.size()];
+            for (Map.Entry<ColumnInfo, Object> entry : map.entrySet()) {
+                row[columns.indexOf(entry.getKey())] = entry.getValue();
+            }
+            rs.addRow(row);
+        }
+        return rs;
+    }
+
+    @Override
+    public String toString() {
+        return columns + ": " + data.size();
+    }
+
+}

h2/src/test/org/h2/test/TestAll.java

@@ -97,6 +97,7 @@ import org.h2.test.jdbc.TestConnection;
 import org.h2.test.jdbc.TestCustomDataTypesHandler;
 import org.h2.test.jdbc.TestDatabaseEventListener;
 import org.h2.test.jdbc.TestDriver;
+import org.h2.test.jdbc.TestGetGeneratedKeys;
 import org.h2.test.jdbc.TestJavaObject;
 import org.h2.test.jdbc.TestJavaObjectSerializer;
 import org.h2.test.jdbc.TestLimitUpdates;
@@ -814,6 +815,7 @@ kill -9 `jps -l | grep "org.h2.test." | cut -d " " -f 1`
         addTest(new TestPreparedStatement());
         addTest(new TestResultSet());
         addTest(new TestStatement());
+        addTest(new TestGetGeneratedKeys());
         addTest(new TestTransactionIsolation());
         addTest(new TestUpdatableResultSet());
         addTest(new TestZloty());

h2/src/test/org/h2/test/db/TestTriggersConstraints.java

@@ -227,7 +227,7 @@ public class TestTriggersConstraints extends TestBase implements Trigger {
         assertEquals(1, count);
 
         ResultSet gkRs;
-        gkRs = pstat.getGeneratedKeys();
+        gkRs = stat.executeQuery("select scope_identity()");
 
         assertTrue(gkRs.next());
         assertEquals(1, gkRs.getInt(1));

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

@@ -27,6 +27,7 @@ import java.sql.Types;
 import java.util.Calendar;
 import java.util.GregorianCalendar;
 import java.util.UUID;
+
 import org.h2.api.ErrorCode;
 import org.h2.api.Trigger;
 import org.h2.engine.SysProperties;
@@ -83,15 +84,12 @@ public class TestPreparedStatement extends TestBase {
         testOffsetDateTime8(conn);
         testInstant8(conn);
         testArray(conn);
-        testUUIDGeneratedKeys(conn);
         testSetObject(conn);
         testPreparedSubquery(conn);
         testLikeIndex(conn);
         testCasewhen(conn);
         testSubquery(conn);
         testObject(conn);
-        testIdentity(conn);
-        testBatchGeneratedKeys(conn);
         testDataTypes(conn);
         testGetMoreResults(conn);
         testBlob(conn);
@@ -521,21 +519,6 @@ public class TestPreparedStatement extends TestBase {
         stat.execute("drop table test_uuid");
     }
 
-    private void testUUIDGeneratedKeys(Connection conn) throws SQLException {
-        Statement stat = conn.createStatement();
-        stat.execute("CREATE TABLE TEST_UUID(id UUID DEFAULT " +
-                "random_UUID() PRIMARY KEY)");
-        stat.execute("INSERT INTO TEST_UUID() VALUES()");
-        ResultSet rs = stat.getGeneratedKeys();
-        rs.next();
-        byte[] data = rs.getBytes(1);
-        assertEquals(16, data.length);
-        stat.execute("INSERT INTO TEST_UUID VALUES(random_UUID())");
-        rs = stat.getGeneratedKeys();
-        assertFalse(rs.next());
-        stat.execute("DROP TABLE TEST_UUID");
-    }
-
     /**
      * A trigger that creates a sequence value.
      */
@@ -572,12 +555,17 @@ public class TestPreparedStatement extends TestBase {
         stat.execute("create sequence seq start with 1000");
         stat.execute("create trigger test_ins after insert on test call \"" +
                 SequenceTrigger.class.getName() + "\"");
-        stat.execute("insert into test values(null)");
+        stat.execute("insert into test values(null)", Statement.RETURN_GENERATED_KEYS);
         ResultSet rs = stat.getGeneratedKeys();
         rs.next();
+        // Generated key
         assertEquals(1, rs.getLong(1));
         stat.execute("insert into test values(100)");
         rs = stat.getGeneratedKeys();
+        // No generated keys
+        assertFalse(rs.next());
+        // Value from sequence from trigger
+        rs = stat.executeQuery("select scope_identity()");
         rs.next();
         assertEquals(100, rs.getLong(1));
         stat.execute("drop sequence seq");
@@ -1263,83 +1251,6 @@ public class TestPreparedStatement extends TestBase {
 
     }
 
-    private void testIdentity(Connection conn) throws SQLException {
-        Statement stat = conn.createStatement();
-        stat.execute("CREATE SEQUENCE SEQ");
-        stat.execute("CREATE TABLE TEST(ID INT)");
-        PreparedStatement prep;
-        prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)");
-        prep.execute();
-        ResultSet rs = prep.getGeneratedKeys();
-        rs.next();
-        assertEquals(1, rs.getInt(1));
-        assertFalse(rs.next());
-
-        prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)",
-                Statement.RETURN_GENERATED_KEYS);
-        prep.execute();
-        rs = prep.getGeneratedKeys();
-        rs.next();
-        assertEquals(2, rs.getInt(1));
-        assertFalse(rs.next());
-
-        prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)",
-                new int[] { 1 });
-        prep.execute();
-        rs = prep.getGeneratedKeys();
-        rs.next();
-        assertEquals(3, rs.getInt(1));
-        assertFalse(rs.next());
-
-        prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)",
-                new String[] { "ID" });
-        prep.execute();
-        rs = prep.getGeneratedKeys();
-        rs.next();
-        assertEquals(4, rs.getInt(1));
-        assertFalse(rs.next());
-
-        prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)",
-                ResultSet.TYPE_FORWARD_ONLY,
-                ResultSet.CONCUR_READ_ONLY,
-                ResultSet.HOLD_CURSORS_OVER_COMMIT);
-        prep.execute();
-        rs = prep.getGeneratedKeys();
-        rs.next();
-        assertEquals(5, rs.getInt(1));
-        assertFalse(rs.next());
-
-        stat.execute("DROP TABLE TEST");
-        stat.execute("DROP SEQUENCE SEQ");
-    }
-
-    private void testBatchGeneratedKeys(Connection conn) throws SQLException {
-        Statement stat = conn.createStatement();
-        stat.execute("CREATE SEQUENCE SEQ");
-        stat.execute("CREATE TABLE TEST(ID INT)");
-        PreparedStatement prep = conn.prepareStatement(
-                "INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)");
-        prep.addBatch();
-        prep.addBatch();
-        prep.addBatch();
-        prep.executeBatch();
-        ResultSet keys = prep.getGeneratedKeys();
-        keys.next();
-        assertEquals(1, keys.getLong(1));
-        keys.next();
-        assertEquals(2, keys.getLong(1));
-        keys.next();
-        assertEquals(3, keys.getLong(1));
-        assertFalse(keys.next());
-        stat.execute("DROP TABLE TEST");
-        stat.execute("DROP SEQUENCE SEQ");
-    }
-
     private int getLength() {
         return getSize(LOB_SIZE, LOB_SIZE_BIG);
     }

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

@@ -573,12 +573,14 @@ public class TestResultSet extends TestBase {
         ResultSet rs;
         stat.execute("CREATE TABLE TEST(ID IDENTITY NOT NULL, NAME VARCHAR NULL)");
 
-        stat.execute("INSERT INTO TEST(NAME) VALUES('Hello')");
+        stat.execute("INSERT INTO TEST(NAME) VALUES('Hello')",
+                Statement.RETURN_GENERATED_KEYS);
         rs = stat.getGeneratedKeys();
         assertTrue(rs.next());
         assertEquals(1, rs.getInt(1));
 
-        stat.execute("INSERT INTO TEST(NAME) VALUES('World')");
+        stat.execute("INSERT INTO TEST(NAME) VALUES('World')",
+                Statement.RETURN_GENERATED_KEYS);
         rs = stat.getGeneratedKeys();
         assertTrue(rs.next());
         assertEquals(2, rs.getInt(1));

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

@@ -351,16 +351,19 @@ public class TestStatement extends TestBase {
         stat.execute("create table test1(id identity, x int)");
         stat.execute("drop table if exists test2");
         stat.execute("create table test2(id identity, x int)");
-        stat.execute("merge into test1(x) key(x) values(5)");
+        stat.execute("merge into test1(x) key(x) values(5)",
+                Statement.RETURN_GENERATED_KEYS);
         ResultSet keys;
         keys = stat.getGeneratedKeys();
         keys.next();
         assertEquals(1, keys.getInt(1));
         stat.execute("insert into test2(x) values(10), (11), (12)");
-        stat.execute("merge into test1(x) key(x) values(5)");
+        stat.execute("merge into test1(x) key(x) values(5)",
+                Statement.RETURN_GENERATED_KEYS);
         keys = stat.getGeneratedKeys();
         assertFalse(keys.next());
-        stat.execute("merge into test1(x) key(x) values(6)");
+        stat.execute("merge into test1(x) key(x) values(6)",
+                Statement.RETURN_GENERATED_KEYS);
         keys = stat.getGeneratedKeys();
         keys.next();
         assertEquals(2, keys.getInt(1));
@@ -371,7 +374,8 @@ public class TestStatement extends TestBase {
         Statement stat = conn.createStatement();
         stat.execute("CREATE SEQUENCE SEQ");
         stat.execute("CREATE TABLE TEST(ID INT)");
-        stat.execute("INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)");
+        stat.execute("INSERT INTO TEST VALUES(NEXT VALUE FOR SEQ)",
+                Statement.RETURN_GENERATED_KEYS);
         ResultSet rs = stat.getGeneratedKeys();
         rs.next();
         assertEquals(1, rs.getInt(1));

h2/src/test/org/h2/test/server/TestWeb.java

@@ -447,7 +447,7 @@ public class TestWeb extends TestBase {
             assertContains(result, "There is currently no running statement");
             result = client.get(url,
                     "query.do?sql=@generated insert into test(id) values(test_sequence.nextval)");
-            assertContains(result, "SCOPE_IDENTITY()");
+            assertContains(result, "<tr><th>ID</th></tr><tr><td>1</td></tr>");
             result = client.get(url, "query.do?sql=@maxrows 2000");
             assertContains(result, "Max rowcount is set");
             result = client.get(url, "query.do?sql=@password_hash user password");