Set requested generated keys to client together with result of command

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

@@ -14,7 +14,9 @@ import org.h2.engine.Session;
 import org.h2.expression.ParameterInterface;
 import org.h2.message.DbException;
 import org.h2.message.Trace;
+import org.h2.result.LocalResult;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.util.MathUtils;
 
 /**
@@ -238,7 +240,7 @@ public abstract class Command implements CommandInterface {
     }
 
     @Override
-    public int executeUpdate(Object generatedKeysRequest) {
+    public ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest) {
         long start = 0;
         Database database = session.getDatabase();
         Object sync = database.isMultiThreaded() ? (Object) session : (Object) database;
@@ -257,7 +259,14 @@ public abstract class Command implements CommandInterface {
                 while (true) {
                     database.checkPowerOff();
                     try {
-                        return update();
+                        int updateCount = update();
+                        if (!Boolean.FALSE.equals(generatedKeysRequest)) {
+                            return new ResultWithGeneratedKeys.WithKeys(updateCount,
+                                    LocalResult.read(session,
+                                            session.getGeneratedKeys().getKeys(),
+                                            Integer.MAX_VALUE));
+                        }
+                        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.
@@ -519,7 +520,7 @@ public interface CommandInterface {
      *
      * @return the update count
      */
-    int executeUpdate(Object generatedKeysRequest);
+    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(false);
+        int updateCount = command.executeUpdate(false).getUpdateCount();
         executeRemaining();
         return updateCount;
     }

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

@@ -17,6 +17,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 +195,14 @@ public class CommandRemote implements CommandInterface {
     }
 
     @Override
