#2521 - Move horrible DB2 hack into DB2GetKeys

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 {
 
-  /**
+  private static final DB2GetKeys DB2_HACK = new DB2GetKeys();
-   * The underlying statement.
+
-   */
   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<BatchPostExecute> list = new ArrayList<>();
   private final String sql;
@@ -44,18 +36,6 @@ public final class BatchedPstmt implements SpiProfileTransactionEvent {
   private int[] results;
   private List<InputStream> inputStreams;
 
-  private static Class<? extends PreparedStatement> DB2_PREPARED_STATEMENT;
-  private static Method GET_DB_GENERATED_KEYS;
-  
-  static {
-    try {
-      DB2_PREPARED_STATEMENT = (Class<? extends PreparedStatement>) 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 {
+  private void getGeneratedKeys() throws SQLException {
-    int index = 0;
+    if (DB2_HACK.getGeneratedKeys(pstmt, list)) {
-    if (DB2_PREPARED_STATEMENT != null) {
+      return;
-      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;
-      }
     }
-
+    int index = 0;
     try (ResultSet rset = pstmt.getGeneratedKeys()) {
       while (rset.next()) {
         Object idValue = rset.getObject(1);

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<? extends PreparedStatement> DB2_PREPARED_STATEMENT;
+  private static Method GET_DB_GENERATED_KEYS;
+
+  static {
+    try {
+      DB2_PREPARED_STATEMENT = (Class<? extends PreparedStatement>) 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<BatchPostExecute> 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;
+  }
+}