Pass the whole inserted row to GeneratedKeys to ensure that latest values will be returned

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

@@ -166,7 +166,7 @@ public class Insert extends Prepared implements ResultTarget {
                             Value v = c.convert(e.getValue(session), session.getDatabase().getMode());
                             newRow.setValue(index, v);
                             if (e instanceof SequenceValue) {
-                                generatedKeys.add(c, v);
+                                generatedKeys.add(c);
                             }
                         } catch (DbException ex) {
                             throw setRow(ex, x, getSQL(expr));
@@ -187,7 +187,7 @@ public class Insert extends Prepared implements ResultTarget {
                             continue;
                         }
                     }
-                    generatedKeys.confirmRow();
+                    generatedKeys.confirmRow(newRow);
                     session.log(table, UndoLogRecord.INSERT, newRow);
                     table.fireAfterRow(session, null, newRow, false);
                 }

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

@@ -86,10 +86,10 @@ 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.size() > 0) {
             // process values in list
             count = 0;
-            GeneratedKeys generatedKeys = session.getGeneratedKeys();
             generatedKeys.initialize(targetTable);
             for (int x = 0, size = valuesExpressionList.size(); x < size; x++) {
                 setCurrentRowNumber(x + 1);
@@ -106,7 +106,7 @@ public class Merge extends Prepared {
                             Value v = c.convert(e.getValue(session));
                             newRow.setValue(index, v);
                             if (e instanceof SequenceValue) {
-                                generatedKeys.add(c, v);
+                                generatedKeys.add(c);
                             }
                         } catch (DbException ex) {
                             throw setRow(ex, count, getSQL(expr));
@@ -124,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);
@@ -179,7 +180,7 @@ public class Merge extends Prepared {
                 if (!done) {
                     targetTable.lock(session, true, false);
                     targetTable.addRow(session, row);
-                    session.getGeneratedKeys().confirmRow();
+                    session.getGeneratedKeys().confirmRow(row);
                     session.log(targetTable, UndoLogRecord.INSERT, row);
                     targetTable.fireAfterRow(session, null, row, false);
                 }

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

@@ -11,67 +11,102 @@ import java.util.HashMap;
 import java.util.Map;
 
 import org.h2.message.DbException;
+import org.h2.result.Row;
 import org.h2.table.Column;
 import org.h2.table.Table;
 import org.h2.tools.SimpleResultSet;
 import org.h2.util.MathUtils;
 import org.h2.util.New;
 import org.h2.value.DataType;
-import org.h2.value.Value;
 
 /**
- * Generated keys.
+ * Class for gathering and processing of generated keys.
  */
 public final class GeneratedKeys {
+    /**
+     * Data for result set with generated keys.
+     */
     private final ArrayList<Map<Column, Object>> data = New.arrayList();
 
-    private final ArrayList<Object> row = New.arrayList();
+    /**
+     * Columns with generated keys in the current row.
+     */
+    private final ArrayList<Column> row = New.arrayList();
 
-    private final ArrayList<Column> columns = 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;
 
-    public void add(Column column, Value value) {
+    /**
+     * Remembers columns with generated keys.
+     *
+     * @param column
+     *            table column
+     */
+    public void add(Column column) {
         if (Boolean.FALSE.equals(generatedKeysRequest)) {
             return;
         }
         row.add(column);
-        row.add(value.getObject());
     }
 
+    /**
+     * 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();
-        columns.clear();
+        allColumns.clear();
         table = null;
     }
 
-    public void initialize(Table table) {
-        this.table = table;
-    }
-
-    public void confirmRow() {
+    /**
+     * 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 == 2) {
-                Column column = (Column) row.get(0);
-                data.add(Collections.singletonMap(column, row.get(1)));
-                if (!columns.contains(column)) {
-                    columns.add(column);
+            if (size == 1) {
+                Column column = row.get(0);
+                data.add(Collections.singletonMap(column, tableRow.getValue(column.getColumnId()).getObject()));
+                if (!allColumns.contains(column)) {
+                    allColumns.add(column);
                 }
             } else {
                 HashMap<Column, Object> map = new HashMap<>();
-                for (int i = 0; i < size; i += 2) {
-                    Column column = (Column) row.get(i);
-                    map.put(column, row.get(i + 1));
-                    if (!columns.contains(column)) {
-                        columns.add(column);
+                for (Column column : row) {
+                    map.put(column, tableRow.getValue(column.getColumnId()).getObject());
+                    if (!allColumns.contains(column)) {
+                        allColumns.add(column);
                     }
                 }
                 data.add(map);
@@ -80,13 +115,18 @@ public final class GeneratedKeys {
         }
     }
 
+    /**
+     * Returns generated keys.
+     *
+     * @return result set with generated keys
+     */
     public SimpleResultSet getKeys() {
         SimpleResultSet rs = new SimpleResultSet();
         if (Boolean.FALSE.equals(generatedKeysRequest)) {
             return rs;
         }
         if (Boolean.TRUE.equals(generatedKeysRequest)) {
-            for (Column column : columns) {
+            for (Column column : allColumns) {
                 rs.addColumn(column.getName(), DataType.convertTypeToSQLType(column.getType()),
                         MathUtils.convertLongToInt(column.getPrecision()), column.getScale());
             }
@@ -95,13 +135,13 @@ public final class GeneratedKeys {
                 int[] indices = (int[]) generatedKeysRequest;
                 Column[] columns = table.getColumns();
                 int cnt = columns.length;
-                this.columns.clear();
+                this.allColumns.clear();
                 for (int idx : indices) {
                     if (idx >= 1 && idx <= cnt) {
                         Column column = columns[idx - 1];
                         rs.addColumn(column.getName(), DataType.convertTypeToSQLType(column.getType()),
                                 MathUtils.convertLongToInt(column.getPrecision()), column.getScale());
-                        this.columns.add(column);
+                        this.allColumns.add(column);
                     }
                 }
             } else {
@@ -110,13 +150,13 @@ public final class GeneratedKeys {
         } else if (generatedKeysRequest instanceof String[]) {
             if (table != null) {
                 String[] names = (String[]) generatedKeysRequest;
-                this.columns.clear();
+                this.allColumns.clear();
                 for (String name : names) {
                     try {
                         Column column = table.getColumn(name);
                         rs.addColumn(column.getName(), DataType.convertTypeToSQLType(column.getType()),
                                 MathUtils.convertLongToInt(column.getPrecision()), column.getScale());
-                        this.columns.add(column);
+                        this.allColumns.add(column);
                     } catch (DbException e) {
                     }
                 }
@@ -130,9 +170,9 @@ public final class GeneratedKeys {
             return rs;
         }
         for (Map<Column, Object> map : data) {
-            Object[] row = new Object[columns.size()];
+            Object[] row = new Object[allColumns.size()];
             for (Map.Entry<Column, Object> entry : map.entrySet()) {
-                int idx = columns.indexOf(entry.getKey());
+                int idx = allColumns.indexOf(entry.getKey());
                 if (idx >= 0) {
                     row[idx] = entry.getValue();
                 }
@@ -142,13 +182,29 @@ public final class GeneratedKeys {
         return rs;
     }
 
+    /**
+     * 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 columns + ": " + data.size();
+        return allColumns + ": " + data.size();
     }
 
 }

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

@@ -261,7 +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), v);
+                        session.getGeneratedKeys().add(table.getColumn(i));
                         newRow.setValue(i, v);
                     }
                 }

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

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