-    public int executeUpdate(Object generatedKeys) {
+    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,20 +211,21 @@ public class CommandRemote implements CommandInterface {
                     session.traceOperation("COMMAND_EXECUTE_UPDATE", id);
                     transfer.writeInt(SessionRemote.COMMAND_EXECUTE_UPDATE).writeInt(id);
                     sendParameters(transfer);
-                    if (session.getClientVersion() >= Constants.TCP_PROTOCOL_VERSION_17) {
-                        if (Boolean.FALSE.equals(generatedKeys)) {
+                    if (supportsGeneratedKeys) {
+                        if (Boolean.FALSE.equals(generatedKeysRequest)) {
                             transfer.writeInt(0);
-                        } else if (Boolean.TRUE.equals(generatedKeys)) {
+                            readGeneratedKeys = false;
+                        } else if (Boolean.TRUE.equals(generatedKeysRequest)) {
                             transfer.writeInt(1);
-                        } else if (generatedKeys instanceof int[]) {
-                            int[] keys = (int[]) generatedKeys;
+                        } else if (generatedKeysRequest instanceof int[]) {
+                            int[] keys = (int[]) generatedKeysRequest;
                             transfer.writeInt(2);
                             transfer.writeInt(keys.length);
                             for (int key : keys) {
                                 transfer.writeInt(key);
                             }
-                        } else if (generatedKeys instanceof String[]) {
-                            String[] keys = (String[]) generatedKeys;
+                        } else if (generatedKeysRequest instanceof String[]) {
+                            String[] keys = (String[]) generatedKeysRequest;
                             transfer.writeInt(3);
                             transfer.writeInt(keys.length);
                             for (String key : keys) {
@@ -227,11 +233,20 @@ public class CommandRemote implements CommandInterface {
                             }
                         } else {
                             transfer.writeInt(0);
+                            readGeneratedKeys = false;
                         }
                     }
                     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);
                 }
@@ -239,7 +254,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/engine/Session.java

@@ -1254,8 +1254,8 @@ public class Session extends SessionWithState {
      */
     public void setCurrentCommand(Command command, Object generatedKeysRequest) {
         this.currentCommand = command;
-        // Preserve generated keys in case of a new query so they can be read with
-        // CALL GET_GENERATED_KEYS()
+        // Preserve generated keys in case of a new query due to possible nested
+        // queries in update
         if (command != null && !command.isQuery()) {
             getGeneratedKeys().clear(generatedKeysRequest);
         }
@@ -1809,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

@@ -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/expression/Function.java

@@ -115,7 +115,7 @@ public class Function extends Expression implements FunctionCall {
     public static final int DATABASE = 150, USER = 151, CURRENT_USER = 152,
             IDENTITY = 153, SCOPE_IDENTITY = 154, AUTOCOMMIT = 155,
             READONLY = 156, DATABASE_PATH = 157, LOCK_TIMEOUT = 158,
-            DISK_SPACE_USED = 159, SIGNAL = 160, GET_GENERATED_KEYS = 161;
+            DISK_SPACE_USED = 159, SIGNAL = 160;
 
     private static final Pattern SIGNAL_PATTERN = Pattern.compile("[0-9A-Z]{5}");
 
@@ -491,7 +491,6 @@ public class Function extends Expression implements FunctionCall {
         addFunctionNotDeterministic("DISK_SPACE_USED", DISK_SPACE_USED,
                 1, Value.LONG);
         addFunctionWithNull("SIGNAL", SIGNAL, 2, Value.NULL);
-        addFunction("GET_GENERATED_KEYS", GET_GENERATED_KEYS, 0, Value.RESULT_SET);
         addFunction("H2VERSION", H2VERSION, 0, Value.STRING);
 
         // TableFunction
@@ -925,9 +924,6 @@ public class Function extends Expression implements FunctionCall {
         case DISK_SPACE_USED:
             result = ValueLong.get(getDiskSpaceUsed(session, v0));
             break;
-        case GET_GENERATED_KEYS:
-            result = ValueResultSet.get(session.getGeneratedKeys().getKeys());
-            break;
         case CAST:
         case CONVERT: {
             v0 = v0.convertTo(dataType);

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

@@ -85,7 +85,6 @@ public class JdbcConnection extends TraceObject
     private CommandInterface getReadOnly, getGeneratedKeys;
     private CommandInterface setLockMode, getLockMode;
     private CommandInterface setQueryTimeout, getQueryTimeout;
-    private boolean oldGetGeneratedKeys;
 
     private int savepointId;
     private String catalog;
@@ -1583,21 +1582,6 @@ public class JdbcConnection extends TraceObject
      * INTERNAL
      */
     ResultSet getGeneratedKeys(JdbcStatement stat, int id) {
-        if (!oldGetGeneratedKeys) {
-            try {
-                getGeneratedKeys = prepareCommand("CALL GET_GENERATED_KEYS()", getGeneratedKeys);
-                ResultInterface result = getGeneratedKeys.executeQuery(Integer.MAX_VALUE, true);
-                ResultSet rs = new JdbcResultSet(this, stat, getGeneratedKeys, result,
-                        id, false, true, false);
-                return rs;
-            } catch (DbException ex) {
-                if (ex.getErrorCode() == ErrorCode.FUNCTION_NOT_FOUND_1) {
-                    oldGetGeneratedKeys = true;
-                } else {
-                    throw ex;
-                }
-            }
-        }
         getGeneratedKeys = prepareCommand(
                 "SELECT SCOPE_IDENTITY() "
                         + "WHERE SCOPE_IDENTITY() IS NOT NULL",

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

@@ -31,6 +31,7 @@ import org.h2.expression.ParameterInterface;
 import org.h2.message.DbException;
 import org.h2.message.TraceObject;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
 import org.h2.util.DateTimeUtils;
 import org.h2.util.IOUtils;
 import org.h2.util.MergedResultSet;
@@ -195,7 +196,14 @@ public class JdbcPreparedStatement extends JdbcStatement implements
         synchronized (session) {
             try {
                 setExecutingStatement(command);
-                updateCount = command.executeUpdate(generatedKeysRequest);
+                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);
             }
@@ -238,7 +246,13 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                                     updatable, cachedColumnLabelMap);
                         } else {
                             returnsResultSet = false;
-                            updateCount = command.executeUpdate(generatedKeysRequest);
+                            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) {
@@ -1263,7 +1277,8 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                     }
                     try {
                         result[i] = executeUpdateInternal();
-                        ResultSet rs = conn.getGeneratedKeys(this, id);
+                        // Cannot use own implementation, it returns batch identities
+                        ResultSet rs = super.getGeneratedKeys();
                         batchIdentities.add(rs);
                     } catch (Exception re) {
                         SQLException e = logAndConvert(re);

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

@@ -18,6 +18,8 @@ import org.h2.engine.SysProperties;
 import org.h2.message.DbException;
 import org.h2.message.TraceObject;
 import org.h2.result.ResultInterface;
+import org.h2.result.ResultWithGeneratedKeys;
+import org.h2.tools.SimpleResultSet;
 import org.h2.util.New;
 import org.h2.util.ParserUtil;
 import org.h2.util.StringUtils;
@@ -33,6 +35,7 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
     protected int maxRows;
     protected int fetchSize = SysProperties.SERVER_RESULT_SET_FETCH_SIZE;
     protected int updateCount;
+    protected JdbcResultSet generatedKeys;
     protected final int resultSetType;
     protected final int resultSetConcurrency;
     protected final boolean closedByResultSet;
@@ -163,7 +166,14 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
             synchronized (session) {
                 setExecutingStatement(command);
                 try {
-                    updateCount = command.executeUpdate(generatedKeysRequest);
+                    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);
                 }
@@ -219,7 +229,13 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
                                 closedByResultSet, scrollable, updatable);
                     } else {
                         returnsResultSet = false;
-                        updateCount = command.executeUpdate(generatedKeysRequest);
+                        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) {
@@ -816,6 +832,13 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
                 debugCodeAssign("ResultSet", TraceObject.RESULT_SET, id, "getGeneratedKeys()");
             }
             checkClosed();
+            if (generatedKeys != null) {
+                return generatedKeys;
+            }
+            if (session.isSupportsGeneratedKeys()) {
+                return new SimpleResultSet();
+            }
+            // Old server, so use SCOPE_IDENTITY()
             return conn.getGeneratedKeys(this, id);
         } catch (Exception e) {
             throw logAndConvert(e);
@@ -1197,11 +1220,15 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
                 if (resultSet != null) {
                     resultSet.closeInternal();
                 }
+                if (generatedKeys != null) {
+                    generatedKeys.closeInternal();
+                }
             }
         } finally {
             cancelled = false;
             resultSet = null;
             updateCount = -1;
+            generatedKeys = null;
         }
     }
 

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/server/TcpServerThread.java

@@ -31,6 +31,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;
@@ -353,12 +354,15 @@ 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 (clientVersion >= Constants.TCP_PROTOCOL_VERSION_17) {
+            if (supportsGeneratedKeys) {
                 int type = transfer.readInt();
                 switch (type) {
                 default:
                     generatedKeysRequest = false;
+                    writeGeneratedKeys = false;
                     break;
                 case 1:
                     generatedKeysRequest = true;
@@ -382,12 +386,12 @@ public class TcpServerThread implements Runnable {
                 }
                 }
             } else {
-                generatedKeysRequest = true;
+                generatedKeysRequest = false;
             }
             int old = session.getModificationId();
-            int updateCount;
+            ResultWithGeneratedKeys result;
             synchronized (session) {
-                updateCount = command.executeUpdate(generatedKeysRequest);
+                result = command.executeUpdate(generatedKeysRequest);
             }
             int status;
             if (session.isClosed()) {
@@ -396,8 +400,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/test/org/h2/test/jdbc/TestPreparedStatement.java

@@ -1355,7 +1355,7 @@ public class TestPreparedStatement extends TestBase {
                 "insert into test values (30), (next value for seq),"
                         + " (next value for seq), (next value for seq), (20)",
                 PreparedStatement.RETURN_GENERATED_KEYS);
-            prep.executeUpdate();
+        prep.executeUpdate();
         ResultSet rs = prep.getGeneratedKeys();
         rs.next();
         assertEquals(1L, rs.getLong(1));