Do not use SimpleResultSet for generated keys

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

@@ -1592,7 +1592,7 @@ public class JdbcConnection extends TraceObject
     /**
      * INTERNAL
      */
-    ResultSet getGeneratedKeys(JdbcStatement stat, int id) {
+    JdbcResultSet getGeneratedKeys(JdbcStatement stat, int id) {
         getGeneratedKeys = prepareCommand(
                 "SELECT SCOPE_IDENTITY() "
                         + "WHERE SCOPE_IDENTITY() IS NOT NULL",

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

@@ -31,11 +31,11 @@ import org.h2.command.CommandInterface;
 import org.h2.expression.ParameterInterface;
 import org.h2.message.DbException;
 import org.h2.message.TraceObject;
+import org.h2.result.MergedResult;
 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;
 import org.h2.util.Utils;
 import org.h2.value.DataType;
 import org.h2.value.Value;
@@ -63,7 +63,7 @@ public class JdbcPreparedStatement extends JdbcStatement implements
     protected CommandInterface command;
     private final String sqlStatement;
     private ArrayList<Value[]> batchParameters;
-    private MergedResultSet batchIdentities;
+    private MergedResult batchIdentities;
     private HashMap<String, Integer> cachedColumnLabelMap;
     private final Object generatedKeysRequest;
 
@@ -1259,7 +1259,7 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                 // Empty batch is allowed, see JDK-4639504 and other issues
                 batchParameters = Utils.newSmallArrayList();
             }
-            batchIdentities = new MergedResultSet();
+            batchIdentities = new MergedResult();
             int size = batchParameters.size();
             int[] result = new int[size];
             boolean error = false;
@@ -1279,7 +1279,7 @@ public class JdbcPreparedStatement extends JdbcStatement implements
                         result[i] = executeUpdateInternal();
                         // Cannot use own implementation, it returns batch identities
                         ResultSet rs = super.getGeneratedKeys();
-                        batchIdentities.add(rs);
+                        batchIdentities.add(((JdbcResultSet) rs).result);
                     } catch (Exception re) {
                         SQLException e = logAndConvert(re);
                         if (next == null) {
@@ -1308,7 +1308,11 @@ public class JdbcPreparedStatement extends JdbcStatement implements
     @Override
     public ResultSet getGeneratedKeys() throws SQLException {
         if (batchIdentities != null) {
-            return batchIdentities.getResult();
+            int id = getNextId(TraceObject.RESULT_SET);
+            if (isDebugEnabled()) {
+                debugCodeAssign("ResultSet", TraceObject.RESULT_SET, id, "getGeneratedKeys()");
+            }
+            generatedKeys = new JdbcResultSet(conn, this, null, batchIdentities.getResult(), id, false, true, false);
         }
         return super.getGeneratedKeys();
     }

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

@@ -84,7 +84,7 @@ public class JdbcResultSet extends TraceObject implements ResultSet, JdbcResultS
     private final boolean closeStatement;
     private final boolean scrollable;
     private final boolean updatable;
-    private ResultInterface result;
+    ResultInterface result;
     private JdbcConnection conn;
     private JdbcStatement stat;
     private int columnCount;

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

@@ -19,7 +19,7 @@ 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.result.SimpleResult;
 import org.h2.util.ParserUtil;
 import org.h2.util.StringUtils;
 import org.h2.util.Utils;
@@ -866,21 +866,20 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
     @Override
     public ResultSet getGeneratedKeys() throws SQLException {
         try {
-            int id = getNextId(TraceObject.RESULT_SET);
+            int id = generatedKeys != null ? generatedKeys.getTraceId() : getNextId(TraceObject.RESULT_SET);
             if (isDebugEnabled()) {
                 debugCodeAssign("ResultSet", TraceObject.RESULT_SET, id, "getGeneratedKeys()");
             }
             checkClosed();
-            if (!conn.scopeGeneratedKeys()) {
-                if (generatedKeys != null) {
-                    return generatedKeys;
-                }
-                if (session.isSupportsGeneratedKeys()) {
-                    return new SimpleResultSet();
+            if (generatedKeys == null) {
+                if (!conn.scopeGeneratedKeys() && session.isSupportsGeneratedKeys()) {
+                    generatedKeys = new JdbcResultSet(conn, this, null, new SimpleResult(), id, false, true, false);
+                } else {
+                    // Compatibility mode or an old server, so use SCOPE_IDENTITY()
+                    generatedKeys = conn.getGeneratedKeys(this, id);
                 }
             }
-            // Compatibility mode or an old server, so use SCOPE_IDENTITY()
-            return conn.getGeneratedKeys(this, id);
+            return generatedKeys;
         } catch (Exception e) {
             throw logAndConvert(e);
         }

h2/src/main/org/h2/result/MergedResult.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;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.h2.util.Utils;
+import org.h2.value.Value;
+
+/**
+ * Merged result. Used to combine several results into one. Merged result will
+ * contain rows from all appended results. Results are not required to have the
+ * same lists of columns, but required to have compatible column definitions,
+ * for example, if one result has a {@link java.sql.Types#VARCHAR} column
+ * {@code NAME} then another results that have {@code NAME} column should also
+ * define it with the same type.
+ */
+public final class MergedResult {
+    private final ArrayList<Map<SimpleResult.Column, Value>> data = Utils.newSmallArrayList();
+
+    private final ArrayList<SimpleResult.Column> columns = Utils.newSmallArrayList();
+
+    /**
+     * Appends a result.
+     *
+     * @param result
+     *            result to append
+     */
+    public void add(ResultInterface result) {
+        int count = result.getVisibleColumnCount();
+        if (count == 0) {
+            return;
+        }
+        SimpleResult.Column[] cols = new SimpleResult.Column[count];
+        for (int i = 0; i < count; i++) {
+            SimpleResult.Column c = new SimpleResult.Column(result.getAlias(i), result.getColumnName(i),
+                    result.getColumnType(i), result.getColumnPrecision(i), result.getColumnScale(i),
+                    result.getDisplaySize(i));
+            cols[i] = c;
+            if (!columns.contains(c)) {
+                columns.add(c);
+            }
+        }
+        while (result.next()) {
+            if (count == 1) {
+                data.add(Collections.singletonMap(cols[0], result.currentRow()[0]));
+            } else {
+                HashMap<SimpleResult.Column, Value> map = new HashMap<>();
+                for (int i = 0; i < count; i++) {
+                    SimpleResult.Column ci = cols[i];
+                    map.put(ci, result.currentRow()[i]);
+                }
+                data.add(map);
+            }
+        }
+    }
+
+    /**
+     * Returns merged results.
+     *
+     * @return result with rows from all appended result sets
+     */
+    public SimpleResult getResult() {
+        SimpleResult result = new SimpleResult();
+        for (SimpleResult.Column c : columns) {
+            result.addColumn(c);
+        }
+        for (Map<SimpleResult.Column, Value> map : data) {
+            Value[] row = new Value[columns.size()];
+            for (Map.Entry<SimpleResult.Column, Value> entry : map.entrySet()) {
+                row[columns.indexOf(entry.getKey())] = entry.getValue();
+            }
+            result.addRow(row);
+        }
+        return result;
+    }
+
+    @Override
+    public String toString() {
+        return columns + ": " + data.size();
+    }
+
+}

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

@@ -17,7 +17,7 @@ import org.h2.value.Value;
  */
 public class SimpleResult implements ResultInterface {
 
-    private static final class Column {
+    static final class Column {
 
         final String alias;
 
@@ -33,6 +33,9 @@ public class SimpleResult implements ResultInterface {
 
         Column(String alias, String columnName, int columnType, long columnPrecision, int columnScale,
                 int displaySize) {
+            if (alias == null || columnName == null) {
+                throw new NullPointerException();
+            }
             this.alias = alias;
             this.columnName = columnName;
             this.columnType = columnType;
@@ -41,6 +44,33 @@ public class SimpleResult implements ResultInterface {
             this.displaySize = displaySize;
         }
 
+        @Override
+        public int hashCode() {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + alias.hashCode();
+            result = prime * result + columnName.hashCode();
+            return result;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (this == obj)
+                return true;
+            if (obj == null || getClass() != obj.getClass())
+                return false;
+            Column other = (Column) obj;
+            return alias.equals(other.alias) && columnName.equals(other.columnName);
+        }
+
+        @Override
+        public String toString() {
+            if (alias.equals(columnName)) {
+                return columnName;
+            }
+            return columnName + ' ' + alias;
+        }
+
     }
 
     private final ArrayList<Column> columns;
@@ -66,8 +96,12 @@ public class SimpleResult implements ResultInterface {
 
     public void addColumn(String alias, String columnName, int columnType, long columnPrecision, int columnScale,
             int displaySize) {
+        addColumn(new Column(alias, columnName, columnType, columnPrecision, columnScale, displaySize));
+    }
+
+    void addColumn(Column column) {
         assert rows.isEmpty();
-        columns.add(new Column(alias, columnName, columnType, columnPrecision, columnScale, displaySize));
+        columns.add(column);
     }
 
     public void addRow(Value... values) {

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 {
-    private final ArrayList<Map<SimpleColumnInfo, Object>> data = Utils.newSmallArrayList();
-
-    private final ArrayList<SimpleColumnInfo> columns = Utils.newSmallArrayList();
-
-    /**
-     * 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;
-        }
-        SimpleColumnInfo[] info = new SimpleColumnInfo[cols];
-        for (int i = 1; i <= cols; i++) {
-            SimpleColumnInfo ci = new SimpleColumnInfo(meta.getColumnName(i), meta.getColumnType(i),
-                    meta.getColumnTypeName(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<SimpleColumnInfo, Object> map = new HashMap<>();
-                for (int i = 1; i <= cols; i++) {
-                    SimpleColumnInfo 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 (SimpleColumnInfo ci : columns) {
-            rs.addColumn(ci.name, ci.type, ci.typeName, ci.precision, ci.scale);
-        }
-        for (Map<SimpleColumnInfo, Object> map : data) {
-            Object[] row = new Object[columns.size()];
-            for (Map.Entry<SimpleColumnInfo, Object> entry : map.entrySet()) {
-                row[columns.indexOf(entry.getKey())] = entry.getValue();
-            }
-            rs.addRow(row);
-        }
-        return rs;
-    }
-
-    @Override
-    public String toString() {
-        return columns + ": " + data.size();
-    }
-
-}