diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmt.java b/ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmt.java index 0a14d485a..6a1a8afd2 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmt.java +++ b/ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmt.java @@ -6,8 +6,6 @@ import io.ebeaninternal.api.SpiTransaction; import java.io.IOException; import java.io.InputStream; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Method; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; @@ -23,18 +21,12 @@ import java.util.List; */ public final class BatchedPstmt implements SpiProfileTransactionEvent { - /** - * The underlying statement. - */ + private static final DB2GetKeys DB2_HACK = new DB2GetKeys(); + private PreparedStatement pstmt; - - /** - * True if an insert that uses generated keys. - */ private final boolean isGenKeys; - /** - * The list of BatchPostExecute used to perform post processing. + * The list of BatchPostExecute used to perform post-processing. */ private final List list = new ArrayList<>(); private final String sql; @@ -44,18 +36,6 @@ public final class BatchedPstmt implements SpiProfileTransactionEvent { private int[] results; private List inputStreams; - private static Class DB2_PREPARED_STATEMENT; - private static Method GET_DB_GENERATED_KEYS; - - static { - try { - DB2_PREPARED_STATEMENT = (Class) Class.forName("com.ibm.db2.jcc.DB2PreparedStatement"); - GET_DB_GENERATED_KEYS = DB2_PREPARED_STATEMENT.getDeclaredMethod("getDBGeneratedKeys"); - } catch (ClassNotFoundException | NoSuchMethodException | SecurityException cnf) { - // NOP - } - - } /** * Create with a given statement. */ @@ -183,35 +163,11 @@ public final class BatchedPstmt implements SpiProfileTransactionEvent { } } - protected void getGeneratedKeys() throws SQLException { - int index = 0; - if (DB2_PREPARED_STATEMENT != null) { - PreparedStatement db2Stmt = null; - if (DB2_PREPARED_STATEMENT.isInstance(pstmt)) { - db2Stmt = pstmt; - } else if (pstmt.isWrapperFor(DB2_PREPARED_STATEMENT)) { - db2Stmt = pstmt.unwrap(DB2_PREPARED_STATEMENT); - } - if (db2Stmt != null) { - // WTF: https://stackoverflow.com/questions/41725492/how-to-get-auto-generated-keys-of-batch-insert-statement - ResultSet[] result; - try { - result = (ResultSet[]) GET_DB_GENERATED_KEYS.invoke(db2Stmt); - } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) { - throw new SQLException("Could not get generated keys for DB2", e); - } - for (int i = 0; i < result.length; i++) { - while (result[i].next()) { - ResultSet rset = result[i]; - Object idValue = rset.getObject(1); - list.get(index).setGeneratedKey(idValue); - index++; - } - } - return; - } + private void getGeneratedKeys() throws SQLException { + if (DB2_HACK.getGeneratedKeys(pstmt, list)) { + return; } - + int index = 0; try (ResultSet rset = pstmt.getGeneratedKeys()) { while (rset.next()) { Object idValue = rset.getObject(1); diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/persist/DB2GetKeys.java b/ebean-core/src/main/java/io/ebeaninternal/server/persist/DB2GetKeys.java new file mode 100644 index 000000000..38d0fe7b4 --- /dev/null +++ b/ebean-core/src/main/java/io/ebeaninternal/server/persist/DB2GetKeys.java @@ -0,0 +1,56 @@ +package io.ebeaninternal.server.persist; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; + +/** + * Horrible DB2 hack. + */ +final class DB2GetKeys { + + private static Class DB2_PREPARED_STATEMENT; + private static Method GET_DB_GENERATED_KEYS; + + static { + try { + DB2_PREPARED_STATEMENT = (Class) Class.forName("com.ibm.db2.jcc.DB2PreparedStatement"); + GET_DB_GENERATED_KEYS = DB2_PREPARED_STATEMENT.getDeclaredMethod("getDBGeneratedKeys"); + } catch (ClassNotFoundException | NoSuchMethodException | SecurityException cnf) { + // NOP + } + } + + static boolean getGeneratedKeys(PreparedStatement pstmt, List list) throws SQLException { + if (DB2_PREPARED_STATEMENT != null) { + PreparedStatement db2Stmt = null; + if (DB2_PREPARED_STATEMENT.isInstance(pstmt)) { + db2Stmt = pstmt; + } else if (pstmt.isWrapperFor(DB2_PREPARED_STATEMENT)) { + db2Stmt = pstmt.unwrap(DB2_PREPARED_STATEMENT); + } + if (db2Stmt != null) { + // WTF: https://stackoverflow.com/questions/41725492/how-to-get-auto-generated-keys-of-batch-insert-statement + ResultSet[] result; + try { + result = (ResultSet[]) GET_DB_GENERATED_KEYS.invoke(db2Stmt); + } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) { + throw new SQLException("Could not get generated keys for DB2", e); + } + int index = 0; + for (ResultSet resultSet : result) { + while (resultSet.next()) { + Object idValue = resultSet.getObject(1); + list.get(index).setGeneratedKey(idValue); + index++; + } + } + return true; + } + } + return false; + } +}